"""
Suno vocal cover utilities based on working implementation
Uses cover task with detailed tags and v4.5 model
"""
import os
import time
import requests
from pathlib import Path
from typing import Dict, Tuple
from suno_utils import get_aligned_lyrics
from suno_step_by_step_utils import StepLogger, upload_audio_step_by_step, _headers

API_BASE = "https://studio-api.staging.suno.com"

def generate_vocal_cover_v45(
    clip_id: str, 
    transcript: str,
    style_genre: str,
    logger: StepLogger,
    audio_weight: float = 1.0,
    style_weight: float = 1.0,
    weirdness_constraint: float = 0.0
) -> Dict:
    """
    Generate a vocal-preserving cover using v4.5 model with detailed tags
    
    Args:
        clip_id: The uploaded audio clip ID
        transcript: The full transcript text (e.g. lyrics or dialogue)
        style_genre: Basic style/genre (e.g. "drum and bass", "house music")
        logger: StepLogger instance
        audio_weight: How much of original audio to preserve (0-1, default 1.0)
        style_weight: Style influence (0-1, default 1.0)
        weirdness_constraint: Creativity constraint (0-1, default 0.0)
    
    Returns:
        Generation response data
    """
    logger.log("VOCAL_COVER", "INFO", f"Starting vocal cover generation with clip: {clip_id}")
    logger.log("VOCAL_COVER", "DEBUG", f"Style genre: {style_genre}")
    logger.log("VOCAL_COVER", "DEBUG", f"Transcript length: {len(transcript)} chars")
    
    # Format the prompt with dialogue tag if needed
    if transcript and not transcript.startswith("["):
        if "\n" in transcript or "!" in transcript or "?" in transcript:
            prompt = f"[Dialogue]\n{transcript}"
        else:
            prompt = transcript
    else:
        prompt = transcript
    
    # Create detailed tags that explicitly mention vocals
    detailed_tags = generate_detailed_tags(style_genre)
    
    generation_payload = {
        "prompt": prompt,  # Full transcript goes here
        "generation_type": "TEXT",
        "tags": detailed_tags,  # Detailed description with vocal mentions
        "negative_tags": "instrumental, spoken word, speaking",  # Avoid instrumental versions
        "mv": "chirp-auk",  # v4.5 model
        "task": "cover",  # Using cover, not overpainting!
        "cover_clip_id": clip_id,  # Standard cover_clip_id field
        "metadata": {
            "control_sliders": {
                "audio_weight": audio_weight,
                "style_weight": style_weight,
                "weirdness_constraint": weirdness_constraint
            },
            "is_remix": True
        }
    }
    
    logger.log("VOCAL_COVER", "DEBUG", "Generation payload", generation_payload)
    
    start_time = time.time()
    r = requests.post(f"{API_BASE}/api/generate/v2", json=generation_payload, headers=_headers())
    elapsed = time.time() - start_time
    
    logger.log("VOCAL_COVER", "DEBUG", f"Response status: {r.status_code} (took {elapsed:.2f}s)")
    
    if not r.ok:
        logger.log("VOCAL_COVER", "ERROR", f"Generation failed: {r.status_code}")
        logger.log("VOCAL_COVER", "ERROR", f"Response: {r.text}")
        raise RuntimeError(f"Vocal cover generation failed: {r.text}")
    
    gen_data = r.json()
    gen_id = gen_data.get("id")
    
    logger.log("VOCAL_COVER", "SUCCESS", f"Vocal cover generation started: {gen_id}")
    logger.log("VOCAL_COVER", "DEBUG", "Full response", gen_data)
    
    return gen_data


