#!/usr/bin/env python3
"""Batch Step 3 (Suno) + Step 4 (SeeDream) for all 11 songs."""
import json, os, subprocess, urllib.request, urllib.error, time
from pathlib import Path

BATCH_DIR = Path('/tmp/music-pipeline/batch')
SUNO_SH = Path('/root/.openclaw/workspace/music-pipeline3000/skills/suno/suno.sh')
REPO_DIR = Path('/root/.openclaw/workspace/music-pipeline3000')

# Load env
for line in (REPO_DIR / '.env').read_text().splitlines():
    line = line.strip()
    if line and not line.startswith('#') and '=' in line:
        k, v = line.split('=', 1)
        os.environ.setdefault(k, v)

SEEDREAM_TOKEN = os.environ.get('SEEDREAM_TOKEN', '')
SEEDREAM_MODEL = os.environ.get('SEEDREAM_MODEL', 'doubao-seedream-4-5-251128')

# Cover prompt templates based on song mood
def generate_cover_prompt(title, prompt_str, mood_hint=""):
    """Generate a cover art prompt based on song mood."""
    # Extract mood from prompt
    mood = ""
    for part in prompt_str.split(';'):
        if 'Mood:' in part:
            mood = part.split('Mood:')[1].strip()
            break
    
    return f"A minimalist album cover art, abstract artistic style, no text, no words, no letters, no human faces. Theme: '{title}'. Mood: {mood}. Cinematic lighting, rich color palette, professional photography quality, 4K, clean composition, suitable for music streaming platform cover art."


def run_suno(prompt_str, lyrics_str):
    """Run Suno generation and return result."""
    try:
        result = subprocess.run(
            [str(SUNO_SH), 'custom', prompt_str, lyrics_str],
            capture_output=True, text=True, timeout=300
        )
        return result.stdout
    except Exception as e:
        return f"ERROR: {e}"


def run_seedream(cover_prompt):
    """Call SeeDream API for cover generation."""
    if not SEEDREAM_TOKEN:
        return "ERROR: SEEDREAM_TOKEN not set"
    
    payload = {
        "action": "generate",
        "model": SEEDREAM_MODEL,
        "prompt": cover_prompt,
        "size": "2048x2048",
        "watermark": False
    }
    
    req = urllib.request.Request(
        'https://api.acedata.cloud/seedream/images',
        data=json.dumps(payload).encode(),
        headers={
            'accept': 'application/json',
            'authorization': f'Bearer {SEEDREAM_TOKEN}',
            'content-type': 'application/json',
        },
        method='POST'
    )
    
    try:
        with urllib.request.urlopen(req, timeout=180) as resp:
            data = json.loads(resp.read().decode())
            if data.get('success') and data.get('data'):
                return data['data'][0].get('image_url', 'NO_URL')
            return f"API_ERROR: {json.dumps(data, ensure_ascii=False)[:200]}"
    except Exception as e:
        return f"ERROR: {e}"


def main():
    dirs = sorted([d for d in BATCH_DIR.iterdir() if d.is_dir() and (d / 'step2.json').exists()])
    
    results = []
    for d in dirs:
        name = d.name
        step2 = json.loads((d / 'step2.json').read_text())
        title = step2.get('title', '')
        prompt_str = step2.get('prompt', '')
        lyrics_str = step2.get('lyrics', '')
        
        print(f"\n{'='*50}")
        print(f"Processing: {name} -> {title}")
        print(f"{'='*50}")
        
        # Generate cover prompt
        cover_prompt = generate_cover_prompt(title, prompt_str)
        step2['cover'] = {
            'template_id': 1,
            'template_name': 'cinematic_abstract',
            'scene_description': f'Abstract art representing the mood of "{title}"',
            'cover_prompt': cover_prompt,
            'selection_reason': 'Auto-generated based on song mood and title'
        }
        # Save updated step2
        (d / 'step2.json').write_text(json.dumps(step2, ensure_ascii=False, indent=2))
        
        # Step 3: Suno
        print(f"[Step 3] Running Suno for {title}...")
        suno_output = run_suno(prompt_str, lyrics_str)
        (d / 'step3_suno.txt').write_text(suno_output)
        
        # Extract URLs from suno output
        mp3_urls = []
        for line in suno_output.split('\n'):
            if 'cdn' in line and '.mp3' in line:
                # Try to extract URL
                import re
                urls = re.findall(r'https?://[^\s\)\"]+\.mp3[^\s\)\"]*', line)
                mp3_urls.extend(urls)
        
        # Also check for audio_url in JSON-like output
        if not mp3_urls:
            try:
                for line in suno_output.split('\n'):
                    if 'audio_url' in line or 'source_audio_url' in line:
                        urls = re.findall(r'https?://[^\s\"\,]+', line)
                        mp3_urls.extend(urls)
            except:
                pass
        
        print(f"[Step 3] Found {len(mp3_urls)} audio URLs")
        
        # Step 4: SeeDream cover
        print(f"[Step 4] Generating cover for {title}...")
        cover_url = run_seedream(cover_prompt)
        print(f"[Step 4] Cover: {cover_url[:80]}...")
        
        # Save cover URL
        step2['cover']['image_url'] = cover_url
        (d / 'step2.json').write_text(json.dumps(step2, ensure_ascii=False, indent=2))
        
        results.append({
            'dir': name,
            'title': title,
            'mp3_urls': mp3_urls,
            'cover_url': cover_url,
            'suno_ok': len(mp3_urls) > 0,
            'cover_ok': cover_url.startswith('http') and 'ERROR' not in cover_url,
        })
        
        # Small delay between songs
        time.sleep(2)
    
    # Summary
    print(f"\n{'='*50}")
    print("SUMMARY")
    print(f"{'='*50}")
    for r in results:
        s3 = '✅' if r['suno_ok'] else '❌'
        s4 = '✅' if r['cover_ok'] else '❌'
        print(f"{r['dir']} | {r['title']} | Suno: {s3} ({len(r['mp3_urls'])} urls) | Cover: {s4}")


if __name__ == '__main__':
    main()