"""Model search and download tools backed by ComfyUI-Model-Manager."""

from __future__ import annotations

import asyncio
import re
from typing import Annotated, Any, Literal
from urllib.parse import urlparse

import httpx
from mcp.server.fastmcp import FastMCP
from mcp.types import ToolAnnotations
from pydantic import Field

from comfyui_mcp.audit import AuditLogger
from comfyui_mcp.client import ComfyUIClient
from comfyui_mcp.config import ModelSearchSettings
from comfyui_mcp.model_manager import ModelManagerDetector
from comfyui_mcp.pagination import paginate
from comfyui_mcp.security.download_validator import DownloadValidator
from comfyui_mcp.security.rate_limit import RateLimiter
from comfyui_mcp.security.sanitizer import PathSanitizer

_HF_API = "https://huggingface.co/api/models"
_CIVITAI_API = "https://civitai.com/api/v1/models"

# Extensions we look for when finding the primary model file in HuggingFace repos
_MODEL_EXTENSIONS = {".safetensors", ".ckpt", ".pt", ".pth", ".bin"}

_HF_REPO_RE = re.compile(r"^[A-Za-z0-9_.\-]+/[A-Za-z0-9_.\-]+$")


async def _search_civitai(
    query: str,
    model_type: str,
    limit: int,
    api_key: str,
    http: httpx.AsyncClient,
) -> list[dict[str, Any]]:
    """Search CivitAI for models."""
    params: dict[str, str | int] = {"query": query, "limit": limit}
    if model_type:
        params["types"] = model_type
    params["sort"] = "Most Downloaded"

    headers: dict[str, str] = {}
    if api_key:
        headers["Authorization"] = f"Bearer {api_key}"

    r = await http.get(_CIVITAI_API, params=params, headers=headers, timeout=30)
    r.raise_for_status()
    data = r.json()

    results: list[dict[str, Any]] = []
    for item in data.get("items", []):
        versions = item.get("modelVersions", [])
        if not versions:
            continue
        latest = versions[0]
        files = latest.get("files", [])
        size_kb = files[0].get("sizeKB", 0) if files else 0
        filename = files[0].get("name", "") if files else ""
        results.append(
            {
                "name": item.get("name", ""),
                "type": item.get("type", ""),
                "url": latest.get("downloadUrl", ""),
                "filename": filename,
                "size_mb": round(size_kb / 1024, 1),
                "downloads": item.get("stats", {}).get("downloadCount", 0),
                "rating": item.get("stats", {}).get("rating", 0),
                "source": "civitai",
            }
        )
    return results


async def _fetch_hf_model_detail(
    http: httpx.AsyncClient,
    model: dict[str, Any],
    headers: dict[str, str],
) -> dict[str, Any] | None:
    """Fetch detail for a single HuggingFace model. Returns None for invalid IDs."""
    model_id = model.get("id", "")
    if not isinstance(model_id, str) or not _HF_REPO_RE.match(model_id):
        return None
    detail_url = f"{_HF_API}/{model_id}"
    try:
        dr = await http.get(detail_url, headers=headers, timeout=30)
        dr.raise_for_status()
        detail = dr.json()
    except httpx.HTTPError:
        detail = {}

    siblings = detail.get("siblings", [])
    model_file = None
    model_size = 0
    for sib in siblings:
        fname = sib.get("rfilename", "")
        ext = "." + fname.rsplit(".", 1)[-1] if "." in fname else ""
        if ext.lower() in _MODEL_EXTENSIONS:
            size = sib.get("size", 0)
            if size > model_size:
                model_size = size
                model_file = fname

    download_url = ""
    if model_file:
        download_url = f"https://huggingface.co/{model_id}/resolve/main/{model_file}"

    return {
        "name": model_id,
        "type": model.get("pipeline_tag", ""),
        "url": download_url,
        "filename": model_file or "",
        "size_mb": round(model_size / (1024 * 1024), 1) if model_size else 0,
        "downloads": model.get("downloads", 0),
        "likes": model.get("likes", 0),
        "source": "huggingface",
    }


