"""
downloader.py — asset download engine for comfy-resolve

Handles downloading model files from HuggingFace, Civitai, and direct URLs
with the following features:

- Concurrent downloads via ThreadPoolExecutor
- Live per-file progress bars via rich.progress
- Resume support for interrupted downloads (Range requests)
- Exponential back-off retry on transient network errors
- Rate-limit awareness (Retry-After header)
- SHA-256 integrity check — warns on mismatch, never silently deletes
- Disk-space pre-flight check
- Local and already-skipped assets are filtered before progress bars appear
"""

from __future__ import annotations

import hashlib
import logging
import os
import shutil
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass, field
from pathlib import Path
from threading import Lock
from typing import Any

import requests  # type: ignore
from rich.console import Console
from rich.progress import (
    BarColumn,
    DownloadColumn,
    Progress,
    TaskID,
    TextColumn,
    TimeRemainingColumn,
    TransferSpeedColumn,
)

from .sources.base import ResolutionResult, ResolvedSource

logger = logging.getLogger(__name__)
console = Console()


# ---------------------------------------------------------------------------
# Data classes
# ---------------------------------------------------------------------------

@dataclass
class DownloadTask:
    """A single file to be downloaded."""
    result: ResolutionResult
    destination: Path
    asset_type: str


@dataclass
class DownloadResult:
    """
    Outcome of a single download attempt.

    ``success`` is True even when a SHA-256 mismatch is detected — the file
    is kept and the caller is warned via ``sha256_warning``.  Use ``error``
    for hard failures only.
    """
    task: DownloadTask
    success: bool
    skipped: bool
    error: str | None = None
    sha256_warning: str | None = None   # set when hash mismatches but file is kept


# ---------------------------------------------------------------------------
# Downloader
# ---------------------------------------------------------------------------

