"""Prebuilt audiocpp_server installs from audio.cpp's GitHub releases.

Upstream (0xShug0/audio.cpp) publishes ready-to-run ``audiocpp_server``
binaries for every release, built with ``AUDIOCPP_DEPLOYMENT_BUILD=ON``
(the model-spec catalog is embedded, so the binary is self-contained):

- ``audio-v<tag>-bin-macos-arm64-metal.tar.gz`` — Apple Silicon, real Metal
- ``audio-v<tag>-bin-macos-x64-metal.tar.gz`` — Intel Macs (upstream's CI
  Intel runners have no usable GPU, so this asset is CPU-only despite the
  name; either way the tool records macOS installs as the ``cpu`` backend)
- ``audio-v<tag>-bin-windows-x64-{cpu,vulkan,cuda12.4,cuda13.3}.zip`` —
  Windows; the CUDA variants additionally need the matching
  ``audio-v<tag>-cudart-windows-x64-cuda<ver>.zip`` runtime DLLs because
  those builds ship the CUDA backend as ``ggml-cuda.dll``
  (``ENGINE_ENABLE_CPU_ALL_VARIANTS``), loaded from the binary's directory.

Installing a release into the same ``build/<dir>/bin/`` layout the source
builds use means binary discovery (``find_audiocpp_server_bin`` /
``built_server_binary``), the ``-metal-`` -> ``cpu`` backend token mapping,
the managed-server launch, and the model-manager tooling from the git
checkout all work unchanged. A ``prebuilt.json`` marker in the build
directory records what was installed so ``build.update()`` can re-download
a newer release instead of rebuilding from source.

Everything downloads through urllib (no new dependencies); the GitHub API
publishes a ``sha256:<hex>`` digest per asset, which is verified before
extraction. Files saved by urllib carry no macOS quarantine attribute, so
the ad-hoc-signed binaries run without Gatekeeper prompts.
"""

import hashlib
import json
import os
import platform as _platform
import re
import shutil
import sys
import tarfile
import tempfile
import time
import urllib.error
import urllib.request
import zipfile
from pathlib import Path
from typing import List, Optional, Tuple

from backends import common

RELEASES_API_URL = \
    "https://api.github.com/repos/0xShug0/audio.cpp/releases/latest"

# GitHub rejects API/HTTP requests without a User-Agent header.
USER_AGENT = "tts-audiobook-generator (+https://github.com/0xShug0/audio.cpp)"

API_TIMEOUT = 30

# The CUDA runtime variant selected when the driver version cannot be
# detected: CUDA 12.4 runs on the widest range of installed drivers.
DEFAULT_CUDA_VARIANT = "12.4"

# nvidia-smi driver major version from which the CUDA 13.x builds run.
CUDA_13_MIN_DRIVER = 580

_ASSET_RE = re.compile(
    r"^audio-v[^/]+-bin-(?P<platform>macos|windows)-(?P<arch>arm64|x64)"
    r"-(?P<variant>[a-z0-9.]+)\.(?P<ext>tar\.gz|zip)$")

_CUDART_RE = re.compile(
    r"^audio-v[^/]+-cudart-windows-x64-cuda(?P<variant>[a-z0-9.]+)\.zip$")

# Download progress reporting: emit a line at most every 4 MiB.
_REPORT_STEP = 4 * 1024 * 1024


def prebuilt_supported(backend: Optional[str], *,
                       platform: Optional[str] = None) -> bool:
    """Whether a prebuilt release asset can exist for PLATFORM/BACKEND.

    Pure metadata (no network): macOS releases exist for the ``cpu``
    backend only (Metal is recorded as ``cpu``), Windows releases exist
    for ``cpu``, ``vulkan``, and ``cuda`` (HIP has no release asset and
    keeps building from source), and no other platform ships prebuilt
    binaries — Linux keeps using ``build_linux.sh``.
    """
    if platform is None:
        platform = sys.platform
    if platform == "darwin":
        return backend == "cpu"
    if platform == "win32":
        return backend in ("cpu", "vulkan", "cuda")
    return False


