"""Julio XTTS Voice - Train + Generate Noticiero"""
import os, sys, subprocess, threading, glob, gc, re, time, json

# Monkey-patch TTS TOS prompt (global solution)
import TTS.utils.manage as _tts_manage
_tts_manage.ModelManager.ask_tos = lambda self, path: True

import torch

stage = {"status": "ready", "output": "", "gen_status": "idle", "gen_output": "", "final_audio": ""}
print(f"Dataset: {len(glob.glob('dataset/wavs/*.wav'))} wavs", flush=True)

MODEL_HUB = "Converso72/julio-xtts-voice"
MODEL_CACHE = "/app/model_cache/best_model.pth"

def find_best():
    # Check ft_output first
    m = sorted(glob.glob("ft_output/**/best_model*.pth", recursive=True))
    if m:
        return m[-1]
    # Check cache
    if os.path.exists(MODEL_CACHE):
        return MODEL_CACHE
    # Try to download from hub
    try:
        from huggingface_hub import hf_hub_download
        os.makedirs(os.path.dirname(MODEL_CACHE), exist_ok=True)
        path = hf_hub_download(repo_id=MODEL_HUB, filename="best_model.pth", local_dir=os.path.dirname(MODEL_CACHE))
        return path
    except:
        pass
    return None

def get_status():
    import torch
    c = torch.cuda.is_available()
    g = torch.cuda.get_device_name(0) if c else "NONE"
    b = find_best()
    mi = f"\nModel: {os.path.basename(b)} ({os.path.getsize(b)/1e6:.0f} MB)" if b else ""
    return f"GPU: {g}\nCUDA: {c}\n" + mi

def start_train():
    if not os.path.exists("train.csv"):
        return "Missing train.csv"
    stage["status"] = "running"
    stage["output"] = ""
    def _run():
        try:
            r = subprocess.run([sys.executable, "train.py"], capture_output=True, text=True, timeout=6000)
            stage["output"] = (r.stdout[-2000:] + "\nSTDERR:\n" + r.stderr[-1000:]).strip()
        except Exception as e:
            stage["output"] = f"Error: {e}"
        stage["status"] = "done"
    threading.Thread(target=_run, daemon=True).start()
    return "Training started..."

def refresh():
    if stage["status"] == "ready":
        return get_status()
    elif stage["status"] == "running":
        return "Training... ⏳"
    return stage["output"] or "Done!"

