"""Load model pack from path or CDN. Internal use by WakeWordDetector."""

from __future__ import annotations

import asyncio
import hashlib
import io
import os
import tarfile
from pathlib import Path

import httpx

CDN_BASE = "https://cdn.sonnetics.com/models"
CHUNK_SIZE = 2048  # matches sonnetics-js / Rust engine


def _normalize_path(p: str) -> str:
    """Strip leading ./ so keys match Rust expectations."""
    return p[2:] if p.startswith("./") else p


def _cache_dir() -> Path:
    if os.name == "nt":
        base = Path(os.environ.get("APPDATA", Path.home() / "AppData" / "Roaming"))
    else:
        base = Path.home() / ".cache"
    return base / "sonnetics" / "cache"


def _cache_path(key: str) -> Path:
    h = hashlib.sha256(key.encode()).hexdigest()
    return _cache_dir() / h


def _extract_tar_gz(tar_gz: bytes) -> dict[str, bytes]:
    """Extract tar.gz into dict of path -> bytes."""
    with tarfile.open(fileobj=io.BytesIO(tar_gz), mode="r:gz") as tf:
        out: dict[str, bytes] = {}
        for member in tf.getmembers():
            if member.isfile():
                f = tf.extractfile(member)
                if f:
                    out[_normalize_path(member.name)] = f.read()
        return out


def _get_cached(key: str) -> dict[str, bytes] | None:
    dir_path = _cache_path(key)
    if not dir_path.exists():
        return None
    out: dict[str, bytes] = {}
    for p in dir_path.rglob("*"):
        if p.is_file():
            rel = p.relative_to(dir_path).as_posix()
            out[rel] = p.read_bytes()
    return out if out else None


def _set_cached(key: str, files: dict[str, bytes]) -> None:
    dir_path = _cache_path(key)
    dir_path.mkdir(parents=True, exist_ok=True)
    for path, data in files.items():
        full = dir_path / path
        full.parent.mkdir(parents=True, exist_ok=True)
        full.write_bytes(data)


async def _load_from_url(url: str) -> dict[str, bytes]:
    cached = _get_cached(url)
    if cached:
        return cached
    async with httpx.AsyncClient() as client:
        resp = await client.get(url)
        resp.raise_for_status()
        tar_gz = resp.content
    files = _extract_tar_gz(tar_gz)
    _set_cached(url, files)
    return files


def _load_from_path(path: str | Path) -> dict[str, bytes]:
    path = Path(path)
    if not path.name.endswith(".tar.gz"):
        raise ValueError("Path must point to a .tar.gz file")
    key = f"file:{path}"
    cached = _get_cached(key)
    if cached:
        return cached
    files = _extract_tar_gz(path.read_bytes())
    _set_cached(key, files)
    return files


def _load_model_sync(
    *,
    path: str | Path | None = None,
    model_id: str | None = None,
) -> dict[str, bytes]:
    """Load model from path or model_id. Exactly one required. Sync (blocks on CDN fetch)."""
    if (path is None) == (model_id is None):
        raise ValueError("Provide exactly one of path or model_id")
    if path is not None:
        return _load_from_path(path)
    assert model_id is not None
    url = f"{CDN_BASE}/sonnetics-model-{model_id}.tar.gz"
    return asyncio.run(_load_from_url(url))


async def _load_model_async(
    *,
    path: str | Path | None = None,
    model_id: str | None = None,
) -> dict[str, bytes]:
    """Load model from path or model_id. Exactly one required. Async."""
    if (path is None) == (model_id is None):
        raise ValueError("Provide exactly one of path or model_id")
    if path is not None:
        return _load_from_path(path)
    assert model_id is not None
    url = f"{CDN_BASE}/sonnetics-model-{model_id}.tar.gz"
    return await _load_from_url(url)
