import os
# Auto-agree to Coqui TOS (Must be before imports)
os.environ["COQUI_TOS_AGREED"] = "1"

import json
import base64
import tempfile
import requests
import torch
import numpy as np

# Singleton for lazy loading
_model = None

def _allow_xtts_config_pickle():
    """Allow loading XTTS configs with torch >=2.6 safe loading."""
    add_safe = getattr(torch.serialization, "add_safe_globals", None)
    if not add_safe:
        return
    allowed = []
    try:
        from TTS.tts.configs.xtts_config import XttsConfig
        from TTS.tts.models.xtts import XttsAudioConfig
        allowed += [XttsConfig, XttsAudioConfig]
    except Exception as e:
        print(f"⚠️ Could not register safe globals for XTTS config: {e}")
    try:
        import TTS.config.shared_configs as shared_configs
        allowed += [v for v in shared_configs.__dict__.values() if isinstance(v, type)]
    except Exception as e:
        print(f"⚠️ Could not register shared config globals: {e}")
    try:
        import TTS.tts.models.xtts as xtts_models
        allowed += [v for v in xtts_models.__dict__.values() if isinstance(v, type)]
    except Exception as e:
        print(f"⚠️ Could not register XTTS model globals: {e}")
    if allowed:
        add_safe(allowed)

def get_model():
    global _model
    if _model is None:
        print("⏳ Loading XTTS Model (Lazy Load)....")
        # Lazy Import to prevent startup hang
        from TTS.api import TTS
        _allow_xtts_config_pickle()
        
        device = "cuda" if torch.cuda.is_available() else "cpu"
        print(f"🔧 XTTS Running on: {device}")
        
        # Load Model (download if needed)
        # Using default XTTS v2 model
        _model = TTS("tts_models/multilingual/multi-dataset/xtts_v2").to(device)
        print("✅ XTTS Model loaded successfully.")
    return _model

def synthesize(job: dict):
    from langdetect import detect
    model = get_model()
    ...
    
    text = job.get("input")
    if not text:
        raise ValueError("No input text provided")
        
    # Language handling
    language = job.get("language")
    if not language:
        try:
            # Simple detection
            detected = detect(text)
            # XTTS expects 2-letter codes usually. 
            # We assume detected is valid or mapped if needed.
            # Supported: en, es, fr, de, it, pt, pl, tr, ru, nl, cs, ar, zh-cn, ja, hu, ko
            language = detected
            print(f"🌍 Auto-detected language: {language}")
        except:
            language = "en"
            print("⚠️ Language detection failed, using 'en'")
            
    # Speaker Handling
    speaker_wav = None
    temp_files = []
    
    try:
        # Priority 1: Direct URL
        if job.get("voice_sample_url"):
            try:
                print(f"⬇️ Downloading voice sample from {job['voice_sample_url']}")
                r = requests.get(job["voice_sample_url"], timeout=10)
                r.raise_for_status()
                t = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
                t.write(r.content)
                t.close()
                speaker_wav = t.name
                temp_files.append(t.name)
            except Exception as e:
                print(f"❌ Failed to download voice sample: {e}")

        # Priority 2: Base64
        if not speaker_wav and job.get("voice_sample_base64"):
            try:
                b64 = job["voice_sample_base64"]
                decoded = base64.b64decode(b64)
                t = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
                t.write(decoded)
                t.close()
                speaker_wav = t.name
                temp_files.append(t.name)
            except Exception as e:
                print(f"❌ Failed to decode base64 voice: {e}")

        # Priority 3: Registry / Local File
        if not speaker_wav:
            voice_id = job.get("voice", "auto")
            if voice_id and voice_id != "auto":
                # Look in worker/voices/
                base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
                voice_path = os.path.join(base_dir, "voices", f"{voice_id}.wav")
                
                # Check for other extensions if wav missing
                if not os.path.exists(voice_path):
                    for ext in [".mp3", ".ogg", ".m4a"]:
                        p = os.path.join(base_dir, "voices", f"{voice_id}{ext}")
                        if os.path.exists(p):
                            voice_path = p
                            break
                            
                if os.path.exists(voice_path):
                    speaker_wav = voice_path
                    print(f"🗣️ Using registered voice: {voice_id}")
                else:
                    print(f"⚠️ Voice '{voice_id}' not found in registry.")

        # Priority 4: Default/Auto Voice
        if not speaker_wav:
            # Fallback to a default file if it exists
            base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
            default_path = os.path.join(base_dir, "voices", "default.wav")
            if os.path.exists(default_path):
                speaker_wav = default_path
                print("⚠️ Using default.wav")
            else:
                # If completely nothing, we can't synthesize with XTTS
                # Unless we use speaker_idxs (only for multi-speaker models w/o cloning?)
                # XTTS v2 IS zero-shot, needs reference.
                raise ValueError("No speaker reference found (url, base64, registry, or default.wav)")

        # Run Inference
        print(f"🎤 Synthesizing: '{text[:30]}...' Lang: {language}")
        
        # XTTS API returns List[float]
        wav = model.tts(text=text, speaker_wav=speaker_wav, language=language)
        
        return wav

    finally:
        # Cleanup
        for f in temp_files:
            try:
                os.remove(f)
            except:
                pass