class Downloader:
    """Downloads assets from HuggingFace, Civitai, and direct URLs."""

    def __init__(self, config: dict[str, Any]) -> None:
        """
        Initialise the downloader.

        Args:
            config: comfy-resolve configuration dictionary.
        """
        self.config      = config
        self.max_retries = config.get("download", {}).get("max_retries", 3)
        self.concurrent  = config.get("download", {}).get("concurrent", 2)

    # ── Public API ──────────────────────────────────────────────────────────

    def download_all(self, tasks: list[DownloadTask]) -> list[DownloadResult]:
        """
        Download all tasks, showing a live multi-bar progress display.

        Local and already-skipped tasks are resolved immediately and never
        appear in the progress display.  Only tasks that require a network
        transfer get a live progress row.

        Args:
            tasks: List of download tasks (may include local/skip entries).

        Returns:
            List of DownloadResult in completion order.
        """
        results: list[DownloadResult] = []
        lock = Lock()

        # Split into immediate (no download needed) vs network tasks up-front
        # so the progress display is never cluttered with pre-resolved items.
        immediate: list[DownloadTask] = []
        network:   list[DownloadTask] = []

        for t in tasks:
            src_type = t.result.source.source_type if t.result.source else None
            if src_type == "local" or self._should_skip(t):
                immediate.append(t)
            else:
                network.append(t)

        # Resolve immediate tasks without touching the progress display
        for t in immediate:
            results.append(DownloadResult(task=t, success=True, skipped=True))

        if not network:
            return results

        # Live progress display for network downloads
        with Progress(
            TextColumn("[bold cyan]{task.fields[filename]}", justify="right"),
            BarColumn(bar_width=40),
            DownloadColumn(),
            TransferSpeedColumn(),
            TimeRemainingColumn(),
            console=console,
            transient=False,
        ) as progress:

            def run(dl_task: DownloadTask) -> DownloadResult:
                prog_id = progress.add_task(
                    "", filename=dl_task.result.input_name, total=None
                )
                result = self.download_one(dl_task, progress, prog_id)
                with lock:
                    results.append(result)
                return result

            with ThreadPoolExecutor(max_workers=self.concurrent) as executor:
                futures = [executor.submit(run, t) for t in network]
                try:
                    for f in as_completed(futures):
                        f.result()  # re-raises any unexpected exception
                except KeyboardInterrupt:
                    console.print("\n[yellow]Download cancelled — stopping...[/yellow]")
                    executor.shutdown(wait=False, cancel_futures=True)
                    raise

        return results

    def download_one(
        self,
        task: DownloadTask,
        progress: Progress | None = None,
        prog_id: TaskID | None = None,
    ) -> DownloadResult:
        """
        Download a single task, driving *progress* if provided.

        Dispatches to the appropriate source-specific method based on
        ``task.result.source.source_type``.

        Args:
            task:     The download task.
            progress: Rich Progress instance (optional).
            prog_id:  Task ID within *progress* (optional).

        Returns:
            DownloadResult describing the outcome.
        """
        def _finish(
            success: bool,
            skipped: bool,
            error: str | None = None,
            sha256_warning: str | None = None,
        ) -> DownloadResult:
            if progress is not None and prog_id is not None:
                if skipped:
                    progress.update(prog_id, description="[yellow]SKIPPED[/yellow]")
                elif sha256_warning:
                    progress.update(prog_id, description="[yellow]⚠ HASH MISMATCH[/yellow]")
                elif success:
                    progress.update(prog_id, description="[green]✓ DONE[/green]")
                else:
                    progress.update(prog_id, description="[red]✗ FAILED[/red]")
                progress.stop_task(prog_id)
            return DownloadResult(
                task=task,
                success=success,
                skipped=skipped,
                error=error,
                sha256_warning=sha256_warning,
            )

        # Guard: should already be filtered, but handle defensively
        if self._should_skip(task):
            return _finish(True, True)

        if not task.result.source:
            return _finish(False, False, error="No source available")

        task.destination.mkdir(parents=True, exist_ok=True)

        source_type = task.result.source.source_type
        try:
            if source_type == "huggingface":
                return self._download_huggingface(task, progress, prog_id)
            elif source_type == "direct_url":
                return self._download_direct(task, progress, prog_id)
            elif source_type == "civitai":
                return self._download_civitai(task, progress, prog_id)
            elif source_type == "local":
                return _finish(True, True)
            else:
                return _finish(False, False, error=f"Unknown source type: {source_type}")
        except Exception as exc:
            logger.exception("Unexpected error downloading %s", task.result.input_name)
            return _finish(False, False, error=str(exc))

    # ── Source-specific downloaders ─────────────────────────────────────────

    def _download_huggingface(
        self,
        task: DownloadTask,
        progress: Progress | None = None,
        prog_id: TaskID | None = None,
    ) -> DownloadResult:
        """
        Stream a file directly from the HuggingFace CDN.

        Bypasses ``hf_hub_download`` entirely to avoid the double-copy
        cache behaviour introduced in recent versions of huggingface_hub,
        which caused apparent hangs on large files.

        Args:
            task:     Download task with a ``huggingface`` source.
            progress: Rich Progress instance (optional).
            prog_id:  Progress task ID (optional).

        Returns:
            DownloadResult describing the outcome.
        """
        source = task.result.source
        if not source:
            return DownloadResult(task=task, success=False, skipped=False,
                                  error="No source")

        # Empty-string tokens must not be sent as "Bearer " — that's an
        # illegal header value and causes an immediate protocol error.
        token = self.config.get("api_keys", {}).get("huggingface") or None

        url = f"https://huggingface.co/{source.repo}/resolve/main/{source.file}"
        headers: dict[str, str] = {}
        if token:
            headers["Authorization"] = f"Bearer {token}"

        dest_file = task.destination / source.file

        for attempt in range(self.max_retries):
            try:
                resume_pos = dest_file.stat().st_size if dest_file.exists() else 0
                req_headers = dict(headers)
                if resume_pos > 0:
                    req_headers["Range"] = f"bytes={resume_pos}-"
                    logger.debug(
                        "Resuming %s from byte %d", source.file, resume_pos
                    )

                response = requests.get(
                    url, headers=req_headers, stream=True, timeout=30
                )

                if response.status_code == 401:
                    return DownloadResult(
                        task=task, success=False, skipped=False,
                        error="Unauthorised — check your HuggingFace API token.",
                    )
                if response.status_code == 403:
                    return DownloadResult(
                        task=task, success=False, skipped=False,
                        error=(
                            f"Access denied (gated model). "
                            f"Visit https://huggingface.co/{source.repo} "
                            f"to request access."
                        ),
                    )
                if response.status_code == 429:
                    retry_after = int(response.headers.get("Retry-After", 60))
                    if attempt < self.max_retries - 1:
                        logger.warning(
                            "Rate limited — waiting %ds before retry.", retry_after
                        )
                        time.sleep(retry_after)
                        continue
                    return DownloadResult(
                        task=task, success=False, skipped=False,
                        error=f"Rate limited. Retry after {retry_after}s.",
                    )
                if response.status_code not in (200, 206):
                    if attempt < self.max_retries - 1:
                        time.sleep(2 ** attempt)
                        continue
                    return DownloadResult(
                        task=task, success=False, skipped=False,
                        error=f"HTTP {response.status_code}",
                    )

                # Set progress total now that we know Content-Length
                total = int(response.headers.get("Content-Length", 0))
                if resume_pos > 0 and response.status_code == 206:
                    total += resume_pos
                if progress is not None and prog_id is not None:
                    progress.update(prog_id, total=total or None, completed=resume_pos)

                mode = (
                    "ab"
                    if resume_pos > 0 and response.status_code == 206
                    else "wb"
                )
                try:
                    with open(dest_file, mode) as fh:
                        for chunk in response.iter_content(chunk_size=65536):
                            if chunk:
                                fh.write(chunk)
                                if progress is not None and prog_id is not None:
                                    progress.advance(prog_id, len(chunk))
                except KeyboardInterrupt:
                    if dest_file.exists():
                        dest_file.unlink()
                        logger.debug("Removed partial file: %s", dest_file)
                    raise

                if total and progress is not None and prog_id is not None:
                    progress.update(prog_id, completed=total)

                return self._integrity_check(task, dest_file, progress, prog_id)

            except requests.exceptions.RequestException as exc:
                if attempt < self.max_retries - 1:
                    wait = 2 ** attempt
                    logger.warning(
                        "Network error on attempt %d/%d for %s: %s — retrying in %ds",
                        attempt + 1, self.max_retries, source.file, exc, wait,
                    )
                    time.sleep(wait)
                    continue
                return DownloadResult(
                    task=task, success=False, skipped=False,
                    error=f"Network error: {exc}",
                )

        return DownloadResult(
            task=task, success=False, skipped=False, error="Max retries exceeded"
        )

    def _download_direct(
        self,
        task: DownloadTask,
        progress: Progress | None = None,
        prog_id: TaskID | None = None,
    ) -> DownloadResult:
        """
        Download from a direct URL with resume support.

        Args:
            task:     Download task with a ``direct_url`` source.
            progress: Rich Progress instance (optional).
            prog_id:  Progress task ID (optional).

        Returns:
            DownloadResult describing the outcome.
        """
        source = task.result.source
        if not source or not source.download_url:
            return DownloadResult(
                task=task, success=False, skipped=False, error="No download URL"
            )

        url       = source.download_url
        dest_file = task.destination / source.file
        headers: dict[str, str] = {}

        for attempt in range(self.max_retries):
            try:
                resume_pos = dest_file.stat().st_size if dest_file.exists() else 0
                req_headers = dict(headers)
                if resume_pos > 0:
                    req_headers["Range"] = f"bytes={resume_pos}-"

                response = requests.get(
                    url, headers=req_headers, stream=True, timeout=30
                )

                if response.status_code == 429:
                    retry_after = int(response.headers.get("Retry-After", 60))
                    if attempt < self.max_retries - 1:
                        time.sleep(retry_after)
                        continue
                    return DownloadResult(
                        task=task, success=False, skipped=False,
                        error=f"Rate limited. Retry after {retry_after}s.",
                    )
                if response.status_code not in (200, 206):
                    if attempt < self.max_retries - 1:
                        time.sleep(2 ** attempt)
                        continue
                    return DownloadResult(
                        task=task, success=False, skipped=False,
                        error=f"HTTP {response.status_code}",
                    )

                total = int(response.headers.get("Content-Length", 0))
                if resume_pos > 0 and response.status_code == 206:
                    total += resume_pos
                if progress is not None and prog_id is not None:
                    progress.update(prog_id, total=total or None, completed=resume_pos)

                mode = (
                    "ab"
                    if resume_pos > 0 and response.status_code == 206
                    else "wb"
                )
                try:
                    with open(dest_file, mode) as fh:
                        for chunk in response.iter_content(chunk_size=65536):
                            if chunk:
                                fh.write(chunk)
                                if progress is not None and prog_id is not None:
                                    progress.advance(prog_id, len(chunk))
                except KeyboardInterrupt:
                    if dest_file.exists():
                        dest_file.unlink()
                        logger.debug("Removed partial file: %s", dest_file)
                    raise

                if total and progress is not None and prog_id is not None:
                    progress.update(prog_id, completed=total)

                return self._integrity_check(task, dest_file, progress, prog_id)

            except requests.exceptions.RequestException as exc:
                if attempt < self.max_retries - 1:
                    time.sleep(2 ** attempt)
                    continue
                return DownloadResult(
                    task=task, success=False, skipped=False,
                    error=f"Network error: {exc}",
                )

        return DownloadResult(
            task=task, success=False, skipped=False, error="Max retries exceeded"
        )

    def _download_civitai(
        self,
        task: DownloadTask,
        progress: Progress | None = None,
        prog_id: TaskID | None = None,
    ) -> DownloadResult:
        """
        Download from Civitai, injecting the API token as a query parameter.

        Args:
            task:     Download task with a ``civitai`` source.
            progress: Rich Progress instance (optional).
            prog_id:  Progress task ID (optional).

        Returns:
            DownloadResult describing the outcome.
        """
        source = task.result.source
        if not source or not source.download_url:
            return DownloadResult(
                task=task, success=False, skipped=False, error="No download URL"
            )

        token = (
            os.getenv("COMFY_CIVITAI_TOKEN")
            or self.config.get("api_keys", {}).get("civitai")
            or None
        )
        if not token:
            return DownloadResult(
                task=task, success=False, skipped=False,
                error="Civitai API token required. Set api_keys.civitai in config.",
            )

        # Append token as query parameter (Civitai's auth mechanism)
        sep = "&" if "?" in source.download_url else "?"
        authed_url = f"{source.download_url}{sep}token={token}"

        # Re-use direct download logic with the token-injected URL
        authed_source = ResolvedSource(
            source_type=source.source_type,
            repo=source.repo,
            file=source.file,
            download_url=authed_url,
            sha256=source.sha256,
            file_size_bytes=source.file_size_bytes,
        )
        authed_result = ResolutionResult(
            input_name=task.result.input_name,
            resolved_name=task.result.resolved_name,
            source=authed_source,
            confidence=task.result.confidence,
            error=task.result.error,
            asset_type=task.result.asset_type,
        )
        authed_task = DownloadTask(
            result=authed_result,
            destination=task.destination,
            asset_type=task.asset_type,
        )
        return self._download_direct(authed_task, progress, prog_id)

    # ── Integrity ───────────────────────────────────────────────────────────

    def _integrity_check(
        self,
        task: DownloadTask,
        dest_file: Path,
        progress: Progress | None,
        prog_id: TaskID | None,
    ) -> DownloadResult:
        """
        Verify SHA-256 integrity if a hash is available.

        On mismatch the file is **kept** and the result carries a
        ``sha256_warning`` so the caller can inform the user.  The file is
        never silently deleted — the user decides what to do.

        On match (or when no expected hash is available) a clean success
        result is returned.

        Args:
            task:      The original download task.
            dest_file: Path to the completed download.
            progress:  Rich Progress instance (for status updates).
            prog_id:   Progress task ID.

        Returns:
            DownloadResult — success=True in all cases; sha256_warning set on mismatch.
        """
        source = task.result.source
        if not source or not source.sha256:
            # No hash to check — treat as clean success
            return DownloadResult(task=task, success=True, skipped=False)

        if progress is not None and prog_id is not None:
            progress.update(prog_id, description="[cyan]verifying…[/cyan]")

        actual = self._compute_sha256(dest_file)
        if actual.lower() == source.sha256.lower():
            logger.debug("SHA-256 verified for %s", dest_file.name)
            return DownloadResult(task=task, success=True, skipped=False)

        warning = (
            f"SHA-256 mismatch for {dest_file.name}\n"
            f"  Expected : {source.sha256.lower()}\n"
            f"  Actual   : {actual.lower()}\n"
            f"  The file has been kept. Verify manually or re-run with --no-cache."
        )
        logger.warning(warning)
        return DownloadResult(
            task=task,
            success=True,          # file exists and may still be usable
            skipped=False,
            sha256_warning=warning,
        )

    def _compute_sha256(self, path: Path) -> str:
        """
        Compute the SHA-256 hex digest of a file.

        Args:
            path: Path to the file.

        Returns:
            Lowercase hex string.
        """
        h = hashlib.sha256()
        with open(path, "rb") as fh:
            for chunk in iter(lambda: fh.read(65536), b""):
                h.update(chunk)
        return h.hexdigest()

    # ── Skip logic ──────────────────────────────────────────────────────────

    def _should_skip(self, task: DownloadTask) -> bool:
        """
        Return True if the destination file already exists and either has no
        expected hash or its hash matches — i.e. downloading is unnecessary.

        Args:
            task: Download task to evaluate.

        Returns:
            True if the download should be skipped.
        """
        source = task.result.source
        if not source:
            return False

        dest_file = task.destination / source.file
        if not dest_file.exists():
            return False

        # File exists but no hash — assume it's correct and skip
        if not source.sha256:
            return True

        # File exists and hash available — only skip if hash matches
        return self._compute_sha256(dest_file).lower() == source.sha256.lower()

    # ── Disk space ──────────────────────────────────────────────────────────

    def check_disk_space(
        self, tasks: list[DownloadTask]
    ) -> tuple[bool, int, int]:
        """
        Pre-flight disk-space check for a list of download tasks.

        Only counts files that will actually be downloaded (not skipped).
        Adds a 10 % safety buffer to the required figure.

        Args:
            tasks: Download tasks to evaluate.

        Returns:
            ``(enough_space, required_bytes, available_bytes)``
        """
        if not tasks:
            return True, 0, 0

        required = sum(
            t.result.source.file_size_bytes
            for t in tasks
            if t.result.source
            and t.result.source.file_size_bytes
            and not self._should_skip(t)
        )

        available = self._available_space(tasks[0].destination)
        return available >= int(required * 1.1), required, available

    def _available_space(self, path: Path) -> int:
        """
        Return free bytes on the volume containing *path*.

        Walks up the directory tree until an existing ancestor is found so
        that this never crashes when the destination doesn't exist yet.

        Args:
            path: Target path (need not exist).

        Returns:
            Free bytes available.
        """
        check = path
        while not check.exists():
            check = check.parent
        return shutil.disk_usage(check).free