def _arch_token(machine: Optional[str] = None) -> Optional[str]:
    """The release asset's arch token for this machine, or None."""
    if machine is None:
        machine = _platform.machine()
    machine = machine.lower()
    if machine in ("arm64", "aarch64"):
        return "arm64"
    if machine in ("x86_64", "amd64"):
        return "x64"
    return None


def _platform_token(platform: Optional[str] = None) -> Optional[str]:
    if platform is None:
        platform = sys.platform
    if platform == "darwin":
        return "macos"
    if platform == "win32":
        return "windows"
    return None


def _default_cuda_variant() -> str:
    """The CUDA runtime variant to download, from the installed driver.

    ``nvidia-smi`` reports the driver version; CUDA 13.x builds need a
    580+ driver, older (or undetectable) drivers get the CUDA 12.4 build.
    """
    try:
        proc = common.run_console_subprocess_quiet(
            ["nvidia-smi", "--query-gpu=driver_version",
             "--format=csv,noheader"], timeout=10)
    except Exception:
        return DEFAULT_CUDA_VARIANT
    if proc is None or proc.returncode != 0:
        return DEFAULT_CUDA_VARIANT
    for line in proc.stdout.decode("utf-8", errors="replace").splitlines():
        major = line.strip().split(".")[0]
        if major.isdigit():
            return "13.3" if int(major) >= CUDA_13_MIN_DRIVER \
                else DEFAULT_CUDA_VARIANT
    return DEFAULT_CUDA_VARIANT


def install_dir(audiocpp_dir: Path, backend: Optional[str], *,
                platform: Optional[str] = None) -> Optional[Path]:
    """The ``bin`` directory a prebuilt install for BACKEND lands in.

    Mirrors the source-build directory names so backend detection and
    binary discovery keep working: ``build/macos-metal-release/bin`` on
    macOS and ``build/windows-<backend>-release/bin`` on Windows. None
    when the platform/backend has no prebuilt asset (see
    ``prebuilt_supported``).
    """
    token = _platform_token(platform)
    if token == "macos":
        if backend != "cpu":
            return None
        return audiocpp_dir / "build" / "macos-metal-release" / "bin"
    if token == "windows":
        if backend not in ("cpu", "vulkan", "cuda"):
            return None
        return (audiocpp_dir / "build"
                / f"windows-{backend}-release" / "bin")
    return None


def marker_path(audiocpp_dir: Path, backend: Optional[str], *,
                platform: Optional[str] = None) -> Optional[Path]:
    """The ``prebuilt.json`` marker path for this install, or None."""
    bin_dir = install_dir(audiocpp_dir, backend, platform=platform)
    if bin_dir is None:
        return None
    return bin_dir.parent / "prebuilt.json"


def installed_release(audiocpp_dir: Path, backend: Optional[str], *,
                      platform: Optional[str] = None) -> Optional[dict]:
    """What prebuilt release is installed for BACKEND, or None.

    Reads the ``prebuilt.json`` marker written by ``install_prebuilt``
    (``{"tag", "asset", "sha256", "installed_at"}``). Any unreadable or
    missing marker means "not a prebuilt install" — a source build lives
    there instead and ``update()`` keeps rebuilding it.
    """
    path = marker_path(audiocpp_dir, backend, platform=platform)
    if path is None or not path.is_file():
        return None
    try:
        data = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, ValueError):
        return None
    if not isinstance(data, dict) or not data.get("tag"):
        return None
    return data


def fetch_latest_release(*, timeout: int = API_TIMEOUT) -> Optional[dict]:
    """The latest audio.cpp GitHub release (API JSON), or None.

    Unauthenticated requests suffice (60/hour); callers treat None as
    "cannot check right now" and keep whatever is installed.
    """
    request = urllib.request.Request(
        RELEASES_API_URL,
        headers={"User-Agent": USER_AGENT,
                 "Accept": "application/vnd.github+json"})
    try:
        with urllib.request.urlopen(request, timeout=timeout) as response:
            return json.loads(response.read().decode("utf-8"))
    except (OSError, ValueError) as exc:
        print(f"[ERROR] Could not reach {RELEASES_API_URL}: {exc}")
        return None


