"""
Suno audio download utilities
"""
import os
import time
import requests
from pathlib import Path
from typing import Optional

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

def _headers():
    token = os.getenv("SUNO_TOKEN")
    if not token:
        raise RuntimeError("SUNO_TOKEN env var missing")
    return {"Authorization": f"Bearer {token}"}

def download_suno_clip(clip_id: str, output_path: Path, max_wait: int = 300) -> bool:
    """
    Download a Suno clip by ID. Returns True if successful.
    """
    print(f"[download] Starting download for clip: {clip_id}")
    
    # First, get clip info to check if it's ready
    clip_url = f"{API_BASE}/api/feed/?ids={clip_id}"
    
    start_time = time.time()
    poll_count = 0
    
    while time.time() - start_time < max_wait:
        poll_count += 1
        
        r = requests.get(clip_url, headers=_headers())
        if not r.ok:
            print(f"[download] Failed to get clip info: {r.status_code}")
            time.sleep(5)
            continue
            
        data = r.json()
        if not data or len(data) == 0:
            print(f"[download] No clip data returned")
            time.sleep(5)
            continue
            
        clip_info = data[0]
        status = clip_info.get("status", "unknown")
        audio_url = clip_info.get("audio_url", "")
        
        print(f"[download] Poll #{poll_count} - Status: {status}, Audio URL: {'Yes' if audio_url else 'No'}")
        
        if status in ["error", "failed"]:
            print(f"[download] Clip generation failed: {status}")
            return False
            
        if audio_url:
            # Download the audio
            print(f"[download] Downloading from: {audio_url}")
            
            r = requests.get(audio_url, stream=True)
            if r.ok:
                with open(output_path, "wb") as f:
                    for chunk in r.iter_content(chunk_size=8192):
                        f.write(chunk)
                
                file_size = output_path.stat().st_size
                print(f"[download] Download complete: {output_path} ({file_size / 1024 / 1024:.2f} MB)")
                return True
            else:
                print(f"[download] Failed to download audio: {r.status_code}")
                
        time.sleep(5)
    
    print(f"[download] Timed out waiting for audio after {poll_count} polls")
    return False

def wait_and_download_clips(clip_ids: list[str], output_dir: Path, prefix: str = "clip") -> dict[str, Path]:
    """
    Wait for multiple clips to be ready and download them.
    Returns a dict mapping clip_id to downloaded file path.
    """
    results = {}
    
    for i, clip_id in enumerate(clip_ids):
        output_path = output_dir / f"{prefix}_{i}.mp3"
        print(f"\n[download] Processing clip {i+1}/{len(clip_ids)}: {clip_id}")
        
        if download_suno_clip(clip_id, output_path):
            results[clip_id] = output_path
        else:
            print(f"[download] Failed to download clip: {clip_id}")
    
    return results