"""
sources/civitai.py — Civitai asset resolution source

The Civitai **search** API is fully public — no token required.
A token is only needed at *download* time.  This module reflects that:
- ``search()`` always runs regardless of token presence.
- ``has_download_token()`` lets callers check before attempting a download.

Query strategy
--------------
Civitai indexes models by *display name*, not by filename.  Raw filenames
like ``SexGod_Nudity_LTX2_v1_5.safetensors`` often return no hits because
they contain version suffixes, underscores, and quantisation tags that the
index doesn't see.

We therefore build a *search term cascade*:

  1. Full cleaned stem  (underscores→spaces, version suffixes stripped)
  2. First N words of the cleaned stem  (handles very long names)
  3. Raw stem as fallback

Each strategy is tried in order and we stop as soon as we get ≥1 hit.
"""

from __future__ import annotations

import logging
import os
import re
from typing import Any

import requests  # type: ignore
from rapidfuzz import fuzz  # type: ignore

from .base import BaseSource, ResolutionResult, ResolvedSource

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Patterns stripped from filenames before searching
# ---------------------------------------------------------------------------

#: Version-like suffixes to remove, e.g. ``_v1_5``, ``_V2_step_06000``,
#: ``_bf16``, ``_Q4_0``, ``_Q8_0``, ``_fp16``, ``_fpmixed``.
_VERSION_RE = re.compile(
    r"""
    (?:
        [_\-]v\d+(?:[_.\-]\d+)*       # _v1, _v1_5, -v2.3
      | [_\-]V\d+(?:[_.\-]\d+)*       # _V1, _V2_step_06000
      | [_\-]step[_\-]\d+             # _step_06000
      | [_\-](?:bf|fp)\d+(?:mixed)?   # _bf16, _fp16, _fpmixed
      | [_\-][Qq]\d+[_\-][KkMm0]      # _Q4_0, _Q4_K_M, _q8_0
      | [_\-]rank[_\-]\d+             # _rank_242
      | [_\-]avg                      # _avg
      | [_\-]dynamic                  # _dynamic
      | [_\-]fro\d+                   # _fro095
      | [_\-]resized                  # _resized
      | [_\-]distill(?:ed)?           # _distill, _distilled
    )+
    $                                  # only strip from the end
    """,
    re.VERBOSE | re.IGNORECASE,
)


def _clean_query(filename: str) -> str:
    """
    Derive a human-friendly search term from a model filename.

    Steps:
    1. Strip file extension.
    2. Remove trailing version / quantisation suffixes.
    3. Replace underscores and hyphens with spaces.
    4. Collapse multiple spaces.

    Examples::

        "SexGod_Nudity_LTX2_v1_5.safetensors"
            → "SexGod Nudity LTX2"

        "ltx-2-19b-distilled-lora-resized_dynamic_fro095_avg_rank_242_bf16.safetensors"
            → "ltx 2 19b"   (after aggressive suffix stripping)

        "Qwen-Rapid-NSFW-v20_Q4_0.gguf"
            → "Qwen Rapid NSFW"

    Args:
        filename: Raw model filename (with or without extension).

    Returns:
        Cleaned search string.
    """
    # 1. Strip extension
    stem = filename
    for ext in (".safetensors", ".ckpt", ".pt", ".bin", ".pth", ".gguf", ".sft"):
        if stem.lower().endswith(ext):
            stem = stem[: -len(ext)]
            break

    # 2. Strip version / quantisation suffixes from the end
    stem = _VERSION_RE.sub("", stem)

    # 3. Underscores / hyphens → spaces
    stem = re.sub(r"[_\-]+", " ", stem)

    # 4. Collapse whitespace
    return " ".join(stem.split())


def _query_cascade(filename: str) -> list[str]:
    """
    Return an ordered list of search queries to try, from most specific
    to most general.

    Args:
        filename: Raw model filename.

    Returns:
        List of query strings (deduplicated, non-empty).
    """
    cleaned = _clean_query(filename)
    words   = cleaned.split()

    candidates: list[str] = []

    # Full cleaned name
    if cleaned:
        candidates.append(cleaned)

    # First 4 words (handles very long descriptive names)
    if len(words) > 4:
        candidates.append(" ".join(words[:4]))

    # First 3 words
    if len(words) > 3:
        candidates.append(" ".join(words[:3]))

    # First 2 words (broadest useful query)
    if len(words) >= 2:
        candidates.append(" ".join(words[:2]))

    # Raw stem as absolute fallback (no cleaning)
    raw_stem = filename.rsplit(".", 1)[0] if "." in filename else filename
    candidates.append(raw_stem)

    # Deduplicate while preserving order
    seen: set[str] = set()
    result: list[str] = []
    for c in candidates:
        if c and c not in seen:
            seen.add(c)
            result.append(c)

    return result


# ---------------------------------------------------------------------------
# Source class
# ---------------------------------------------------------------------------