async def _search_huggingface(
    query: str,
    model_type: str,
    limit: int,
    token: str,
    http: httpx.AsyncClient,
) -> list[dict[str, Any]]:
    """Search HuggingFace for models with concurrent detail fetches."""
    params: dict[str, str | int] = {"search": query, "limit": limit}
    if model_type:
        params["pipeline_tag"] = model_type

    headers: dict[str, str] = {}
    if token:
        headers["Authorization"] = f"Bearer {token}"

    r = await http.get(_HF_API, params=params, headers=headers, timeout=30)
    r.raise_for_status()
    models = r.json()

    results = await asyncio.gather(*[_fetch_hf_model_detail(http, m, headers) for m in models])
    return [r for r in results if r is not None]


def register_model_tools(
    mcp: FastMCP,
    client: ComfyUIClient,
    audit: AuditLogger,
    read_limiter: RateLimiter,
    file_limiter: RateLimiter,
    sanitizer: PathSanitizer,
    detector: ModelManagerDetector,
    validator: DownloadValidator,
    search_settings: ModelSearchSettings,
    search_http: httpx.AsyncClient,
) -> dict[str, Any]:
    """Register model search and download tools."""
    tool_fns: dict[str, Any] = {}

    @mcp.tool(
        annotations=ToolAnnotations(
            readOnlyHint=True,
            destructiveHint=False,
            idempotentHint=True,
            openWorldHint=True,
        )
    )
    async def comfyui_search_models(
        query: Annotated[
            str,
            Field(
                description="Search query (model name, style, etc.)",
                min_length=1,
                max_length=200,
            ),
        ],
        source: Annotated[
            Literal["civitai", "huggingface"],
            Field(description='Where to search — "civitai" or "huggingface"'),
        ] = "civitai",
        model_type: Annotated[
            str,
            Field(
                description="Filter by type. CivitAI: Checkpoint, LORA, TextualInversion, etc. "
                "HuggingFace: text-to-image, etc.",
                max_length=100,
            ),
        ] = "",
        limit: Annotated[
            int,
            Field(
                description="Maximum results to return (0 for default page size; "
                "values above server max are capped)",
                ge=0,
            ),
        ] = 5,
        offset: Annotated[int, Field(description="Starting index for pagination", ge=0)] = 0,
    ) -> dict[str, Any]:
        """Search for models on HuggingFace or CivitAI.

        Returns:
            JSON with search results including name, download URL, size, and stats.
            Use comfyui_download_model with the URL to install a model.
        """
        read_limiter.check("search_models")

        # Input validation
        stripped_query = query.strip()
        if not stripped_query:
            raise ValueError("query must not be empty")
        if len(stripped_query) > 200:
            raise ValueError("query must not exceed 200 characters")
        if model_type and len(model_type) > 100:
            raise ValueError("model_type must not exceed 100 characters")

        if source not in ("civitai", "huggingface"):
            raise ValueError("source must be 'civitai' or 'huggingface'")

        cap = search_settings.max_search_results

        await audit.async_log(
            tool="search_models",
            action="searching",
            extra={"query": stripped_query, "source": source, "model_type": model_type},
        )

        if source == "civitai":
            results = await _search_civitai(
                stripped_query, model_type, cap, search_settings.civitai_api_key, search_http
            )
        else:
            results = await _search_huggingface(
                stripped_query, model_type, cap, search_settings.huggingface_token, search_http
            )

        await audit.async_log(
            tool="search_models",
            action="searched",
            extra={"source": source, "result_count": len(results)},
        )

        result = paginate(
            results, offset, limit, default_limit=5, max_limit=search_settings.max_search_results
        )
        result["query"] = stripped_query
        result["source"] = source
        return result

    tool_fns["comfyui_search_models"] = comfyui_search_models

    @mcp.tool(
        annotations=ToolAnnotations(
            readOnlyHint=False,
            destructiveHint=False,
            idempotentHint=False,
            openWorldHint=True,
        )
    )
    async def comfyui_download_model(
        url: Annotated[
            str,
            Field(description="Direct download URL (must be from an allowed domain)"),
        ],
        folder: Annotated[
            str,
            Field(description='Target model folder (e.g. "checkpoints", "loras", "vae")'),
        ],
        filename: Annotated[
            str,
            Field(description="Filename to save as (optional — inferred from URL if empty)"),
        ] = "",
    ) -> dict[str, Any]:
        """Download a model from HuggingFace or CivitAI via ComfyUI-Model-Manager.

        Returns:
            JSON with download task status. Use comfyui_get_download_tasks to check progress.
        """
        file_limiter.check("download_model")

        # Validate URL domain and path pattern
        validator.validate_url(url)

        # Validate folder
        sanitizer.validate_path_segment(folder, label="folder")
        await detector.validate_folder(folder)

        # Infer filename from URL if not provided
        if not filename:
            path = urlparse(url).path
            filename = path.rsplit("/", 1)[-1] if "/" in path else ""
            if not filename:
                raise ValueError("Could not infer filename from URL. Please provide a filename.")

        # Validate extension and filename
        validator.validate_extension(filename)
        sanitizer.validate_filename(filename)

        # Infer download platform from URL
        hostname = (urlparse(url).hostname or "").lower()
        if "civitai" in hostname:
            platform = "civitai"
        elif "huggingface" in hostname:
            platform = "huggingface"
        else:
            platform = "other"

        await audit.async_log(
            tool="download_model",
            action="downloading",
            extra={"url": url, "folder": folder, "filename": filename, "platform": platform},
        )

        result = await client.create_download_task(
            model_type=folder,
            path_index=0,
            fullname=filename,
            download_platform=platform,
            download_url=url,
            size_bytes=0,
        )

        await audit.async_log(
            tool="download_model",
            action="download_started",
            extra={"result": result, "folder": folder, "filename": filename},
        )

        return result

    tool_fns["comfyui_download_model"] = comfyui_download_model

    @mcp.tool(
        annotations=ToolAnnotations(
            readOnlyHint=True,
            destructiveHint=False,
            idempotentHint=True,
            openWorldHint=True,
        )
    )
    async def comfyui_get_download_tasks() -> dict[str, Any]:
        """Check the status of active model downloads.

        Returns:
            JSON with list of download tasks including progress, speed, and status.
        """
        read_limiter.check("get_download_tasks")
        await detector.get_folders()  # Ensure Model Manager is available

        await audit.async_log(tool="get_download_tasks", action="checking")

        tasks = await client.get_download_tasks()

        await audit.async_log(
            tool="get_download_tasks",
            action="checked",
            extra={"task_count": len(tasks)},
        )

        return {"tasks": tasks}

    tool_fns["comfyui_get_download_tasks"] = comfyui_get_download_tasks

    @mcp.tool(
        annotations=ToolAnnotations(
            readOnlyHint=False,
            destructiveHint=True,
            idempotentHint=False,
            openWorldHint=True,
        )
    )
    async def comfyui_cancel_download(task_id: str) -> dict[str, Any]:
        """Cancel and remove a model download task.

        Args:
            task_id: ID of the download task to cancel
        """
        file_limiter.check("cancel_download")
        await detector.get_folders()  # Ensure Model Manager is available

        await audit.async_log(
            tool="cancel_download",
            action="cancelling",
            extra={"task_id": task_id},
        )

        result = await client.delete_download_task(task_id)

        success = bool(result.get("success", True)) if isinstance(result, dict) else True

        await audit.async_log(
            tool="cancel_download",
            action="cancelled",
            extra={"task_id": task_id, "result": result, "success": success},
        )

        return {"success": success, "task_id": task_id, "result": result}

    tool_fns["comfyui_cancel_download"] = comfyui_cancel_download

    return tool_fns