def select_assets(assets: List[dict], backend: Optional[str], *,
                  platform: Optional[str] = None,
                  machine: Optional[str] = None,
                  cuda_variant: Optional[str] = None
                  ) -> Optional[Tuple[dict, Optional[dict]]]:
    """Pick the release assets for this machine and BACKEND.

    Returns ``(main, extra)`` — the server binary archive plus the CUDA
    runtime DLL archive when BACKEND is ``cuda`` — or None when no asset
    matches (unsupported platform, HIP, an unexpected arch). The CUDA
    variant defaults to the driver-appropriate one (``_default_cuda_variant``).
    """
    token = _platform_token(platform)
    arch = _arch_token(machine)
    if token is None or arch is None or not prebuilt_supported(
            backend, platform=platform):
        return None
    variant: Optional[str]
    cuda: Optional[str] = None
    if token == "macos":
        variant = "metal"
    elif backend == "cuda":
        # The binary asset names carry a "cuda" prefix on the variant
        # token (``...-bin-windows-x64-cuda12.4.zip``); the cudart asset
        # names use the bare version (``...-cudart-...-cuda12.4.zip``).
        cuda = cuda_variant or _default_cuda_variant()
        variant = "cuda" + cuda
    else:
        variant = backend
    main = _find_asset(assets, token, arch, variant)
    if main is None:
        return None
    extra = None
    if token == "windows" and backend == "cuda":
        extra = _find_cudart_asset(assets, cuda or "")
        if extra is None:
            return None
    return main, extra


def _find_asset(assets: List[dict], platform_token: str, arch: str,
                variant: str) -> Optional[dict]:
    for asset in assets:
        match = _ASSET_RE.match(str(asset.get("name", "")))
        if match and match.group("platform") == platform_token \
                and match.group("arch") == arch \
                and match.group("variant") == variant:
            return asset
    return None


def _find_cudart_asset(assets: List[dict], variant: str) -> Optional[dict]:
    for asset in assets:
        match = _CUDART_RE.match(str(asset.get("name", "")))
        if match and match.group("variant") == variant:
            return asset
    return None


def _download(url: str, dest: Path, *, emit=None, cancel=None) -> int:
    """Stream URL to DEST with progress lines, honoring CANCEL.

    Returns 0 on success, 130 when cancelled (the partial file is
    removed), 1 on any network or filesystem error. EMIT receives a
    progress line at most every 4 MiB (the in-TUI task view sink when
    set; console paths print nothing until the install summary).
    """
    say = emit if emit is not None else print
    request = urllib.request.Request(url, headers={"User-Agent": USER_AGENT})
    try:
        with urllib.request.urlopen(request, timeout=API_TIMEOUT) as response:
            total = response.headers.get("Content-Length")
            total = int(total) if total else None
            with dest.open("wb") as fh:
                downloaded = 0
                next_report = _REPORT_STEP
                while True:
                    if common.cancel_requested(cancel):
                        fh.close()
                        dest.unlink(missing_ok=True)
                        say("[INFO] Download cancelled.")
                        return 130
                    chunk = response.read(1 << 20)
                    if not chunk:
                        break
                    fh.write(chunk)
                    downloaded += len(chunk)
                    if emit is not None and downloaded >= next_report:
                        next_report += _REPORT_STEP
                        if total:
                            say(f"[INFO] Downloading: "
                                f"{downloaded / (1 << 20):.1f} / "
                                f"{total / (1 << 20):.1f} MB "
                                f"({100 * downloaded // total}%)")
                        else:
                            say(f"[INFO] Downloading: "
                                f"{downloaded / (1 << 20):.1f} MB")
        if emit is not None:
            say(f"[INFO] Downloaded {dest.name} "
                f"({dest.stat().st_size / (1 << 20):.1f} MB).")
        return 0
    except urllib.error.HTTPError as exc:
        print(f"[ERROR] Download failed ({exc.code} {exc.reason}): {url}")
        dest.unlink(missing_ok=True)
        return 1
    except OSError as exc:
        print(f"[ERROR] Download failed: {exc}")
        dest.unlink(missing_ok=True)
        return 1


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


