"""Pinned, resumable asset setup for ResearchMate's pocket-tts default voice and Piper fallback voice."""

from __future__ import annotations

import argparse
import hashlib
import json
from dataclasses import asdict, dataclass
from pathlib import Path


@dataclass(frozen=True)
class DownloadSpec:
    name: str
    repo_id: str
    revision: str
    destination: str
    allow_patterns: tuple[str, ...] = ()


def download_specs() -> list[DownloadSpec]:
    return [
        DownloadSpec(
            "pocket-tts",
            "kyutai/pocket-tts",
            "4c8ad48f8a003909bc4f1122cbe88a4252124621",
            "pocket",
        ),
        DownloadSpec(
            "piper-en-us",
            "rhasspy/piper-voices",
            "e21c7de8d4eab79b902f0d61e662b3f21664b8d2",
            "piper",
            (
                "en/en_US/lessac/medium/en_US-lessac-medium.onnx",
                "en/en_US/lessac/medium/en_US-lessac-medium.onnx.json",
                "en/en_US/lessac/medium/MODEL_CARD",
            ),
        ),
    ]


def setup_assets(root: Path) -> dict:
    from huggingface_hub import snapshot_download

    root.mkdir(parents=True, exist_ok=True)
    rows: list[dict] = []
    for spec in download_specs():
        destination = root / spec.destination
        try:
            snapshot_download(
                repo_id=spec.repo_id,
                revision=spec.revision,
                allow_patterns=list(spec.allow_patterns) or None,
                local_dir=destination,
            )
            rows.append({
                **asdict(spec),
                "status": "installed",
                "bytes": _directory_size(destination),
                "checksums": _checksums(destination),
            })
        except Exception as exc:
            rows.append({**asdict(spec), "status": "error", "error": str(exc)})

    manifest = {"schema_version": 1, "root": str(root), "assets": rows, "total_bytes": _directory_size(root)}
    (root / "installation-manifest.json").write_text(json.dumps(manifest, indent=2), encoding="utf-8")
    return manifest


def _directory_size(path: Path) -> int:
    return sum(item.stat().st_size for item in path.rglob("*") if item.is_file()) if path.exists() else 0


def _sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def _checksums(path: Path) -> dict[str, str]:
    return {
        item.relative_to(path).as_posix(): _sha256(item)
        for item in path.rglob("*") if item.is_file() and not item.name.startswith(".cache")
    }


def main() -> None:
    parser = argparse.ArgumentParser(description="Download pinned Pocket TTS and Piper assets.")
    parser.add_argument("--root", type=Path, default=Path.home() / ".studybuddy" / "models")
    parser.add_argument("--yes", action="store_true", help="Accept the model download without prompting")
    args = parser.parse_args()
    if not args.yes:
        answer = input("Download Pocket TTS and Piper assets? [y/N] ")
        if answer.strip().lower() not in {"y", "yes"}:
            return
    manifest = setup_assets(args.root)
    print(json.dumps(manifest, indent=2))


if __name__ == "__main__":
    main()