#!/usr/bin/env python3
# Copyright 2026 Umbra Tools (umbratools.com)
# SPDX-License-Identifier: Apache-2.0

"""Test a Scottish voice for a given speaker once project discovery completes.

Waits for ``analysis_status.json`` to reach ``complete`` (or ``error``), then:
  1. loads speakers.json and locates the requested speaker (primary/alias, case-insensitive),
  2. starts the Qwen3-TTS server if it isn't already up,
  3. generates a voice reference via VoiceDesign using the speaker's demographics +
     a mid-low, adult, medium-Scottish description,
  4. writes it next to the other project voices, drops a copy into the shared
     web samples dir for listening, appends it to the samples index page, and
     pins the description so the real "Generate All Voices" run uses this accent.

Usage:
  scripts/test_speaker_scottish.py --speaker Hallr --project "Last Dwarf"
"""

from __future__ import annotations

import argparse
import io
import json
import logging
import os
import subprocess
import sys
import time
from pathlib import Path

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

import httpx
import soundfile as sf

from umbraaudio.config import load_config

TTS_URL = os.getenv("UMBRA_TTS_URL", "http://localhost:8881")
SAMPLES_DIR = ROOT / "umbraaudio" / "web" / "static" / "samples"
MAX_WAIT_SECONDS = 6 * 3600
POLL_INTERVAL = 20

VOICE_REFERENCE_TEXT = (
    "When the sunlight strikes raindrops in the air, they act as a prism and form a rainbow. "
    "The rainbow is a division of white light into many beautiful colors. These take the shape "
    "of a long round arch, with its path high above, and its two ends apparently beyond the "
    "horizon. There is, according to legend, a boiling pot of gold at one end. People look, "
    "but no one ever finds it. When a man looks for something beyond his reach, his friends "
    "say he is looking for the pot of gold at the end of the rainbow."
)


def log(msg: str) -> None:
    print(msg, flush=True)


def build_description(speaker: dict) -> str:
    gender = ((speaker.get("demographics") or {}).get("gender") or "").strip().lower()
    if "female" in gender:
        return (
            "An adult female voice with a medium Scottish accent, mid-low pitch, "
            "warm and resonant, clear and intelligible."
        )
    if "male" in gender or "man" in gender:
        return (
            "An adult male voice with a medium Scottish accent, mid-low pitch, "
            "warm and slightly gravelly, clear and intelligible."
        )
    return (
        "An adult voice with a medium Scottish accent, mid-low pitch, "
        "warm and slightly gravelly, clear and intelligible."
    )


def wait_for_analysis(status_path: Path) -> dict:
    started = time.monotonic()
    while time.monotonic() - started < MAX_WAIT_SECONDS:
        try:
            with open(status_path) as f:
                st = json.load(f)
        except (OSError, json.JSONDecodeError):
            st = {"status": "unknown"}
        status = st.get("status")
        if status == "complete":
            log(f"Discovery complete: {st.get('detail')}")
            return st
        if status == "error":
            raise RuntimeError(f"Discovery failed: {st.get('error') or st.get('detail')}")
        log(f"[poll] discovering — {st.get('detail')} ({st.get('progress')}%)")
        time.sleep(POLL_INTERVAL)
    raise TimeoutError("Discovery did not complete within the wait window")


def find_speaker(speakers_path: Path, target: str) -> tuple[dict, str]:
    with open(speakers_path) as f:
        data = json.load(f)
    speakers = data.get("speakers", [])
    tl = target.lower()
    for s in speakers:
        if s.get("primary_name", "").lower() == tl:
            return s, s["primary_name"]
        for a in s.get("aliases", []):
            if a.lower() == tl:
                return s, s["primary_name"]
    raise KeyError(
        f"Speaker '{target}' not found. Known: {[s.get('primary_name') for s in speakers]}"
    )