def _verify_checksum(path: Path, asset: dict, *, emit=None) -> bool:
    """Verify PATH against the asset's published sha256 digest.

    The GitHub API publishes ``digest: "sha256:<hex>"``; a missing digest
    only warns (older releases), a mismatched one fails the install.
    """
    expected = str(asset.get("digest") or "")
    if not expected.startswith("sha256:"):
        print(f"[WARNING] {path.name} has no published checksum; "
              "skipping verification")
        return True
    actual = _sha256(path)
    if actual != expected[len("sha256:"):]:
        (emit if emit is not None else print)(
            f"[ERROR] {path.name}: checksum mismatch "
            f"(expected {expected}, got sha256:{actual}); not installing "
            "a tampered or corrupted download.")
        return False
    (emit if emit is not None else print)(
        f"[OK] {path.name}: checksum verified.")
    return True


def _validate_member(name: str) -> str:
    """Normalize an archive member name; raise ValueError on traversal."""
    normalized = name.replace("\\", "/")
    path = Path(normalized)
    if path.is_absolute() or ".." in path.parts or path.drive:
        raise ValueError(f"unsafe archive member: {name!r}")
    return normalized


def _extract(archive: Path, dest: Path, *, emit=None) -> None:
    """Extract a release archive into DEST (tar.gz or zip, safely)."""
    say = emit if emit is not None else print
    say(f"[INFO] Extracting {archive.name} into {dest}...")
    if tarfile.is_tarfile(archive):
        with tarfile.open(archive) as tf:
            try:
                tf.extractall(dest, filter="data")
            except TypeError:  # Python < 3.12: no filter= parameter
                for member in tf.getmembers():
                    _validate_member(member.name)
                tf.extractall(dest)
        return
    with zipfile.ZipFile(archive) as zf:
        for name in zf.namelist():
            _validate_member(name)
        zf.extractall(dest)


def _sync_checkout_to_release(checkout: Path, tag: str, *,
                              emit=None, cancel=None) -> None:
    """Point the checkout's tracked files at the installed release tag.

    The prebuilt binary embeds the release tag's model-spec catalog; the
    wizard reads model_specs/ and runs tools/model_manager_v2.py from the
    checkout, so detaching to the same tag keeps the tooling, the model
    catalog, and the binary consistent. Non-fatal: a failure (offline
    fetch, local edits) leaves the checkout on its current commit — the
    binary still works, worst case the model list drifts slightly from
    what the binary knows.
    """
    fetch_rc = common.run_console_subprocess(
        ["git", "-C", str(checkout), "fetch", "--tags", "origin"],
        emit=emit, cancel=cancel)
    if fetch_rc != 0:
        print(f"[WARNING] Could not fetch audio.cpp tags (exit {fetch_rc}); "
              f"the checkout stays on its current commit (binary is {tag}).")
        return
    detach_rc = common.run_console_subprocess(
        ["git", "-C", str(checkout), "checkout", "--detach", tag],
        emit=emit, cancel=cancel)
    if detach_rc != 0:
        print(f"[WARNING] Could not check out {tag} (exit {detach_rc}); "
              "the checkout's model catalog may differ from the binary's.")


def _ensure_executables(bin_dir: Path) -> None:
    """Make the extracted binaries executable (POSIX only).

    GitHub's artifact round-trip does not preserve tar permission bits,
    so the released binaries can land mode 644; the hub launches
    ``audiocpp_server`` directly, which then fails with EACCES. Every
    extension-less ``audiocpp_*`` entry in the bin directory (server,
    cli, gguf converter) gets the exec bits; data files (metallib,
    model_specs/, tools/) are untouched.
    """
    if os.name != "posix":
        return
    for entry in bin_dir.iterdir():
        if entry.is_file() and not entry.suffix \
                and entry.name.startswith("audiocpp_"):
            entry.chmod(entry.stat().st_mode | 0o111)


