"""
ROCm backend download, assembly, and verification.

Downloads two archives from GitHub Releases:
  1. Server core (voicebox-server-rocm.tar.gz) — the exe + non-AMD deps,
     versioned with the app.
  2. ROCm libs (rocm-libs-{version}.tar.gz) — AMD runtime libraries,
     versioned independently (only redownloaded on ROCm toolkit bump).

Both archives are extracted into {data_dir}/backends/rocm/ which forms the
complete PyInstaller --onedir directory structure that torch expects.
"""

import asyncio
import hashlib
import json
import logging
import os
import shutil
import sys
import tarfile
from pathlib import Path

from .. import __version__
from ..config import get_data_dir
from ..utils.progress import get_progress_manager

logger = logging.getLogger(__name__)

GITHUB_RELEASES_URL = "https://github.com/jamiepine/voicebox/releases/download"

PROGRESS_KEY = "rocm-backend"

# The current expected ROCm libs version.  Bump this when we change the
# ROCm toolkit version or torch's ROCm dependency changes (e.g. rocm7.2 -> rocm7.4).
ROCM_LIBS_VERSION = "rocm7.2-v1"

# Prevents concurrent download_rocm_binary() calls from racing on the same
# temp file.  The auto-update background task and the manual HTTP endpoint
# can both invoke download_rocm_binary(); without this lock the progress-
# manager status check is a TOCTOU race.
_download_lock = asyncio.Lock()


def get_backends_dir() -> Path:
    """Directory where downloaded backend binaries are stored."""
    d = get_data_dir() / "backends"
    d.mkdir(parents=True, exist_ok=True)
    return d


def get_rocm_dir() -> Path:
    """Directory where the ROCm backend (onedir) is extracted."""
    d = get_backends_dir() / "rocm"
    d.mkdir(parents=True, exist_ok=True)
    return d


def get_rocm_exe_name() -> str:
    """Platform-specific ROCm executable filename."""
    if sys.platform == "win32":
        return "voicebox-server-rocm.exe"
    return "voicebox-server-rocm"


def get_rocm_binary_path() -> Path | None:
    """Return path to the ROCm executable if it exists inside the onedir."""
    p = get_rocm_dir() / get_rocm_exe_name()
    if p.exists():
        return p
    return None


def get_rocm_libs_manifest_path() -> Path:
    """Path to the rocm-libs.json manifest inside the ROCm dir."""
    return get_rocm_dir() / "rocm-libs.json"


def get_installed_rocm_libs_version() -> str | None:
    """Read the installed ROCm libs version from rocm-libs.json, or None."""
    manifest_path = get_rocm_libs_manifest_path()
    if not manifest_path.exists():
        return None
    try:
        data = json.loads(manifest_path.read_text())
        return data.get("version")
    except Exception as e:
        logger.warning(f"Could not read rocm-libs.json: {e}")
        return None


def is_rocm_active() -> bool:
    """Check if the current process is the ROCm binary.

    The ROCm binary sets this env var on startup (see server.py).
    """
    return os.environ.get("VOICEBOX_BACKEND_VARIANT") == "rocm"


def get_rocm_status() -> dict:
    """Get current ROCm backend status for the API."""
    progress_manager = get_progress_manager()
    rocm_path = get_rocm_binary_path()
    progress = progress_manager.get_progress(PROGRESS_KEY)
    rocm_libs_version = get_installed_rocm_libs_version()

    return {
        "available": rocm_path is not None,
        "active": is_rocm_active(),
        "binary_path": str(rocm_path) if rocm_path else None,
        "rocm_libs_version": rocm_libs_version,
        "downloading": progress is not None and progress.get("status") == "downloading",
        "download_progress": progress,
    }


def _needs_server_download(version: str | None = None) -> bool:
    """Check if the server core archive needs to be (re)downloaded."""
    rocm_path = get_rocm_binary_path()
    if not rocm_path:
        return True
    # Check if the binary version matches the expected app version
    installed = get_rocm_binary_version()
    expected = version or __version__
    if expected.startswith("v"):
        expected = expected[1:]
    return installed != expected


def _needs_rocm_libs_download() -> bool:
    """Check if the ROCm libs archive needs to be (re)downloaded."""
    installed = get_installed_rocm_libs_version()
    if installed is None:
        return True
    return installed != ROCM_LIBS_VERSION


