"""
Step-by-step Suno upload and remix utilities with detailed logging
"""
import json
import time
import requests
from pathlib import Path
from datetime import datetime
from typing import Dict, List, Any, Tuple
import os

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

class StepLogger:
    def __init__(self, log_file: Path = None):
        self.logs = []
        self.log_file = log_file
        
    def log(self, step: str, level: str, message: str, data: Dict = None):
        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")[:-3]
        log_entry = {
            "timestamp": timestamp,
            "step": step,
            "level": level,
            "message": message
        }
        if data:
            log_entry["data"] = data
            
        self.logs.append(log_entry)
        
        # Format for file output
        log_text = f"[{timestamp}] [{step}] [{level}] {message}"
        if data:
            log_text += f" | DATA: {json.dumps(data, indent=2)}"
        
        print(log_text)
        
        if self.log_file:
            with open(self.log_file, "a") as f:
                f.write(log_text + "\n")
    
    def get_logs(self):
        return self.logs

def _headers():
    token = os.getenv("SUNO_TOKEN")
    if not token:
        raise RuntimeError("SUNO_TOKEN env var missing")
    return {"Authorization": f"Bearer {token}", "Content-Type": "application/json"}

def upload_audio_step_by_step(audio_path: Path, log_file: Path = None) -> Tuple[str, StepLogger]:
    """Upload audio file step by step with detailed logging"""
    logger = StepLogger(log_file)
    start_time = time.time()
    
    logger.log("MAIN", "INFO", f"Starting upload for: {audio_path}")
    
    # Step 1: Reserve upload slot
    logger.log("STEP 1", "INFO", f"Starting upload reservation for file: {audio_path.name}")
    
    payload = {
        "filename": audio_path.name,
        "content_type": "audio/wav"
    }
    logger.log("STEP 1", "DEBUG", "Request payload", payload)
    
    step1_start = time.time()
    r = requests.post(f"{API_BASE}/api/uploads/audio", json=payload, headers=_headers())
    step1_time = time.time() - step1_start
    
    logger.log("STEP 1", "DEBUG", f"Response status: {r.status_code} (took {step1_time:.2f}s)")
    
    if not r.ok:
        logger.log("STEP 1", "ERROR", f"Failed to reserve upload slot: {r.status_code}")
        raise RuntimeError(f"Upload reservation failed: {r.text}")
    
    upload_data = r.json()
    upload_id = upload_data.get("id") or upload_data.get("upload_id")
    logger.log("STEP 1", "SUCCESS", "Upload slot reserved", {
        k: v if k != "url" else "https://suno-uploads.s3.amazonaws.com/..." 
        for k, v in upload_data.items()
    })
    
    # Step 2: Upload file
    logger.log("STEP 2", "INFO", f"Starting file upload: {audio_path}")
    logger.log("STEP 2", "DEBUG", f"File size: {audio_path.stat().st_size / 1024 / 1024:.2f} MB")
    
    step2_start = time.time()
    
    if "fields" in upload_data:
        # Staging flow - multipart POST
        logger.log("STEP 2", "INFO", "Using staging upload (multipart POST)")
        url = upload_data["url"]
        fields = upload_data["fields"]
        logger.log("STEP 2", "DEBUG", f"Upload key: {fields.get('key')}")
        
        with open(audio_path, "rb") as f:
            files = {"file": (audio_path.name, f, fields.get("Content-Type", "audio/mpeg"))}
            r = requests.post(url, data=fields, files=files)
        
        if not r.ok:
            logger.log("STEP 2", "ERROR", f"File upload failed: {r.status_code}")
            raise RuntimeError(f"File upload failed: {r.text}")
            
        logger.log("STEP 2", "SUCCESS", f"File uploaded via multipart POST (took {time.time() - step2_start:.2f}s)")
    else:
        # Production flow - presigned PUT
        logger.log("STEP 2", "INFO", "Using production upload (presigned PUT)")
        presigned_url = upload_data.get("upload_url")
        
        with open(audio_path, "rb") as f:
            r = requests.put(presigned_url, data=f, headers={"Content-Type": "audio/wav"})
        
        if not r.ok:
            logger.log("STEP 2", "ERROR", f"File upload failed: {r.status_code}")
            raise RuntimeError(f"File upload failed: {r.text}")
            
        logger.log("STEP 2", "SUCCESS", f"File uploaded via PUT (took {time.time() - step2_start:.2f}s)")
    
    # Step 3: Mark upload complete
    logger.log("STEP 3", "INFO", f"Signaling upload complete for ID: {upload_id}")
    
    finish_payload = {
        "upload_type": "audio",
        "upload_filename": audio_path.name,
        "upload_key": upload_data.get("fields", {}).get("key", f"raw_uploads/{upload_id}.mp3")
    }
    logger.log("STEP 3", "DEBUG", "Finish payload", finish_payload)
    
    step3_start = time.time()
    r = requests.post(
        f"{API_BASE}/api/uploads/audio/{upload_id}/upload-finish",
        json=finish_payload,
        headers=_headers()
    )
    step3_time = time.time() - step3_start
    
    logger.log("STEP 3", "DEBUG", f"Response status: {r.status_code} (took {step3_time:.2f}s)")
    
    if not r.ok:
        logger.log("STEP 3", "ERROR", f"Failed to mark upload complete: {r.status_code}")
        raise RuntimeError(f"Upload finish failed: {r.text}")
    
    logger.log("STEP 3", "SUCCESS", "Upload marked as complete")
    
    # Step 4: Poll for processing completion
    logger.log("STEP 4", "INFO", f"Starting status polling for upload ID: {upload_id}")
    
    poll_count = 0
    poll_start = time.time()
    
    while True:
        poll_count += 1
        r = requests.get(f"{API_BASE}/api/uploads/audio/{upload_id}", headers=_headers())
        
        if not r.ok:
            logger.log("STEP 4", "ERROR", f"Status poll failed: {r.status_code}")
            raise RuntimeError(f"Status poll failed: {r.text}")
        
        status_data = r.json()
        status = status_data.get("status", "unknown")
        is_uploaded = status_data.get("is_file_uploaded", False)
        
        logger.log("STEP 4", "DEBUG", 
                  f"Poll #{poll_count} - Status: {status}, File uploaded: {is_uploaded}",
                  {"elapsed": f"{time.time() - poll_start:.1f}s"})
        
        if status == "complete" or status == "error":
            break
            
        if time.time() - poll_start > 180:
            logger.log("STEP 4", "ERROR", "Upload processing timed out")
            raise TimeoutError("Upload processing timed out")
            
        time.sleep(3)
    
    if status == "error":
        logger.log("STEP 4", "ERROR", "Upload processing failed", status_data)
        raise RuntimeError(f"Upload processing failed: {status_data}")
    
    logger.log("STEP 4", "SUCCESS", f"Upload processing complete after {poll_count} polls ({time.time() - poll_start:.1f}s)")
    
    # Step 5: Initialize clip
    logger.log("STEP 5", "INFO", f"Initializing clip from upload ID: {upload_id}")
    
    step5_start = time.time()
    r = requests.post(
        f"{API_BASE}/api/uploads/audio/{upload_id}/initialize-clip",
        json={},
        headers=_headers()
    )
    step5_time = time.time() - step5_start
    
    logger.log("STEP 5", "DEBUG", f"Response status: {r.status_code} (took {step5_time:.2f}s)")
    
    if not r.ok:
        logger.log("STEP 5", "ERROR", f"Clip initialization failed: {r.status_code}")
        raise RuntimeError(f"Clip initialization failed: {r.text}")
    
    clip_data = r.json()
    clip_id = clip_data.get("clip_id") or clip_data.get("id")
    
    if not clip_id:
        logger.log("STEP 5", "ERROR", "No clip ID returned", clip_data)
        raise RuntimeError("No clip ID returned from initialization")
    
    logger.log("STEP 5", "SUCCESS", f"Clip initialized: {clip_id}")
    
    total_time = time.time() - start_time
    logger.log("MAIN", "INFO", f"Upload complete! Clip ID: {clip_id} (Total time: {total_time:.1f}s)")
    
    return clip_id, logger