def generate_detailed_tags(style_genre: str) -> str:
    """
    Generate detailed tags that explicitly mention vocals and describe the style
    
    Args:
        style_genre: Basic genre like "drum and bass", "house music", etc.
    
    Returns:
        Detailed tag description mentioning vocals
    """
    # Map genres to detailed descriptions that mention vocals
    genre_descriptions = {
        "drum and bass": "This drum n bass cover launches straight into rapid-fire percussion and a driving bassline and clear vocals from the first second. Syncopated synth stabs and crisp hi-hats layer with energetic vocals, while doubled harmonies and subtle FX sweeps keep momentum high.",
        "drum n bass": "This drum n bass cover launches straight into rapid-fire percussion and a driving bassline and clear vocals from the first second. Syncopated synth stabs and crisp hi-hats layer with energetic vocals, while doubled harmonies and subtle FX sweeps keep momentum high.",
        "house music": "This house music cover features a four-on-the-floor beat with clear vocals prominently mixed. Deep basslines and warm synth pads support the vocal performance, while classic house piano stabs and filtered sweeps create movement. The vocals remain crisp and upfront throughout.",
        "house": "This house music cover features a four-on-the-floor beat with clear vocals prominently mixed. Deep basslines and warm synth pads support the vocal performance, while classic house piano stabs and filtered sweeps create movement. The vocals remain crisp and upfront throughout.",
        "dubstep": "This dubstep cover opens with clear vocals over sparse atmospherics before the massive sub-bass drops. Syncopated rhythms and wobbling basslines frame the vocal performance, while glitchy effects and dramatic builds maintain energy. The vocals cut through the heavy bass with clarity.",
        "jazz": "This smooth jazz cover features clear vocals supported by warm piano chords and walking basslines. Brushed drums and subtle horn sections complement the vocal melody, while sophisticated harmonies add depth. The intimate vocal delivery remains the centerpiece throughout.",
        "smooth jazz": "This smooth jazz cover features clear vocals supported by warm piano chords and walking basslines. Brushed drums and subtle horn sections complement the vocal melody, while sophisticated harmonies add depth. The intimate vocal delivery remains the centerpiece throughout.",
        "lofi": "This lofi cover wraps the clear vocals in warm vinyl crackle and nostalgic atmospherics. Mellow hip-hop beats and dusty piano samples support the vocal performance, while subtle pitch variations add character. The vocals remain intelligible through the cozy production.",
        "lofi beats": "This lofi cover wraps the clear vocals in warm vinyl crackle and nostalgic atmospherics. Mellow hip-hop beats and dusty piano samples support the vocal performance, while subtle pitch variations add character. The vocals remain intelligible through the cozy production.",
        "chill lofi beats": "This lofi cover wraps the clear vocals in warm vinyl crackle and nostalgic atmospherics. Mellow hip-hop beats and dusty piano samples support the vocal performance, while subtle pitch variations add character. The vocals remain intelligible through the cozy production.",
        "orchestral": "This orchestral cover features clear vocals supported by sweeping string sections and dramatic brass. Timpani rolls and woodwind flourishes punctuate the vocal phrases, while the full orchestra swells during emotional peaks. The vocals soar above the symphonic arrangement.",
        "epic orchestral": "This epic orchestral cover features clear vocals supported by sweeping string sections and dramatic brass. Timpani rolls and woodwind flourishes punctuate the vocal phrases, while the full orchestra swells during emotional peaks. The vocals soar above the symphonic arrangement.",
        "classical": "This classical cover presents clear vocals with refined orchestral accompaniment. Elegant string arrangements and delicate woodwinds support the vocal line, while dynamic contrasts highlight the emotional content. The vocals maintain prominence throughout the sophisticated arrangement.",
        "classical symphony": "This classical symphony cover presents clear vocals with refined orchestral accompaniment. Elegant string arrangements and delicate woodwinds support the vocal line, while dynamic contrasts highlight the emotional content. The vocals maintain prominence throughout the sophisticated arrangement.",
        "heavy metal": "This heavy metal cover features powerful vocals over crushing guitar riffs and thunderous drums. Palm-muted chugs and soaring lead guitars frame the vocal performance, while double-bass drumming drives the intensity. The vocals cut through the wall of sound with commanding presence.",
        "metal": "This heavy metal cover features powerful vocals over crushing guitar riffs and thunderous drums. Palm-muted chugs and soaring lead guitars frame the vocal performance, while double-bass drumming drives the intensity. The vocals cut through the wall of sound with commanding presence."
    }
    
    # Get description or create a generic one
    style_lower = style_genre.lower().strip()
    if style_lower in genre_descriptions:
        return genre_descriptions[style_lower]
    else:
        # Generic description that mentions vocals
        return f"This {style_genre} cover features clear and prominent vocals throughout. The instrumental arrangement supports and enhances the vocal performance while maintaining the characteristic elements of {style_genre}. The vocals remain crisp and intelligible, perfectly balanced with the backing track."