def ensure_tts_server(cfg) -> None:
    try:
        r = httpx.get(f"{TTS_URL}/", timeout=3)
        if r.status_code < 500:
            log("TTS server already running")
            return
    except Exception:
        pass
    log("Starting Qwen3-TTS server...")
    env = dict(os.environ)
    env["QWEN_PROMPT_DIR"] = str(cfg.data_dir / "voice_prompts")
    log_fh = open(ROOT / "logs" / "qwen_tts_server.log", "a")
    subprocess.Popen(
        [sys.executable, "-u", str(ROOT / "scripts" / "qwen_tts_server.py")],
        cwd=str(ROOT),
        env=env,
        stdout=log_fh,
        stderr=subprocess.STDOUT,
    )
    for _ in range(120):
        try:
            r = httpx.get(f"{TTS_URL}/", timeout=3)
            if r.status_code < 500:
                log("TTS server ready")
                return
        except Exception:
            pass
        time.sleep(5)
    raise RuntimeError("TTS server did not become ready")


def generate_ref(description: str, output_path: Path) -> float:
    r = httpx.post(
        f"{TTS_URL}/v1/tts",
        data={
            "text": VOICE_REFERENCE_TEXT,
            "voice_description": description,
            "sample_rate": 24000,
        },
        timeout=900.0,
    )
    r.raise_for_status()
    wav, sr = sf.read(io.BytesIO(r.content), dtype="float32")
    sf.write(str(output_path), wav, sr)
    info = sf.info(str(output_path))
    return info.duration


def pin_description(speakers_path: Path, primary_name: str, description: str) -> None:
    with open(speakers_path) as f:
        data = json.load(f)
    for s in data.get("speakers", []):
        if s.get("primary_name") == primary_name:
            s["voice_description"] = description
            break
    with open(speakers_path, "w") as f:
        json.dump(data, f, indent=2)

    profiles_path = speakers_path.parent / "voice_profiles.json"
    profiles = {}
    if profiles_path.exists():
        with open(profiles_path) as f:
            profiles = json.load(f)
    profiles[primary_name] = {
        "voice_description": description,
        "source": "hallr_scottish_test",
    }
    with open(profiles_path, "w") as f:
        json.dump(profiles, f, indent=2)


def append_to_samples_index(display: str, filename: str, desc: str, tag: str) -> None:
    idx = SAMPLES_DIR / "index.html"
    section = (
        f'\n<h2 style="border-color:#8be9fd">★ {display}</h2>\n'
        f'<div class="row"><label>{desc}</label>'
        f'<audio controls preload="none" src="{filename}"></audio>'
        f'<span class="tag">{tag}</span></div>\n'
    )
    with open(idx) as f:
        html = f.read()
    if filename in html:
        return
    with open(idx, "a") as f:
        f.write(section)


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--project", default="Last Dwarf")
    ap.add_argument("--speaker", default="Hallr")
    args = ap.parse_args()

    cfg = load_config(args.project)
    status_path = cfg.data_dir / "analysis_status.json"
    speakers_path = cfg.data_dir / "speakers.json"
    result = wait_for_analysis(status_path)

    if not speakers_path.exists():
        raise RuntimeError("Discovery report 'complete' but speakers.json is missing")

    speaker, primary_name = find_speaker(speakers_path, args.speaker)
    log(f"Found '{primary_name}': {json.dumps(speaker, indent=2)[:800]}")

    description = build_description(speaker)
    log(f"Voice description:\n  {description}")

    ensure_tts_server(cfg)

    voices_dir = cfg.data_dir / "voices"
    voices_dir.mkdir(parents=True, exist_ok=True)
    ref_name = primary_name.lower().replace(" ", "_") + "_ref.wav"
    output_path = voices_dir / ref_name

    duration = generate_ref(description, output_path)
    log(f"Wrote {output_path} ({duration:.1f}s)")

    pin_description(speakers_path, primary_name, description)
    log("Pinned voice_description on the speaker + voice_profiles.json")

    SAMPLES_DIR.mkdir(parents=True, exist_ok=True)
    sample_name = f"{output_path.stem}_scottish.wav"
    sample_path = SAMPLES_DIR / sample_name
    import shutil

    shutil.copy2(output_path, sample_path)
    append_to_samples_index(
        f"{args.project} — {primary_name} (Scottish test)",
        sample_name,
        "mid-low pitch, adult, medium Scottish accent (waits on VoiceDesign)",
        f"{duration:.0f}s",
    )
    log(f"Sample: http://<webui>/static/samples/{sample_name} (see index.html)")


if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        log(f"HALLR SCOTTISH TEST FAILED: {e}")
        raise