class CivitaiSource(BaseSource):
    """
    Civitai asset resolution source.

    Search is unconditionally available (public API).
    A token is only required for downloading files.
    """

    def __init__(self, token: str | None = None) -> None:
        """
        Initialise the Civitai source.

        Args:
            token: Optional Civitai API token.  Not needed for searching;
                   required when the downloader fetches a file.
        """
        self._token   = token or None
        self._api_base = "https://civitai.com/api/v1"

    # ── BaseSource interface ────────────────────────────────────────────────

    @property
    def name(self) -> str:
        return "civitai"

    def is_available(self, config: dict[str, Any]) -> bool:
        """Always ``True`` — Civitai search is public, no token required."""
        return True

    def has_download_token(self, config: dict[str, Any]) -> bool:
        """
        Return ``True`` if a Civitai API token is configured.

        A token is required to download files but not to search.

        Args:
            config: comfy-resolve configuration dictionary.

        Returns:
            True if a token is available via config or environment variable.
        """
        return bool(
            os.getenv("COMFY_CIVITAI_TOKEN")
            or config.get("api_keys", {}).get("civitai")
        )

    def search(self, filename: str, asset_type: str) -> list[ResolutionResult]:
        """
        Search Civitai for a model file.

        Uses a cascading query strategy: progressively shorter / cleaner
        search terms are tried in order until results are found.

        Token is **not** required for searching.  If no token is set, results
        will include a note that one is needed to download.

        Args:
            filename:   The model filename to search for.
            asset_type: Asset category (e.g. ``"loras"``, ``"checkpoints"``).

        Returns:
            Up to 5 :class:`ResolutionResult` objects sorted by confidence,
            or an empty list if nothing was found.
        """
        token = os.getenv("COMFY_CIVITAI_TOKEN") or self._token

        for query in _query_cascade(filename):
            logger.debug("Civitai search: query=%r  (for %s)", query, filename)
            results = self._search_with_query(filename, query, token)
            if results:
                logger.debug(
                    "Civitai: %d result(s) for %r via query %r",
                    len(results), filename, query,
                )
                return results

        logger.debug("Civitai: no results found for %s", filename)
        return []

    # ── Internal helpers ────────────────────────────────────────────────────

    def _search_with_query(
        self,
        original_filename: str,
        query: str,
        token: str | None,
    ) -> list[ResolutionResult]:
        """
        Execute a single Civitai API search and score the results.

        Args:
            original_filename: The original filename (used for scoring).
            query:             The cleaned search term to send to the API.
            token:             Optional API token (for auth header only).

        Returns:
            Sorted, deduplicated list of up to 5 ResolutionResult objects.
            Returns an empty list on API errors or no hits.
        """
        headers: dict[str, str] = {}
        if token:
            headers["Authorization"] = f"Bearer {token}"

        try:
            response = requests.get(
                f"{self._api_base}/models",
                params={"query": query, "limit": 20},
                headers=headers,
                timeout=10,
            )
        except requests.exceptions.RequestException as exc:
            logger.warning("Civitai network error: %s", exc)
            return []

        if response.status_code == 429:
            retry = response.headers.get("Retry-After", "?")
            logger.warning("Civitai rate-limited. Retry-After: %s", retry)
            return []

        if response.status_code == 401:
            logger.warning("Civitai: invalid token (401)")
            return []

        if response.status_code != 200:
            logger.warning("Civitai API error: HTTP %d", response.status_code)
            return []

        try:
            models = response.json().get("items", [])
        except Exception as exc:
            logger.warning("Civitai: failed to parse response: %s", exc)
            return []

        results: list[ResolutionResult] = []
        fn_lower = original_filename.lower()

        for model in models[:10]:
            model_id = model.get("id")
            if not model_id:
                continue

            for version in model.get("modelVersions", [])[:3]:
                for file_info in version.get("files", []):
                    candidate_name: str = file_info.get("name", "")
                    if not candidate_name:
                        continue

                    # token_set_ratio handles word reordering and partial
                    # overlap much better than plain ratio for model names
                    score = fuzz.token_set_ratio(
                        fn_lower, candidate_name.lower()
                    )

                    if score < 60:
                        continue

                    # Also try matching against the cleaned query so that
                    # "SexGod Nudity LTX2" vs "SexGod LTX 2.3 Female Nudity"
                    # gets a fair shake
                    query_score = fuzz.token_set_ratio(
                        query.lower(), candidate_name.lower()
                    )
                    best_score = max(score, query_score)

                    if best_score < 60:
                        continue

                    # Exact filename match → high; ≥80 → medium; else low
                    if fn_lower == candidate_name.lower():
                        confidence = "high"
                    elif best_score >= 75:
                        confidence = "medium"
                    else:
                        confidence = "low"

                    hashes      = file_info.get("hashes", {})
                    sha256      = hashes.get("SHA256") if isinstance(hashes, dict) else None
                    size_kb     = file_info.get("sizeKB")
                    size_bytes  = int(size_kb * 1024) if size_kb else None
                    download_url = file_info.get("downloadUrl")

                    # Annotate if no token — download will fail without one
                    note = None
                    if not token:
                        note = "Civitai token required to download. Set api_keys.civitai in config."

                    results.append(ResolutionResult(
                        input_name=original_filename,
                        resolved_name=candidate_name,
                        source=ResolvedSource(
                            source_type="civitai",
                            repo=f"civitai:{model_id}",
                            file=candidate_name,
                            download_url=download_url,
                            sha256=sha256,
                            file_size_bytes=size_bytes,
                        ),
                        confidence=confidence,
                        error=note,
                    ))

        if not results:
            return []

        # Sort: confidence first, then score descending
        confidence_rank = {"high": 3, "medium": 2, "low": 1, "unknown": 0}
        results.sort(
            key=lambda r: (
                confidence_rank.get(r.confidence, 0),
                fuzz.token_set_ratio(fn_lower, r.resolved_name.lower()),
            ),
            reverse=True,
        )

        return results[:5]