#!/usr/bin/env python3
from __future__ import annotations

import argparse
import shutil
import subprocess
import sys
import tarfile
import tempfile
import urllib.request
from pathlib import Path

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

from owner_voice_pet.speech_models import (
    DEFAULT_DENOISER_PATH,
    DEFAULT_DENOISER_URL,
    DEFAULT_KWS_DIR,
    DEFAULT_KWS_KEYWORDS,
    DEFAULT_KWS_URL,
    DEFAULT_STT_DIR,
    DEFAULT_STT_URL,
    DEFAULT_VAD_URL,
    check_speech_models,
    write_default_manifest,
)


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Download local speech models for owner_voice_pet.")
    parser.add_argument("--dir", default="models", help="Target model directory. Default: models")
    parser.add_argument("--force", action="store_true", help="Re-download and replace existing model files")
    args = parser.parse_args(argv)

    target = Path(args.dir)
    target.mkdir(parents=True, exist_ok=True)
    download_kws(target, force=args.force)
    download_vad(target, force=args.force)
    download_stt(target, force=args.force)
    download_denoiser(target, force=args.force)
    manifest = write_default_manifest(target)
    status = check_speech_models(target, require_sherpa=False)
    print(f"models_dir={target}")
    print(f"manifest={manifest}")
    print(f"missing_files={list(status.missing_files)}")
    return 0 if not status.missing_files else 1


def download_kws(models_dir: Path, *, force: bool = False) -> Path:
    output_dir = models_dir / "wake" / DEFAULT_KWS_DIR
    if output_dir.exists() and any(output_dir.iterdir()) and not force:
        print(f"skip existing {output_dir}")
    else:
        output_dir.parent.mkdir(parents=True, exist_ok=True)
        last_error: Exception | None = None
        for attempt in range(1, 4):
            with tempfile.TemporaryDirectory() as tmp:
                archive = Path(tmp) / f"{DEFAULT_KWS_DIR}.tar.bz2"
                try:
                    download_file(DEFAULT_KWS_URL, archive)
                    extract_tar_bz2(archive, output_dir.parent)
                    last_error = None
                    break
                except Exception as exc:
                    last_error = exc
                    print(f"download/extract KWS attempt {attempt} failed: {exc}")
                    if output_dir.exists():
                        shutil.rmtree(output_dir)
        if last_error is not None:
            raise last_error
    keywords = models_dir / "wake" / "keywords.txt"
    if force or not keywords.exists():
        keywords.parent.mkdir(parents=True, exist_ok=True)
        keywords.write_text(DEFAULT_KWS_KEYWORDS, encoding="utf-8")
        print(f"wrote {keywords}")
    else:
        print(f"skip existing {keywords}")
    return output_dir


def download_vad(models_dir: Path, *, force: bool = False) -> Path:
    output = models_dir / "vad" / "silero_vad.onnx"
    if output.exists() and not force:
        print(f"skip existing {output}")
        return output
    output.parent.mkdir(parents=True, exist_ok=True)
    download_file(DEFAULT_VAD_URL, output)
    return output


def download_stt(models_dir: Path, *, force: bool = False) -> Path:
    output_dir = models_dir / "stt" / DEFAULT_STT_DIR
    if output_dir.exists() and any(output_dir.iterdir()) and not force:
        print(f"skip existing {output_dir}")
        return output_dir
    output_dir.parent.mkdir(parents=True, exist_ok=True)
    last_error: Exception | None = None
    for attempt in range(1, 4):
        with tempfile.TemporaryDirectory() as tmp:
            archive = Path(tmp) / f"{DEFAULT_STT_DIR}.tar.bz2"
            try:
                download_file(DEFAULT_STT_URL, archive)
                extract_tar_bz2(archive, output_dir.parent)
                return output_dir
            except Exception as exc:
                last_error = exc
                print(f"download/extract attempt {attempt} failed: {exc}")
                if output_dir.exists():
                    shutil.rmtree(output_dir)
    if last_error is not None:
        raise last_error
    return output_dir


def download_denoiser(models_dir: Path, *, force: bool = False) -> Path:
    output = models_dir / DEFAULT_DENOISER_PATH
    if output.exists() and not force:
        print(f"skip existing {output}")
        return output
    output.parent.mkdir(parents=True, exist_ok=True)
    download_file(DEFAULT_DENOISER_URL, output)
    return output


def download_file(url: str, output: Path) -> None:
    tmp_output = output.with_suffix(output.suffix + ".part")
    print(f"download {url}")
    if shutil.which("curl"):
        subprocess.run(
            [
                "curl",
                "-L",
                "--fail",
                "--retry",
                "3",
                "--connect-timeout",
                "20",
                "--max-time",
                "600",
                "-o",
                str(tmp_output),
                url,
            ],
            check=True,
        )
    else:
        with urllib.request.urlopen(url, timeout=600) as response, tmp_output.open("wb") as handle:
            shutil.copyfileobj(response, handle)
    tmp_output.replace(output)
    print(f"wrote {output} ({output.stat().st_size} bytes)")


def extract_tar_bz2(archive: Path, target_dir: Path) -> None:
    with tarfile.open(archive, "r:bz2") as tar:
        members = tar.getmembers()
        resolved_target = target_dir.resolve()
        for member in members:
            destination = (target_dir / member.name).resolve()
            if resolved_target != destination and resolved_target not in destination.parents:
                raise RuntimeError(f"refusing unsafe archive member: {member.name}")
        tar.extractall(target_dir, members=members)


if __name__ == "__main__":
    raise SystemExit(main())