def upload_and_remix_step_by_step(audio_path: Path, prompt: str, log_file: Path = None, audio_weight: float = 0.8) -> Tuple[Dict, StepLogger]:
    """Upload audio and create remix/cover with detailed logging"""
    logger = StepLogger(log_file)
    start_time = time.time()
    
    logger.log("MAIN", "INFO", f"Starting upload + remix for: {audio_path}")
    logger.log("MAIN", "INFO", f"Remix prompt: {prompt}")
    
    # First, upload the audio
    clip_id, _ = upload_audio_step_by_step(audio_path, log_file)
    
    # Optional: Enable remixes for the clip
    logger.log("STEP 5.5", "INFO", f"Enabling remixes for clip: {clip_id}")
    r = requests.post(
        f"{API_BASE}/api/gen/{clip_id}/enable_remixes",
        headers=_headers()
    )
    if not r.ok:
        logger.log("STEP 5.5", "WARNING", f"Failed to enable remixes: {r.status_code}")
    
    # Optional: Set remix type
    logger.log("STEP 5.6", "INFO", f"Setting remix type to 'REMIX' for clip: {clip_id}")
    r = requests.post(
        f"{API_BASE}/api/gen/{clip_id}/update_remix_type",
        json={"type": "REMIX"},
        headers=_headers()
    )
    if not r.ok:
        logger.log("STEP 5.6", "WARNING", f"Failed to set remix type: {r.status_code}")
        logger.log("STEP 5.6", "DEBUG", f"Response: {r.text}")
    
    # Now generate the remix/cover
    logger.log("MAIN", "INFO", f"Starting remix generation with clip: {clip_id}")
    
    # Step 6: Generate cover/remix
    logger.log("STEP 6", "INFO", f"Starting cover generation with clip: {clip_id}")
    logger.log("STEP 6", "DEBUG", f"Prompt: {prompt}")
    logger.log("STEP 6", "DEBUG", f"Audio weight: {audio_weight}")
    
    # Add 'cover' to tags to potentially preserve vocals
    tags = prompt if 'cover' in prompt.lower() else f"{prompt}, cover"
    
    generation_payload = {
        "prompt": "",  # Empty for instrumental covers
        "generation_type": "TEXT",
        "tags": tags,
        "mv": "chirp-bluejay-t2",
        "task": "cover",
        "cover_clip_id": clip_id,
        "metadata": {
            "control_sliders": {
                "audio_weight": audio_weight,
                "style_weight": 0.5,
                "weirdness_constraint": 0.0  # Set to 0 as user requested
            },
            "is_remix": True
        }
    }
    
    logger.log("STEP 6", "DEBUG", "Generation payload", generation_payload)
    
    step6_start = time.time()
    r = requests.post(f"{API_BASE}/api/generate/v2", json=generation_payload, headers=_headers())
    step6_time = time.time() - step6_start
    
    logger.log("STEP 6", "DEBUG", f"Response status: {r.status_code} (took {step6_time:.2f}s)")
    
    if not r.ok:
        logger.log("STEP 6", "ERROR", f"Generation failed: {r.status_code}")
        logger.log("STEP 6", "ERROR", f"Response: {r.text}")
        raise RuntimeError(f"Generation failed: {r.text}")
    
    gen_data = r.json()
    gen_id = gen_data.get("id")
    
    logger.log("STEP 6", "SUCCESS", f"Generation started: {gen_id}")
    logger.log("STEP 6", "DEBUG", "Full response", gen_data)
    
    # Step 7: Poll for generation completion
    logger.log("STEP 7", "INFO", f"Polling generation status: {gen_id}")
    
    poll_count = 0
    poll_start = time.time()
    final_clips = []
    
    while True:
        poll_count += 1
        time.sleep(5)
        
        r = requests.get(f"{API_BASE}/api/generate/requests?ids={gen_id}", headers=_headers())
        
        if not r.ok:
            logger.log("STEP 7", "ERROR", f"Status poll failed: {r.status_code}")
            raise RuntimeError(f"Status poll failed: {r.text}")
        
        status_data = r.json()
        if status_data and len(status_data) > 0:
            gen_status = status_data[0]
            status = gen_status.get("status", "unknown")
            
            logger.log("STEP 7", "DEBUG", f"Poll #{poll_count} - Status: {status} ({time.time() - poll_start:.1f}s elapsed)")
            
            if status in ["complete", "streaming", "error", "failed"]:
                final_clips = gen_status.get("clips", [])
                break
        
        if time.time() - poll_start > 300:
            logger.log("STEP 7", "ERROR", "Generation timed out")
            raise TimeoutError("Generation timed out")
    
    if status in ["error", "failed"]:
        logger.log("STEP 7", "ERROR", "Generation failed", gen_status)
        raise RuntimeError(f"Generation failed: {gen_status}")
    
    logger.log("STEP 7", "SUCCESS", f"Generation complete after {poll_count} polls ({time.time() - poll_start:.1f}s)")
    
    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"Remix complete! Generated {len(final_clips)} clips (Total time: {total_time:.1f}s)")
    
    return result, logger