#!/usr/bin/env python3
"""Start the local CAM++ model service and FunASR WebSocket backend together."""

from __future__ import annotations

import json
import os
import subprocess
import sys
import time
from pathlib import Path
from urllib.error import URLError
from urllib.request import urlopen

from dotenv import load_dotenv

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from scripts.model_manifest import (
    auxiliary_models,
    load_manifest,
    model_directory,
    resolve_auxiliary_model_id,
    resolve_model_id,
)

load_dotenv(PROJECT_ROOT / ".env")


def local_model(requested: str, models_dir: Path, kind: str) -> Path:
    """Resolve ASR/VAD IDs through the shared model manifest."""
    name = requested.strip()
    configured_path = Path(name)
    direct_candidates = (
        [configured_path]
        if configured_path.is_absolute()
        else [PROJECT_ROOT / configured_path, models_dir / configured_path]
    )
    for candidate in direct_candidates:
        if candidate.is_dir() and (candidate / "configuration.json").is_file():
            return candidate.resolve()

    manifest = load_manifest()
    if kind == "asr":
        model_id = resolve_model_id(name, manifest)
    else:
        model_id = resolve_auxiliary_model_id(name, manifest, kind=kind)
    candidate = model_directory(model_id, manifest, models_dir)
    if candidate.is_dir() and (candidate / "configuration.json").is_file():
        return candidate.resolve()
    raise FileNotFoundError(
        f"Local {kind.upper()} model '{name}' is missing; expected: {candidate}"
    )


def local_cam_model(models_dir: Path) -> Path:
    """Require one complete CAM++ speaker verification asset from the manifest."""
    manifest = load_manifest()
    override = os.getenv("CAM_MODEL_PATH", "").strip()
    if override:
        path = Path(override)
        if not path.is_absolute():
            path = models_dir / path
        path = path.resolve()
        configs = [
            config for config in auxiliary_models(manifest).values()
            if config.get("kind") == "speaker_verification"
        ]
        if path.is_dir() and configs and all(
            (path / relative).is_file() for relative in configs[0].get("required_files", [])
        ):
            return path
        raise FileNotFoundError(f"CAM++ model is missing or incomplete: {path}")
    checked = []
    for model_id, config in auxiliary_models(manifest).items():
        if config.get("kind") != "speaker_verification":
            continue
        path = model_directory(model_id, manifest, models_dir)
        required = [path / relative for relative in config.get("required_files", [])]
        if path.is_dir() and all(item.is_file() for item in required):
            return path.resolve()
        checked.append(str(path))
    raise FileNotFoundError("CAM++ speaker model is missing; checked: " + ", ".join(checked))


def wait_for_health(url: str, process: subprocess.Popen, key: str, seconds: int = 300) -> None:
    """Wait until a child has loaded its models, failing when it exits early."""
    deadline = time.monotonic() + seconds
    while time.monotonic() < deadline:
        code = process.poll()
        if code is not None:
            raise RuntimeError(f"Service exited before it became ready ({url}, exit={code})")
        try:
            with urlopen(url, timeout=2) as response:
                data = json.load(response)
            if data.get(key):
                return
        except (OSError, ValueError, URLError):
            pass
        time.sleep(0.5)
    raise TimeoutError(f"Service did not become ready within {seconds}s: {url}")


def stop_child(process: subprocess.Popen | None) -> None:
    """Stop a supervised model or WebSocket process on shutdown."""
    if process is None or process.poll() is not None:
        return
    process.terminate()
    try:
        process.wait(timeout=10)
    except subprocess.TimeoutExpired:
        process.kill()
        process.wait()


def main() -> None:
    """Require ASR, VAD, and CAM++ before exposing the backend WebSocket."""
    models_dir = Path(os.getenv("MODEL_DIR", "models"))
    if not models_dir.is_absolute():
        models_dir = PROJECT_ROOT / models_dir
    models_dir = models_dir.resolve()
    asr = local_model(
        os.getenv("FUNASR_ASR_MODEL", "paraformer-zh-streaming"), models_dir, "asr"
    )
    vad = local_model(os.getenv("FUNASR_VAD_MODEL", "fsmn-vad"), models_dir, "vad")
    cam = local_cam_model(models_dir)
    print(f"Local models: ASR={asr}; VAD={vad}; CAM++={cam}", flush=True)

    env = os.environ.copy()
    env.update({
        "MODEL_DIR": str(models_dir),
        "FUNASR_ASR_MODEL": str(asr),
        "FUNASR_VAD_MODEL": str(vad),
        "CAM_MODEL_PATH": str(cam),
        "AUXILIARY_SERVICE_URL": "http://127.0.0.1:8010",
    })
    auxiliary = None
    websocket = None
    try:
        auxiliary = subprocess.Popen(
            [sys.executable, "-m", "scripts.auxiliary_server"],
            cwd=PROJECT_ROOT, env=env,
        )
        wait_for_health("http://127.0.0.1:8010/health", auxiliary, "speaker_embedding_ready")
        websocket = subprocess.Popen(
            [sys.executable, "-m", "scripts.run_funasr_demo", "--no-browser"],
            cwd=PROJECT_ROOT, env=env,
        )
        web_port = int(env.get("WEB_PORT", "8082"))
        wait_for_health(
            f"http://127.0.0.1:{web_port}/api/config", websocket, "engine"
        )
        print(f"Backend ready: ws://127.0.0.1:{web_port}/ws", flush=True)
        while True:
            for label, process in (("CAM++", auxiliary), ("WebSocket", websocket)):
                code = process.poll()
                if code is not None:
                    raise RuntimeError(f"{label} service exited (exit={code})")
            time.sleep(0.5)
    except KeyboardInterrupt:
        pass
    finally:
        stop_child(websocket)
        stop_child(auxiliary)


if __name__ == "__main__":
    main()