def install_prebuilt(audiocpp_dir: Path, backend: Optional[str], *,
                     emit=None, cancel=None,
                     cuda_variant: Optional[str] = None) -> int:
    """Download and install the latest prebuilt audiocpp_server for BACKEND.

    Fetches the latest release, picks this machine's asset(s), downloads
    them to a temp directory, verifies the published sha256 digests, then
    extracts into ``build/<dir>/bin/`` (replacing any previous content —
    including a source build there), syncs the checkout to the release
    tag, and writes the ``prebuilt.json`` marker. Streaming output and
    cancellation behave like the other install steps. Returns 0 on
    success, 130 when cancelled, 1 on any failure.
    """
    say = emit if emit is not None else print
    bin_dir = install_dir(audiocpp_dir, backend)
    if bin_dir is None:
        print(f"[ERROR] No prebuilt audiocpp_server asset exists for "
              f"{backend!r} on this platform; build from source instead.")
        return 1
    say("[INFO] Checking the latest audio.cpp release...")
    release = fetch_latest_release()
    if release is None:
        return 1
    tag = str(release.get("tag_name") or "?")
    pair = select_assets(release.get("assets") or [], backend,
                         cuda_variant=cuda_variant)
    if pair is None:
        print(f"[ERROR] No prebuilt audiocpp_server asset for "
              f"{backend!r} in release {tag}; build from source instead.")
        return 1
    main, extra = pair
    existing = installed_release(audiocpp_dir, backend)
    if existing and existing.get("tag") == tag \
            and existing.get("asset") == main.get("name"):
        say(f"[OK] Prebuilt audiocpp_server is already {tag}.")
        return 0

    tmp = Path(tempfile.mkdtemp(prefix="audiocpp-prebuilt-"))
    try:
        downloads: List[Tuple[dict, Path]] = []
        for asset in (main, extra):
            if asset is None:
                continue
            dest = tmp / str(asset.get("name"))
            say(f"[INFO] Downloading {asset.get('name')} "
                f"({int(asset.get('size') or 0) / (1 << 20):.1f} MB) "
                f"from release {tag}...")
            rc = _download(str(asset.get("browser_download_url")), dest,
                           emit=emit, cancel=cancel)
            if rc != 0:
                return rc
            downloads.append((asset, dest))
        for asset, dest in downloads:
            if not _verify_checksum(dest, asset, emit=emit):
                return 1
        if common.cancel_requested(cancel):
            say("[INFO] Download cancelled.")
            return 130
        build_dir = bin_dir.parent
        if build_dir.is_dir():
            shutil.rmtree(build_dir, ignore_errors=True)
        bin_dir.mkdir(parents=True, exist_ok=True)
        for _asset, dest in downloads:
            _extract(dest, bin_dir, emit=emit)
        _ensure_executables(bin_dir)
        marker = {
            "tag": tag,
            "asset": main.get("name"),
            "sha256": str(main.get("digest") or ""),
            "installed_at": time.strftime("%Y-%m-%dT%H:%M:%SZ",
                                          time.gmtime()),
        }
        (build_dir / "prebuilt.json").write_text(
            json.dumps(marker, indent=2) + "\n", encoding="utf-8")
        if sys.platform == "darwin" \
                and _arch_token() == "x64":
            say("[INFO] Note: upstream's Intel macOS release is built "
                "without Metal (CI limitation); it runs on CPU.")
        say(f"[OK] audiocpp_server {tag} installed at {bin_dir}.")
        _sync_checkout_to_release(audiocpp_dir, tag, emit=emit, cancel=cancel)
        return 0
    finally:
        shutil.rmtree(tmp, ignore_errors=True)