def upload_and_create_vocal_cover_v45(
    clip_id: str,
    transcript: str,
    style_genre: str,
    log_file: Path = None,
    audio_weight: float = 1.0
) -> Tuple[Dict, StepLogger]:
    """
    Upload audio and create vocal-preserving cover with v4.5 model
    
    Args:
        clip_id: uploaded clip
        transcript: Full transcript of the audio (lyrics/dialogue)
        style_genre: Basic genre tags (e.g. "drum and bass")
        log_file: Optional log file path
        audio_weight: Audio preservation weight (0-1)
    
    Returns:
        Tuple of (result dict, logger)
    """
    logger = StepLogger(log_file)
    start_time = time.time()
    
    logger.log("MAIN", "INFO", f"Starting upload + vocal cover v4.5 for: {clip_id}")
    logger.log("MAIN", "INFO", f"Style genre: {style_genre}")
    logger.log("MAIN", "INFO", f"Transcript preview: {transcript[:100]}..." if len(transcript) > 100 else f"Transcript: {transcript}")
    
    # First, upload the audio
   
    # Generate vocal cover with v4.5
    gen_data = generate_vocal_cover_v45(
        clip_id=clip_id,
        transcript=transcript,
        style_genre=style_genre,
        logger=logger,
        audio_weight=audio_weight,
        style_weight=1.0,
        weirdness_constraint=0.0
    )
    upload_id = gen_data.get("up")
    gen_id = gen_data.get("id")
    
    # Poll for completion
    logger.log("POLLING", "INFO", f"Polling generation status: {gen_id}")
    
    poll_count = 0
    poll_start = time.time()
    final_clips = []
    
    while True:
        poll_count += 1
        time.sleep(5)
        
        elapsed = time.time() - poll_start
        
        r = requests.get(f"{API_BASE}/api/generate/requests?ids={gen_id}", headers=_headers())
        
        if r.ok:
            data = r.json()
            
            # Handle both response formats
            if isinstance(data, list) and data:
                # Direct list response
                status = data[0].get("status")
                logger.log("POLLING", "DEBUG", f"Poll #{poll_count} - Status: {status} ({elapsed:.1f}s elapsed)")
                
                if status == "complete":
                    clips = data[0].get("clips", [])
                    logger.log("POLLING", "SUCCESS", f"Generation complete after {poll_count} polls ({elapsed:.1f}s)")
                    final_clips = clips
                    break
                elif status == "error":
                    logger.log("POLLING", "ERROR", "Generation failed")
                    raise RuntimeError("Vocal cover generation failed")
            else:
                # Wrapped response format
                requests_data = data.get("requests", [])
                
                if requests_data:
                    status = requests_data[0].get("status")
                    logger.log("POLLING", "DEBUG", f"Poll #{poll_count} - Status: {status} ({elapsed:.1f}s elapsed)")
                    
                    if status == "complete":
                        clips = requests_data[0].get("clips", [])
                        logger.log("POLLING", "SUCCESS", f"Generation complete after {poll_count} polls ({elapsed:.1f}s)")
                        final_clips = clips
                        break
                    elif status == "error":
                        logger.log("POLLING", "ERROR", "Generation failed")
                        raise RuntimeError("Vocal cover generation failed")
        
        if elapsed > 300:  # 5 minute timeout
            logger.log("POLLING", "ERROR", "Timeout waiting for generation")
            raise TimeoutError("Generation timed out")
    
    total_time = time.time() - start_time
    
    result = {
        "upload_clip_id": clip_id,
        "generation_id": gen_id,
        "generated_clips": final_clips,
        "total_time": total_time
    }
    
    logger.log("MAIN", "SUCCESS", f"Vocal cover complete! Generated {len(final_clips)} clips (Total time: {total_time:.1f}s)")
    
    return result, logger