#!/usr/bin/env python3
"""Generate audio for chapters using Qwen3-TTS with voice cloning from references."""
from __future__ import annotations

import json
import logging
import sqlite3
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent))

from umbraaudio.config import load_config
from umbraaudio.pronunciation import load_phonetic_names
from umbraaudio.voices.lmstudio_lifecycle import get_loaded_model, get_vram_usage_mb, is_running
from umbraaudio.voices.model_coordinator import ModelCoordinator, PipelineStage
from umbraaudio.voices.qwen_tts import VOICE_REFERENCE_TEXT, QwenTTSEngine

logger = logging.getLogger(__name__)


def load_voice_references(voices_dir: Path) -> dict[str, Path]:
    """Load voice reference files, mapping character name -> ref audio path."""
    refs: dict[str, Path] = {}
    if not voices_dir.exists():
        return refs
    for f in voices_dir.glob("*_ref.wav"):
        name = f.stem.replace("_ref", "").replace("_", " ")
        refs[name] = f
    return refs


def load_speakers(data_dir: Path) -> dict[str, dict]:
    """Load speakers from JSON file."""
    path = data_dir / "speakers.json"
    if not path.exists():
        return {}
    with open(path) as f:
        data = json.load(f)
    return {s["primary_name"]: s for s in data.get("speakers", [])}


def get_speaker_name_map(speakers: dict[str, dict]) -> dict[str, str]:
    """Build a mapping from all aliases to primary name."""
    alias_map: dict[str, str] = {}
    for name, info in speakers.items():
        alias_map[name.lower()] = name
        for alias in info.get("aliases", []):
            alias_map[alias.lower()] = name
    return alias_map


def _apply_phonetics(text: str, phonetic: dict[str, str]) -> str:
    if not phonetic:
        return text
    for name, phonetic_spelling in sorted(phonetic.items(), key=lambda x: -len(x[0])):
        text = text.replace(name, phonetic_spelling)
    return text


def resolve_reference(
    speaker_name: str,
    voice_refs: dict[str, Path],
    alias_map: dict[str, str],
    narrator_ref: Path | None = None,
) -> Path | None:
    """Find the voice reference file for a speaker."""
    if speaker_name in voice_refs:
        return voice_refs[speaker_name]
    for name, path in voice_refs.items():
        if name.lower() == speaker_name.lower():
            return path
    canonical = alias_map.get(speaker_name.lower())
    if canonical and canonical in voice_refs:
        return voice_refs[canonical]
    return narrator_ref


def generate_chapter_audio(
    chapter_id: int,
    engine: QwenTTSEngine,
    voice_refs: dict[str, Path],
    alias_map: dict[str, str],
    output_base: Path,
    db_path: Path,
    narrator_ref: Path | None = None,
) -> int:
    """Generate audio for all utterances in a chapter using voice cloning."""
    db = sqlite3.connect(str(db_path))
    db.row_factory = sqlite3.Row

    chapter = db.execute("SELECT * FROM chapters WHERE id = ?", (chapter_id,)).fetchone()
    if not chapter:
        logger.error("Chapter %d not found", chapter_id)
        return 0

    title = chapter["title"] or f"Chapter {chapter['chapter_index']+1}"
    logger.info("Generating audio for: %s", title)

    utterances = db.execute(
        """
        SELECT u.id, u.speaker_name, u.text, u.utterance_index
        FROM utterances u
        WHERE u.chapter_id = ?
        ORDER BY u.utterance_index
        """,
        (chapter_id,),
    ).fetchall()
    logger.info("Found %d utterances", len(utterances))

    chapter_dir = output_base / f"ch{chapter_id}"
    chapter_dir.mkdir(parents=True, exist_ok=True)

    generated = 0
    for utt in utterances:
        utt_id = utt["id"]
        speaker = utt["speaker_name"] or "narrator"
        text = utt["text"]
        idx = utt["utterance_index"]

        if not text or not text.strip():
            continue

        ref_path = resolve_reference(speaker, voice_refs, alias_map, narrator_ref)

        filename = f"utt_{idx:04d}.wav"
        output_path = chapter_dir / filename

        if output_path.exists():
            generated += 1
            continue

        spoken_text = _apply_phonetics(text, phonetic)

        try:
            if ref_path:
                engine.generate(
                    text=spoken_text,
                    output_path=output_path,
                    voice_ref=ref_path,
                    ref_text=VOICE_REFERENCE_TEXT,
                )
            else:
                engine.generate(
                    text=text,
                    output_path=output_path,
                    instruct="A calm, neutral voice speaking clearly.",
                )

            if output_path.exists():
                generated += 1
                if generated % 20 == 0:
                    logger.info("  Generated %d/%d utterances", generated, len(utterances))
            else:
                logger.warning("  Failed: utt_%04d (%s)", idx, speaker)

        except Exception as e:
            logger.error("  Error: utt_%04d (%s): %s", idx, speaker, e)

    logger.info("Generated %d/%d utterances for %s", generated, len(utterances), title)
    db.close()
    return generated