async def _download_and_extract_archive(
    client,
    url: str,
    sha256_url: str | None,
    dest_dir: Path,
    label: str,
    progress_offset: int,
    total_size: int,
):
    """Download a .tar.gz archive and extract it into dest_dir.

    Args:
        client: httpx.AsyncClient
        url: URL of the .tar.gz archive
        sha256_url: URL of the .sha256 checksum file (optional)
        dest_dir: Directory to extract into
        label: Human-readable label for progress updates
        progress_offset: Byte offset for progress reporting (when downloading
            multiple archives sequentially)
        total_size: Total bytes across all downloads (for progress bar)
    """
    progress = get_progress_manager()
    temp_path = dest_dir / f".download-{label.replace(' ', '-')}.tmp"

    # Clean up leftover partial download
    if temp_path.exists():
        temp_path.unlink()

    # Fetch expected checksum (fail-fast: never extract an unverified archive)
    expected_sha = None
    if sha256_url:
        try:
            sha_resp = await client.get(sha256_url)
            sha_resp.raise_for_status()
            expected_sha = sha_resp.text.strip().split()[0]
            logger.info(f"{label}: expected SHA-256: {expected_sha[:16]}...")
        except Exception as e:
            raise RuntimeError(f"{label}: failed to fetch checksum from {sha256_url}") from e

    # Stream download, verify, and extract — always clean up temp file
    downloaded = 0
    try:
        async with client.stream("GET", url) as response:
            response.raise_for_status()
            with open(temp_path, "wb") as f:
                async for chunk in response.aiter_bytes(chunk_size=1024 * 1024):
                    f.write(chunk)
                    downloaded += len(chunk)
                    progress.update_progress(
                        PROGRESS_KEY,
                        current=progress_offset + downloaded,
                        total=total_size,
                        filename=f"Downloading {label}",
                        status="downloading",
                    )

        # Verify integrity
        if expected_sha:
            progress.update_progress(
                PROGRESS_KEY,
                current=progress_offset + downloaded,
                total=total_size,
                filename=f"Verifying {label}...",
                status="downloading",
            )
            sha256 = hashlib.sha256()
            with open(temp_path, "rb") as f:
                while True:
                    data = f.read(1024 * 1024)
                    if not data:
                        break
                    sha256.update(data)
            actual = sha256.hexdigest()
            if actual != expected_sha:
                raise ValueError(
                    f"{label} integrity check failed: expected {expected_sha[:16]}..., got {actual[:16]}..."
                )
            logger.info(f"{label}: integrity verified")

        # Extract (use data filter for path traversal protection on Python 3.12+)
        progress.update_progress(
            PROGRESS_KEY,
            current=progress_offset + downloaded,
            total=total_size,
            filename=f"Extracting {label}...",
            status="downloading",
        )
        with tarfile.open(temp_path, "r:gz") as tar:
            tar.extractall(path=dest_dir, filter="data")

        logger.info(f"{label}: extracted to {dest_dir}")
    finally:
        if temp_path.exists():
            temp_path.unlink()
    return downloaded


async def download_rocm_binary(version: str | None = None):
    """Download the ROCm backend (server core + ROCm libs if needed).

    Downloads both archives from GitHub Releases, extracts them into
    {data_dir}/backends/rocm/, and writes the rocm-libs.json manifest.

    Only downloads what's needed:
    - Server core: always redownloaded (versioned with app)
    - ROCm libs: only if missing or version mismatch

    Args:
        version: Version tag (e.g. "v0.3.0"). Defaults to current app version.
    """
    if _download_lock.locked():
        logger.info("ROCm download already in progress, skipping duplicate request")
        return
    async with _download_lock:
        await _download_rocm_binary_locked(version)


