"""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. When the API is rate-limited, the release tag is resolved from
the release page's redirect instead (no API quota) and the download
proceeds without verification (with a warning). 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"

# The release pages (NOT the API): following /releases/latest's redirect
# reveals the newest tag without spending any of the API's rate limit
# (60 requests/hour unauthenticated, shared per IP), and the asset
# download URLs are deterministic from the tag.
RELEASES_LATEST_URL = "https://github.com/0xShug0/audio.cpp/releases/latest"
RELEASES_DOWNLOAD_URL = "https://github.com/0xShug0/audio.cpp/releases/download"

# 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$")

_TAG_RE = re.compile(r"/releases/tag/(v[^/?#]+)")

# 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 resolve_latest_tag(*, timeout: int = API_TIMEOUT) -> Optional[str]:
    """The latest release tag, resolved WITHOUT the GitHub API.

    ``/releases/latest`` redirects to ``/releases/tag/vX.Y.Z``; reading
    the final URL is an ordinary page hit that costs none of the API's
    rate limit, so "already current" checks and rate-limited installs
    never depend on the API. Returns None when the redirect cannot be
    followed (offline, unexpected page shape).
    """
    request = urllib.request.Request(RELEASES_LATEST_URL,
                                     headers={"User-Agent": USER_AGENT})
    try:
        with urllib.request.urlopen(request, timeout=timeout) as response:
            url = response.geturl() or ""
    except OSError:
        return None
    match = _TAG_RE.search(url)
    return match.group(1) if match else None


def _variant_tokens(backend: Optional[str], *,
                    platform: Optional[str] = None,
                    machine: Optional[str] = None,
                    cuda_variant: Optional[str] = None
                    ) -> Optional[Tuple[str, str, str, Optional[str]]]:
    """The (platform, arch, variant, cudart) tokens for this machine.

    Returns None when this platform/machine/backend combination has no
    release asset naming (see ``prebuilt_supported``). The binary asset's
    variant token carries a "cuda" prefix for CUDA builds
    (``...-bin-windows-x64-cuda12.4.zip``) while the cudart archive uses
    the bare version (``...-cudart-...-cuda12.4.zip``).
    """
    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
    if token == "macos":
        return token, arch, "metal", None
    if backend == "cuda":
        cuda = cuda_variant or _default_cuda_variant()
        return token, arch, "cuda" + cuda, cuda
    return token, arch, backend, 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``).
    """
    tokens = _variant_tokens(backend, platform=platform, machine=machine,
                             cuda_variant=cuda_variant)
    if tokens is None:
        return None
    token, arch, variant, cuda = tokens
    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 synthesize_assets(backend: Optional[str], tag: str, *,
                      platform: Optional[str] = None,
                      machine: Optional[str] = None,
                      cuda_variant: Optional[str] = None
                      ) -> Optional[Tuple[dict, Optional[dict]]]:
    """Asset entries built from the tag alone, without the GitHub API.

    The rate-limited fallback of ``install_prebuilt``: asset names and
    ``releases/download/<tag>/`` URLs are deterministic, but the sha256
    digests are not — callers download these WITHOUT checksum
    verification and must say so. Same (main, extra) contract as
    ``select_assets``.
    """
    tokens = _variant_tokens(backend, platform=platform, machine=machine,
                             cuda_variant=cuda_variant)
    if tokens is None:
        return None
    token, arch, variant, cuda = tokens

    def asset(name: str) -> dict:
        return {"name": name, "size": None, "digest": None,
                "browser_download_url":
                    f"{RELEASES_DOWNLOAD_URL}/{tag}/{name}"}

    ext = ".tar.gz" if token == "macos" else ".zip"
    main = asset(f"audio-{tag}-bin-{token}-{arch}-{variant}{ext}")
    extra = None
    if token == "windows" and backend == "cuda":
        extra = asset(f"audio-{tag}-cudart-windows-{arch}-cuda{cuda}.zip")
    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.

    Resolution degrades gracefully when GitHub's API is unavailable (its
    unauthenticated limit is 60 requests/hour per IP): normally the API
    supplies the release assets and their sha256 digests; when it is
    rate-limited the newest tag is resolved from the release page's
    redirect (no API quota), the asset names are synthesized, and the
    download proceeds WITHOUT checksum verification (with a loud
    warning). Only when neither path can resolve the release does the
    install fail. The archive(s) are downloaded to a temp directory,
    verified, then extracted into ``build/<dir>/bin/`` (replacing any
    previous content — including a source build there), the checkout is
    synced to the release tag, and the ``prebuilt.json`` marker written.
    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()
    pair = select_assets(release.get("assets") or [], backend,
                         cuda_variant=cuda_variant) if release else None
    if release is not None and pair is None:
        # API reachable but nothing matches: a hard mismatch (unsupported
        # backend, renamed assets) — do not paper over it with a
        # synthesized name that would just 404.
        print(f"[ERROR] No prebuilt audiocpp_server asset for "
              f"{backend!r} in release {release.get('tag_name')}; build "
              "from source instead.")
        return 1
    if pair is not None:
        tag = str(release.get("tag_name") or "?")
        main, extra = pair
    else:
        tag = resolve_latest_tag()
        if tag is None:
            print("[ERROR] Could not reach GitHub to resolve the latest "
                  "audio.cpp release (rate limit or offline); try again "
                  "later or build from source.")
            return 1
        pair = synthesize_assets(backend, tag, cuda_variant=cuda_variant)
        if pair is None:
            print(f"[ERROR] No prebuilt audiocpp_server asset exists for "
                  f"{backend!r} in release {tag}; build from source "
                  "instead.")
            return 1
        main, extra = pair
        say("[WARNING] GitHub's release API is unreachable (rate limit?) "
            f"— installing {tag} WITHOUT checksum verification.")
    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)