def main():
    import argparse

    parser = argparse.ArgumentParser(description="Generate chapter audio with Qwen3-TTS")
    parser.add_argument("chapters", nargs="*", type=int, help="Chapter IDs to generate")
    parser.add_argument("--force", action="store_true", help="Overwrite existing audio")
    parser.add_argument("--all", action="store_true", help="Generate all chapters")
    parser.add_argument(
        "--project", type=str, default=None, help="Project name or path"
    )
    args = parser.parse_args()

    logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")

    cfg = load_config(args.project)

    coordinator = ModelCoordinator(cfg)
    if not coordinator.switch_to(PipelineStage.UTTERANCE_GENERATION).get("ok"):
        logger.error("Failed to switch models to UTTERANCE_GENERATION stage")
        sys.exit(1)

    voices_dir = cfg.voices_dir

    voice_refs = load_voice_references(voices_dir)
    speakers = load_speakers(cfg.data_dir)
    alias_map = get_speaker_name_map(speakers)

    phonetic = load_phonetic_names(cfg.data_dir)
    logger.info("Loaded %d confirmed phonetic name mappings", len(phonetic))

    narrator_ref = None
    for name, path in voice_refs.items():
        if name.lower() == "narrator":
            narrator_ref = path
            break
    if not narrator_ref:
        for name, path in voice_refs.items():
            if name.lower() == "kainda":
                narrator_ref = path
                break

    logger.info("Loaded %d voice references", len(voice_refs))
    logger.info("Loaded %d speakers", len(speakers))

    vram = get_vram_usage_mb()
    logger.info("VRAM: %dMB used, %dMB free", vram["used_mb"], vram["free_mb"])

    if is_running():
        model = get_loaded_model()
        if model:
            logger.warning("LM Studio has model loaded: %s", model)
            logger.warning("TTS needs ~4GB VRAM. If OOM occurs, unload the model in LM Studio UI.")

    engine = QwenTTSEngine(device="cuda", sample_rate=24000)

    output_base = cfg.data_dir / "chapter_audio"
    output_base.mkdir(parents=True, exist_ok=True)

    if args.all:
        db = sqlite3.connect(str(cfg.db_path))
        db.row_factory = sqlite3.Row
        chapters = db.execute(
            "SELECT id FROM chapters WHERE project_id=1 ORDER BY chapter_index"
        ).fetchall()
        db.close()
        chapter_ids = [c["id"] for c in chapters]
    elif args.chapters:
        chapter_ids = args.chapters
    else:
        chapter_ids = [5]

    if args.force:
        for ch_id in chapter_ids:
            ch_dir = output_base / f"ch{ch_id}"
            if ch_dir.exists():
                for f in ch_dir.glob("*.wav"):
                    f.unlink()
                logger.info("Cleared existing audio for chapter %d", ch_id)

    total_generated = 0
    for ch_id in chapter_ids:
        count = generate_chapter_audio(
            ch_id, engine, voice_refs, alias_map, output_base, cfg.db_path, narrator_ref
        )
        total_generated += count

    engine.unload_all()
    coordinator.switch_to(PipelineStage.IDLE)

    logger.info("Total generated: %d utterances across %d chapters", total_generated, len(chapter_ids))


if __name__ == "__main__":
    main()