async def _download_rocm_binary_locked(version: str | None = None):
    """Inner implementation of download_rocm_binary, called under _download_lock."""
    import httpx

    if version is None:
        version = f"v{__version__}"

    progress = get_progress_manager()
    rocm_dir = get_rocm_dir()

    need_server = _needs_server_download(version)
    need_libs = _needs_rocm_libs_download()

    if not need_server and not need_libs:
        logger.info("ROCm backend is up to date, nothing to download")
        return

    logger.info(
        f"Starting ROCm backend download for {version} "
        f"(server={'yes' if need_server else 'cached'}, "
        f"libs={'yes' if need_libs else 'cached'})"
    )
    progress.update_progress(
        PROGRESS_KEY,
        current=0,
        total=0,
        filename="Preparing download...",
        status="downloading",
    )

    # Server core and libs archive are both published under the app-version
    # release tag; the libs content version is encoded in the filename only.
    server_base_url = f"{GITHUB_RELEASES_URL}/{version}"
    libs_base_url = server_base_url
    server_archive = "voicebox-server-rocm.tar.gz"
    libs_archive = f"rocm-libs-{ROCM_LIBS_VERSION}.tar.gz"

    # Always stage when any download is needed, then atomically rename over
    # rocm_dir on success. This prevents a failed mid-extraction from leaving
    # rocm_dir in a partially-installed state that still passes the
    # get_rocm_binary_path() existence check. Existing files are pre-copied
    # into staging so partial updates (e.g. libs-only or server-only) preserve
    # whatever isn't being re-downloaded.
    use_staging = need_server or need_libs
    staging_dir = get_backends_dir() / "rocm-staging"

    if use_staging:
        if staging_dir.exists():
            shutil.rmtree(staging_dir)
        staging_dir.mkdir(parents=True, exist_ok=True)
        # Preserve existing files (server or libs) that don't need re-downloading.
        # Extracted archives will overwrite only what we actually download.
        if rocm_dir.exists():
            shutil.copytree(rocm_dir, staging_dir, dirs_exist_ok=True)
        extract_dir = staging_dir
    else:
        extract_dir = rocm_dir

    try:
        async with httpx.AsyncClient(follow_redirects=True, timeout=30.0) as client:
            # Estimate total download size
            total_size = 0
            if need_server:
                try:
                    head = await client.head(f"{server_base_url}/{server_archive}")
                    total_size += int(head.headers.get("content-length", 0))
                except Exception:
                    pass
            if need_libs:
                try:
                    head = await client.head(f"{libs_base_url}/{libs_archive}")
                    total_size += int(head.headers.get("content-length", 0))
                except Exception:
                    pass

            logger.info(f"Total download size: {total_size / 1024 / 1024:.1f} MB")

            offset = 0

            # Download server core
            if need_server:
                server_downloaded = await _download_and_extract_archive(
                    client,
                    url=f"{server_base_url}/{server_archive}",
                    sha256_url=f"{server_base_url}/{server_archive}.sha256",
                    dest_dir=extract_dir,
                    label="ROCm server",
                    progress_offset=offset,
                    total_size=total_size,
                )
                offset += server_downloaded

                # Make executable on Unix
                exe_path = extract_dir / get_rocm_exe_name()
                if sys.platform != "win32" and exe_path.exists():
                    exe_path.chmod(0o755)

            # Download ROCm libs
            if need_libs:
                await _download_and_extract_archive(
                    client,
                    url=f"{libs_base_url}/{libs_archive}",
                    sha256_url=f"{libs_base_url}/{libs_archive}.sha256",
                    dest_dir=extract_dir,
                    label="ROCm libraries",
                    progress_offset=offset,
                    total_size=total_size,
                )

                # Write local rocm-libs.json manifest
                manifest = {"version": ROCM_LIBS_VERSION}
                (extract_dir / "rocm-libs.json").write_text(json.dumps(manifest, indent=2) + "\n")

        # Atomic swap: replace rocm_dir with the fully-extracted staging dir
        if use_staging:
            backup_dir = get_backends_dir() / "rocm-backup"
            if backup_dir.exists():
                shutil.rmtree(backup_dir)
            if rocm_dir.exists():
                rocm_dir.rename(backup_dir)
            try:
                staging_dir.rename(rocm_dir)
            except Exception:
                if backup_dir.exists() and not rocm_dir.exists():
                    backup_dir.rename(rocm_dir)
                raise
            else:
                if backup_dir.exists():
                    shutil.rmtree(backup_dir)

        logger.info(f"ROCm backend ready at {rocm_dir}")
        progress.mark_complete(PROGRESS_KEY)

    except Exception as e:
        if use_staging and staging_dir.exists():
            shutil.rmtree(staging_dir)
        logger.error(f"ROCm backend download failed: {e}")
        progress.mark_error(PROGRESS_KEY, str(e))
        raise


def get_rocm_binary_version() -> str | None:
    """Get the version of the installed ROCm binary, or None if not installed."""
    import subprocess

    rocm_path = get_rocm_binary_path()
    if not rocm_path:
        return None
    try:
        result = subprocess.run(
            [str(rocm_path), "--version"],
            capture_output=True,
            text=True,
            timeout=30,
            cwd=str(rocm_path.parent),  # Run from the onedir directory
        )
        # Output format: "voicebox-server 0.3.0"
        for line in result.stdout.strip().splitlines():
            if "voicebox-server" in line:
                return line.split()[-1]
    except Exception as e:
        logger.warning(f"Could not get ROCm binary version: {e}")
    return None


async def check_and_update_rocm_binary():
    """Check if the ROCm binary is outdated and auto-download if so.

    Called on server startup. Checks both server version and ROCm libs
    version. Downloads only what's needed.
    """
    rocm_path = get_rocm_binary_path()
    if not rocm_path:
        return  # No ROCm binary installed, nothing to update

    if is_rocm_active():
        logger.info("ROCm backend is active; skipping auto-update to avoid replacing the running backend")
        return

    need_server = _needs_server_download()
    need_libs = _needs_rocm_libs_download()

    if not need_server and not need_libs:
        logger.info(f"ROCm binary is up to date (server=v{__version__}, libs={get_installed_rocm_libs_version()})")
        return

    reasons = []
    if need_server:
        rocm_version = get_rocm_binary_version()
        reasons.append(f"server v{rocm_version} != v{__version__}")
    if need_libs:
        installed_libs = get_installed_rocm_libs_version()
        reasons.append(f"libs {installed_libs} != {ROCM_LIBS_VERSION}")

    logger.info(f"ROCm backend needs update ({', '.join(reasons)}). Auto-downloading...")

    try:
        await download_rocm_binary()
    except Exception as e:
        logger.error(f"Auto-update of ROCm binary failed: {e}")


async def delete_rocm_binary() -> bool:
    """Delete the downloaded ROCm backend directory. Returns True if deleted."""
    import shutil

    rocm_dir = get_rocm_dir()
    if rocm_dir.exists() and any(rocm_dir.iterdir()):
        shutil.rmtree(rocm_dir)
        logger.info(f"Deleted ROCm backend directory: {rocm_dir}")
        return True
    return False