# --- GENERATE NOTICIERO ---
def start_generate():
    """Generate 30-min news audio on GPU"""
    best = find_best()
    if not best:
        return "No model trained yet!"
    
    from TTS.api import TTS
    
    # Read script
    script_path = "/app/noticiero_30min.txt"
    if not os.path.exists(script_path):
        return "Script not found! Upload noticiero_30min.txt first."
    
    with open(script_path, "r", encoding="utf-8") as f:
        script = f.read()
    
    sentences = re.split(r"(?<=[.!?])\s+", script.replace("\n", " ").replace("  ", " "))
    sentences = [s.strip() for s in sentences if len(s.strip()) > 15]
    
    stage["gen_status"] = "running"
    stage["gen_output"] = f"Loading model... ({len(sentences)} sentences)"
    
    def _generate():
        try:
            # Load fine-tuned model
            stage["gen_output"] = "Loading XTTS model on GPU..."
            tts = TTS("tts_models/multilingual/multi-dataset/xtts_v2", gpu=True)
            state = torch.load(best, map_location="cpu", weights_only=True)
            inner = state["model"]
            model = tts.synthesizer.tts_model
            model.gpt.load_state_dict({k.replace("xtts.gpt.", ""): v for k, v in inner.items() if k.startswith("xtts.gpt.")}, strict=False)
            model.hifigan_decoder.load_state_dict({k.replace("xtts.hifigan_decoder.", ""): v for k, v in inner.items() if k.startswith("xtts.hifigan_decoder.")}, strict=False)
            
            speaker_wav = "dataset/wavs/julio_001.wav"
            out_dir = "/app/audio_noticiero"
            os.makedirs(out_dir, exist_ok=True)
            
            chunks = []
            start = time.time()
            for i, sent in enumerate(sentences):
                out = os.path.join(out_dir, f"chunk_{i:04d}.wav")
                stage["gen_output"] = f"Generating [{i+1}/{len(sentences)}]..."
                gc.collect()
                torch.cuda.empty_cache()
                try:
                    tts.tts_to_file(text=sent, speaker_wav=speaker_wav, language="es", file_path=out,
                        split_sentences=False, temperature=0.3, repetition_penalty=3.0,
                        top_k=15, top_p=0.7, speed=1.3)
                    chunks.append(out)
                except Exception as e:
                    stage["gen_output"] = f"Error at {i+1}: {e}"
                    break
            
            elapsed = time.time() - start
            
            # Concatenate WAVs with Python (no ffmpeg needed)
            if len(chunks) > 1:
                stage["gen_output"] = f"Concatenating {len(chunks)} chunks..."
                final = os.path.join(out_dir, "noticiero_final.wav")
                stage["final_audio"] = final
                
                with open(chunks[0], "rb") as f:
                    header = f.read(44)
                data_chunks = []
                for ch in chunks:
                    with open(ch, "rb") as f:
                        f.read(44)
                        data_chunks.append(f.read())
                
                total_data = b"".join(data_chunks)
                total_size = 36 + len(total_data)
                
                import struct
                header_patched = bytearray(header)
                struct.pack_into("<I", header_patched, 4, total_size)
                struct.pack_into("<I", header_patched, 40, len(total_data))
                
                with open(final, "wb") as f:
                    f.write(header_patched)
                    f.write(total_data)
                
                sz = os.path.getsize(final)
                dur = sz / 22050 / 2
                stage["gen_output"] = f"✅ GENERATED: {len(chunks)}/{len(sentences)} in {elapsed:.0f}s\nFile: noticiero_final.wav\nSize: {sz/1e6:.1f} MB\nDuration: {dur:.0f}s ({dur/60:.1f} min)\n\nClick the Download button below to get the file!"
            else:
                stage["gen_output"] = "Not enough chunks generated!"
        except Exception as e:
            stage["gen_output"] = f"Error: {e}"
        stage["gen_status"] = "done"
    
    threading.Thread(target=_generate, daemon=True).start()
    return f"Generating {len(sentences)} sentences on GPU..."

def refresh_gen():
    if stage["gen_status"] == "running":
        return stage["gen_output"]
    elif stage["gen_status"] == "done":
        s = stage["gen_output"]
        stage["gen_status"] = "idle"
        return s
    return stage["gen_output"] or "Ready"

# --- UI ---
import gradio as gr
with gr.Blocks(title="Julio XTTS Voice") as demo:
    gr.Markdown("# 🎤 Julio XTTS Voice - Train & Generate")
    with gr.Tab("Training"):
        s = gr.Textbox(label="Status", lines=15)
        with gr.Row():
            gr.Button("🔍 Check").click(fn=refresh, outputs=s)
            gr.Button("🚀 Train").click(fn=start_train, outputs=s)
        f = gr.File(label="Download model")
        gr.Button("📥 Download").click(fn=lambda: find_best(), outputs=f)
    with gr.Tab("Generate Noticiero"):
        g = gr.Textbox(label="Generation Status", lines=15, value="Load script and click Generate")
        with gr.Row():
            gr.Button("🎬 Generate 30-min Audio").click(fn=start_generate, outputs=g)
            gr.Button("🔄 Refresh").click(fn=refresh_gen, outputs=g)
        ga = gr.File(label="Download Generated Audio")
        gr.Button("📥 Download Audio").click(fn=lambda: stage.get("final_audio", "") or None, outputs=ga)

demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))