"""
HuggingFace Source Module

Search and download models from HuggingFace Hub.
"""

import json
import os
import re
import threading
import time
import requests
from typing import Dict, Any, Optional, List, Callable
from urllib.parse import urlparse, quote

from ..progress import report_progress
from ..log_system.log_funcs import create_module_logger
log = create_module_logger(__name__)

from ..path_utils import write_json_atomic, METADATA_DIR
from ..type_utils import check_credential_http

HF_API_URL = "https://huggingface.co/api"
HF_AUTHOR_FALLBACKS = ["Comfy-Org"]
BRAVE_SEARCH_API_URL = "https://api.search.brave.com/res/v1/web/search"
HF_AUTHOR_INDEX_CACHE_TTL_SECONDS = 24 * 60 * 60
HF_AUTHOR_INDEX_CACHE_VERSION = 1

HF_AUTHOR_INDEX_CACHE_PATH = os.path.join(
    METADATA_DIR, "huggingface-author-index.json"
)

# Cache for search results
_search_cache: Dict[str, Any] = {}
_author_index_cache: Dict[str, Dict[str, Any]] = {}
_download_size_cache: Dict[str, Optional[int]] = {}
_author_index_lock = threading.RLock()


def clear_search_cache():
    """Clear cached HuggingFace search results and in-memory indexes."""
    global _search_cache, _author_index_cache
    _search_cache.clear()
    _download_size_cache.clear()
    with _author_index_lock:
        _author_index_cache.clear()


def _report_progress(
    progress_callback: Optional[Callable[[Dict[str, Any]], None]],
    stage: str,
    message: str,
    percent: Optional[float] = None,
    **extra: Any,
) -> None:
    report_progress(
        progress_callback,
        stage,
        message,
        percent,
        error_context="HuggingFace progress callback",
        **extra,
    )


def check_huggingface_token(token: Optional[str]) -> Dict[str, Any]:
    """Check whether a HuggingFace access token is accepted."""
    value = (token or "").strip()
    if not value:
        return {
            "success": False,
            "valid": False,
            "status": "missing",
            "message": "Paste a HuggingFace token first.",
        }

    headers = {"Authorization": f"Bearer {value}"}

    def get_user(data):
        return data.get("name") or data.get("user", {}).get("name") or ""

    result = check_credential_http(
        f"{HF_API_URL}/whoami-v2",
        headers=headers,
        success_message="HuggingFace token is valid.",
        get_username=get_user,
        error_msg_401_403="HuggingFace token is not accepted.",
    )
    if result.get("status_code") == 404:
        result = check_credential_http(
            f"{HF_API_URL}/whoami",
            headers=headers,
            success_message="HuggingFace token is valid.",
            get_username=get_user,
            error_msg_401_403="HuggingFace token is not accepted.",
        )
    return result


def check_brave_search_api_key(api_key: Optional[str]) -> Dict[str, Any]:
    """Check whether a Brave Search API key is accepted."""
    key = (api_key or "").strip()
    if not key:
        return {
            "success": False,
            "valid": False,
            "status": "missing",
            "message": "Paste a Brave Search API key first.",
        }

    headers = {"X-Subscription-Token": key}
    params = {"q": "test", "count": 1}

    def custom_429(response):
        return {
            "success": False,
            "valid": False,
            "status": "limited",
            "message": "Brave Search rate limit was reached.",
            "status_code": 429,
        }

    return check_credential_http(
        BRAVE_SEARCH_API_URL,
        headers=headers,
        params=params,
        success_message="Brave Search API key is valid.",
        error_msg_401_403="Brave Search API key is not accepted.",
        custom_429_handler=custom_429,
    )


def _author_index_cache_key(author: str, headers: Dict[str, str]) -> str:
    auth_key = "token" if headers.get("Authorization") else "public"
    return f"{author}::{auth_key}"


def _is_author_index_fresh(index: Optional[Dict[str, Any]]) -> bool:
    if not index:
        return False

    updated_at = index.get("updated_at")
    if not isinstance(updated_at, (int, float)):
        return False

    return (time.time() - updated_at) < HF_AUTHOR_INDEX_CACHE_TTL_SECONDS


def _read_persistent_author_indexes() -> Dict[str, Any]:
    if not os.path.exists(HF_AUTHOR_INDEX_CACHE_PATH):
        return {"version": HF_AUTHOR_INDEX_CACHE_VERSION, "authors": {}}

    try:
        with open(HF_AUTHOR_INDEX_CACHE_PATH, "r", encoding="utf-8") as f:
            data = json.load(f)
        if data.get("version") != HF_AUTHOR_INDEX_CACHE_VERSION:
            return {"version": HF_AUTHOR_INDEX_CACHE_VERSION, "authors": {}}
        if not isinstance(data.get("authors"), dict):
            data["authors"] = {}
        return data
    except Exception as e:
        log.debug(f"Error reading HuggingFace author index cache: {e}")
        return {"version": HF_AUTHOR_INDEX_CACHE_VERSION, "authors": {}}


def _write_persistent_author_index(author: str, index: Dict[str, Any]) -> None:
    try:
        data = _read_persistent_author_indexes()
        data["authors"][author] = index
        write_json_atomic(
            HF_AUTHOR_INDEX_CACHE_PATH,
            data,
            separators=(",", ":"),
        )
    except Exception as e:
        log.debug(f"Error writing HuggingFace author index cache: {e}")


def _build_author_index_from_models(
    author: str, repos: List[Dict[str, Any]]
) -> Dict[str, Any]:
    files = []
    repo_ids = []

    for repo in repos:
        repo_id = repo.get("id", "")
        if not repo_id:
            continue

        repo_ids.append(repo_id)
        for sibling in repo.get("siblings") or []:
            file_path = sibling.get("rfilename") or sibling.get("path") or ""
            if not file_path:
                continue

            files.append(
                {
                    "repo_id": repo_id,
                    "path": file_path,
                    "filename": os.path.basename(file_path),
                    "size": sibling.get("size"),
                }
            )

    return {
        "author": author,
        "updated_at": time.time(),
        "repo_count": len(repo_ids),
        "file_count": len(files),
        "repos": repo_ids,
        "files": files,
    }


def _fetch_author_index(author: str, headers: Dict[str, str], limit: int = 200):
    try:
        response = requests.get(
            f"{HF_API_URL}/models",
            params={"author": author, "limit": limit, "full": "true"},
            headers=headers,
            timeout=15,
        )
        if response.status_code != 200:
            log.debug(
                f"HuggingFace author index returned {response.status_code} for author={author}"
            )
            return None

        repos = response.json()
        if not isinstance(repos, list):
            return None

        index = _build_author_index_from_models(author, repos)
        log.info(
            f"HuggingFace author index refreshed for author={author}: repos={index['repo_count']}, files={index['file_count']}"
        )
        return index
    except Exception as e:
        log.debug(f"Error refreshing HuggingFace author index for author={author}: {e}")
        return None


def _get_author_index(
    author: str, headers: Dict[str, str], force_refresh: bool = False
) -> Optional[Dict[str, Any]]:
    cache_key = _author_index_cache_key(author, headers)
    can_persist = not headers.get("Authorization")

    with _author_index_lock:
        memory_index = _author_index_cache.get(cache_key)
        if not force_refresh and _is_author_index_fresh(memory_index):
            return memory_index

        persistent_index = None
        if can_persist:
            persistent_index = _read_persistent_author_indexes().get("authors", {}).get(
                author
            )
            if not force_refresh and _is_author_index_fresh(persistent_index):
                _author_index_cache[cache_key] = persistent_index
                log.debug(f"HuggingFace author index cache hit for author={author}")
                return persistent_index

        fresh_index = _fetch_author_index(author, headers=headers)
        if fresh_index:
            _author_index_cache[cache_key] = fresh_index
            if can_persist:
                _write_persistent_author_index(author, fresh_index)
            return fresh_index

        if persistent_index:
            log.warn(
                f"Using stale HuggingFace author index for author={author}; refresh failed"
            )
            _author_index_cache[cache_key] = persistent_index
            return persistent_index

        return None


def get_author_fallback_index_status(author: str = "Comfy-Org") -> Dict[str, Any]:
    """Return persisted public author fallback index status for the options UI."""
    with _author_index_lock:
        index = _read_persistent_author_indexes().get("authors", {}).get(author)

    updated_at = index.get("updated_at") if isinstance(index, dict) else None
    age_seconds = (time.time() - updated_at) if isinstance(updated_at, (int, float)) else None
    return {
        "author": author,
        "exists": bool(index),
        "updated_at": updated_at,
        "age_seconds": age_seconds,
        "ttl_seconds": HF_AUTHOR_INDEX_CACHE_TTL_SECONDS,
        "stale": bool(index) and not _is_author_index_fresh(index),
        "repo_count": index.get("repo_count", 0) if isinstance(index, dict) else 0,
        "file_count": index.get("file_count", 0) if isinstance(index, dict) else 0,
    }


def refresh_author_fallback_index(
    token: Optional[str] = None, author: str = "Comfy-Org"
) -> Dict[str, Any]:
    """Force refresh the public HuggingFace author fallback index."""
    headers = {}

    index = _get_author_index(author, headers=headers, force_refresh=True)
    status = get_author_fallback_index_status(author)
    status["success"] = bool(index)
    status["persisted"] = bool(index)
    return status


def parse_huggingface_url(url: str) -> Optional[Dict[str, str]]:
    """
    Parse a HuggingFace URL to extract repo and filename.

    Handles formats:
    - https://huggingface.co/user/repo/resolve/main/file.safetensors
    - https://huggingface.co/user/repo/blob/main/file.safetensors
    - hf://user/repo/file.safetensors

    Returns:
        Dictionary with 'repo' and 'filename' keys, or None if not HF URL
    """
    if url.startswith("hf://"):
        parts = url[5:].split("/", 2)
        if len(parts) >= 3:
            return {"repo": f"{parts[0]}/{parts[1]}", "filename": parts[2]}
        return None

    parsed = urlparse(url)
    if "huggingface.co" not in parsed.netloc:
        return None

    match = re.match(r"^/([^/]+/[^/]+)/(resolve|blob)/([^/]+)/(.+)$", parsed.path)
    if match:
        return {
            "repo": match.group(1),
            "branch": match.group(3),
            "filename": match.group(4),
        }

    return None


def get_huggingface_download_url(repo: str, filename: str, branch: str = "main") -> str:
    """Generate a direct download URL for a HuggingFace file."""
    return f"https://huggingface.co/{repo}/resolve/{branch}/{quote(filename)}"


def _parse_size_header(value: Any) -> Optional[int]:
    if value is None or value == "":
        return None
    try:
        size = int(str(value).strip())
        return size if size > 0 else None
    except (TypeError, ValueError):
        return None


def _parse_content_range_size(value: Any) -> Optional[int]:
    if not value:
        return None
    match = re.search(r"/(\d+)\s*$", str(value))
    if not match:
        return None
    return _parse_size_header(match.group(1))


def _extract_file_info_size(file_info: Dict[str, Any]) -> Optional[int]:
    if not isinstance(file_info, dict):
        return None

    for key in ("size", "sizeBytes", "size_bytes", "fileSize", "file_size", "bytes"):
        size = _parse_size_header(file_info.get(key))
        if size:
            return size

    lfs_info = file_info.get("lfs") if isinstance(file_info.get("lfs"), dict) else {}
    for key in ("size", "sizeBytes", "size_bytes"):
        size = _parse_size_header(lfs_info.get(key))
        if size:
            return size

    return None


def _normalize_huggingface_size_probe_url(url: str) -> Optional[str]:
    if not isinstance(url, str) or not url.strip():
        return None
    value = url.strip()
    parsed = urlparse(value)
    if "huggingface.co" not in parsed.netloc:
        return None
    if "/blob/" in parsed.path:
        value = value.replace("/blob/", "/resolve/", 1)
    return value


def _response_file_size(response: requests.Response) -> Optional[int]:
    size = _parse_content_range_size(response.headers.get("Content-Range"))
    if size:
        return size

    for key in (
        "x-linked-size",
        "x-file-size",
        "x-goog-stored-content-length",
        "content-length",
    ):
        size = _parse_size_header(response.headers.get(key))
        if size:
            content_type = (response.headers.get("content-type") or "").lower()
            if key == "content-length" and "text/html" in content_type:
                continue
            return size

    return None


def _fetch_remote_file_size_bytes(
    url: str,
    headers: Optional[Dict[str, str]] = None,
    timeout: int = 15,
) -> Optional[int]:
    probe_url = _normalize_huggingface_size_probe_url(url)
    if not probe_url:
        return None

    request_headers = {**(headers or {}), "Accept-Encoding": "identity"}
    cache_scope = "token" if request_headers.get("Authorization") else "public"
    cache_key = f"{cache_scope}::{probe_url}"
    if cache_key in _download_size_cache:
        return _download_size_cache[cache_key]

    size = None
    try:
        response = requests.head(
            probe_url,
            headers=request_headers,
            allow_redirects=True,
            timeout=timeout,
        )
        try:
            if response.status_code < 400:
                size = _response_file_size(response)
        finally:
            response.close()
    except Exception as e:
        log.debug(f"HuggingFace size HEAD failed: url={probe_url}, error={e}")

    if not size:
        try:
            response = requests.get(
                probe_url,
                headers={**request_headers, "Range": "bytes=0-0"},
                allow_redirects=True,
                stream=True,
                timeout=timeout,
            )
            try:
                if response.status_code < 400:
                    size = _response_file_size(response)
            finally:
                response.close()
        except Exception as e:
            log.debug(f"HuggingFace size range request failed: url={probe_url}, error={e}")

    _download_size_cache[cache_key] = size
    return size


def clean_filename_for_search(filename: str) -> str:
    """
    Clean up filename for better search results.
    Remove common suffixes that might prevent matches.
    """
    base = os.path.splitext(filename)[0]
    # Remove common precision/format suffixes
    base = re.sub(
        r"[-_]?(fp16|fp32|fp8|fp4|bf16|e4m3fn|scaled|pruned|emaonly|mixed|q4|q8).*$",
        "",
        base,
        flags=re.IGNORECASE,
    )
    return base


def _build_huggingface_result(
    repo_id: str,
    file_path: str,
    file_info: Dict[str, Any],
    match_type: str,
    headers: Optional[Dict[str, str]] = None,
) -> Dict[str, Any]:
    lfs_info = file_info.get("lfs") if isinstance(file_info.get("lfs"), dict) else {}
    sha256 = (
        file_info.get("sha256")
        or file_info.get("hash")
        or lfs_info.get("sha256")
        or lfs_info.get("oid")
    )
    download_url = get_huggingface_download_url(repo_id, file_path)
    size = _extract_file_info_size(file_info)
    if not size:
        size = _fetch_remote_file_size_bytes(download_url, headers=headers)
    return {
        "source": "huggingface",
        "repo_id": repo_id,
        "filename": os.path.basename(file_path),
        "path": file_path,
        "url": download_url,
        "size": size,
        "match_type": match_type,
        "sha256": sha256,
    }


def _get_repo_tree(
    repo_id: str, headers: Dict[str, str], branch: str = "main"
) -> Optional[List[Dict[str, Any]]]:
    files_url = f"{HF_API_URL}/models/{repo_id}/tree/{quote(branch, safe='')}"

    try:
        response = requests.get(
            files_url,
            params={"recursive": "1"},
            headers=headers,
            timeout=15,
        )
        if response.status_code != 200:
            log.debug(
                f"HuggingFace tree request returned {response.status_code} for repo={repo_id}, branch={branch}"
            )
            return None
        return response.json()
    except Exception as e:
        log.debug(f"Error getting HuggingFace tree for {repo_id}@{branch}: {e}")
        return None


def _find_matching_file_in_repo(
    repo_id: str,
    files: List[Dict[str, Any]],
    filename: str,
    exact_only: bool = False,
    headers: Optional[Dict[str, str]] = None,
) -> Optional[Dict[str, Any]]:
    filename_lower = filename.lower()
    filename_base = os.path.splitext(filename_lower)[0]
    partial_match = None

    for file_info in files:
        file_path = file_info.get("path", "")
        if not file_path:
            continue

        file_name = os.path.basename(file_path)
        file_name_lower = file_name.lower()
        file_base = os.path.splitext(file_name_lower)[0]

        if file_name_lower == filename_lower or file_path.lower().endswith(filename_lower):
            return _build_huggingface_result(
                repo_id,
                file_path,
                file_info,
                "exact",
                headers=headers,
            )

    if exact_only:
        return None

    for file_info in files:
        file_path = file_info.get("path", "")
        if not file_path:
            continue

        file_name = os.path.basename(file_path)
        file_name_lower = file_name.lower()
        file_base = os.path.splitext(file_name_lower)[0]

        if not exact_only:
            if filename_base in file_base or file_base in filename_base:
                if file_path.endswith(".safetensors") or file_path.endswith(".ckpt"):
                    partial_match = _build_huggingface_result(
                        repo_id,
                        file_path,
                        file_info,
                        "partial",
                        headers=headers,
                    )
                    break

    return partial_match


def _find_matching_file_in_author_index(
    index: Dict[str, Any],
    filename: str,
    exact_only: bool = False,
    headers: Optional[Dict[str, str]] = None,
) -> Optional[Dict[str, Any]]:
    files = index.get("files") or []
    filename_lower = filename.lower()
    filename_base = os.path.splitext(filename_lower)[0]
    partial_match = None

    for file_info in files:
        repo_id = file_info.get("repo_id", "")
        file_path = file_info.get("path", "")
        if not repo_id or not file_path:
            continue

        file_name = os.path.basename(file_path)
        file_name_lower = file_name.lower()
        if file_name_lower == filename_lower or file_path.lower().endswith(
            filename_lower
        ):
            return _build_huggingface_result(
                repo_id,
                file_path,
                file_info,
                "exact",
                headers=headers,
            )

    if exact_only:
        return None

    for file_info in files:
        repo_id = file_info.get("repo_id", "")
        file_path = file_info.get("path", "")
        if not repo_id or not file_path:
            continue

        file_name = os.path.basename(file_path)
        file_name_lower = file_name.lower()
        file_base = os.path.splitext(file_name_lower)[0]
        file_path_lower = file_path.lower()

        if filename_base in file_base or file_base in filename_base:
            if file_path_lower.endswith(".safetensors") or file_path_lower.endswith(
                ".ckpt"
            ):
                partial_match = _build_huggingface_result(
                    repo_id,
                    file_path,
                    file_info,
                    "partial",
                    headers=headers,
                )
                break

    return partial_match


def _build_huggingface_search_queries(filename: str) -> List[str]:
    filename_base = os.path.splitext(os.path.basename(filename))[0].lower()
    cleaned_base = clean_filename_for_search(filename).lower()
    simplified_base = re.sub(
        r"[-_]?(fp16|fp32|fp8|fp4|bf16|e4m3fn|mixed|scaled|pruned|emaonly|q4|q8)$",
        "",
        filename_base,
        flags=re.IGNORECASE,
    )

    queries = []
    for query in [filename_base, cleaned_base, simplified_base]:
        query = (query or "").strip(" -_")
        if query and query not in queries:
            queries.append(query)

    return queries


def _get_repos_by_author(
    author: str, headers: Dict[str, str], limit: int = 200
) -> List[str]:
    try:
        response = requests.get(
            f"{HF_API_URL}/models",
            params={"author": author, "limit": limit},
            headers=headers,
            timeout=15,
        )
        if response.status_code != 200:
            log.debug(
                f"HuggingFace author search returned {response.status_code} for author={author}"
            )
            return []

        repos = response.json()
        return [repo.get("id", "") for repo in repos if repo.get("id")]
    except Exception as e:
        log.debug(f"Error getting HuggingFace repos for author={author}: {e}")
        return []


def _search_brave_for_huggingface_candidates(
    filename: str, brave_api_key: str
) -> List[Dict[str, str]]:
    if not brave_api_key:
        return []

    query = f'"{filename}" site:huggingface.co'
    headers = {
        "Accept": "application/json",
        "Accept-Encoding": "gzip",
        "X-Subscription-Token": brave_api_key,
    }
    params = {
        "q": query,
        "count": 20,
        "safesearch": "off",
        "spellcheck": "false",
    }
    candidates = []
    seen = set()

    try:
        response = requests.get(
            BRAVE_SEARCH_API_URL, headers=headers, params=params, timeout=15
        )
        if response.status_code != 200:
            log.warn(
                f"Brave search returned {response.status_code} for query={query}"
            )
            return []

        payload = response.json()
        results = ((payload or {}).get("web") or {}).get("results") or []
        for result in results:
            candidate_url = result.get("url", "")
            if not candidate_url:
                continue

            parsed = parse_huggingface_url(candidate_url)
            if parsed:
                repo_id = parsed.get("repo", "")
                file_path = parsed.get("filename", "")
                branch = parsed.get("branch", "main")
            else:
                parsed_url = urlparse(candidate_url)
                if "huggingface.co" not in parsed_url.netloc:
                    continue
                path_parts = [part for part in parsed_url.path.split("/") if part]
                if len(path_parts) < 2:
                    continue
                repo_id = f"{path_parts[0]}/{path_parts[1]}"
                file_path = filename
                branch = "main"

            key = (repo_id, branch, file_path)
            if not repo_id or key in seen:
                continue

            seen.add(key)
            candidates.append(
                {
                    "repo": repo_id,
                    "branch": branch,
                    "filename": file_path,
                    "url": candidate_url,
                }
            )
    except Exception as e:
        log.debug(f"Brave search error for filename={filename}: {e}")

    return candidates


def search_huggingface_for_file(
    filename: str,
    token: Optional[str] = None,
    exact_only: bool = False,
    brave_api_key: Optional[str] = None,
    use_api_search: bool = True,
    use_comfy_org_fallback: bool = True,
    use_brave_fallback: bool = True,
    force_refresh: bool = False,
    progress_callback: Optional[Callable[[Dict[str, Any]], None]] = None,
) -> Optional[Dict[str, Any]]:
    """
    Search HuggingFace for a specific model file.
    Returns the first repo that contains a matching file.

    Args:
        filename: Filename to search for
        token: Optional HF token
        exact_only: If True, only return exact filename matches (for downloads).
                   If False, also try partial matching (for local file resolution).

    Returns:
        Dict with url, repo, filename if found, None otherwise
    """
    global _search_cache

    brave_key = "brave" if brave_api_key and use_brave_fallback else "nobrave"
    token_key = "token" if token else "notoken"
    methods_key = (
        f"api{int(bool(use_api_search))}_"
        f"comfy{int(bool(use_comfy_org_fallback))}_"
        f"brave{int(bool(use_brave_fallback))}_"
        f"force{int(bool(force_refresh))}"
    )
    cache_key = f"hf_{filename}_exact{exact_only}_{token_key}_{brave_key}_{methods_key}"
    if cache_key in _search_cache:
        log.debug(f"HuggingFace cache hit file={filename} exact={exact_only}")
        _report_progress(
            progress_callback,
            "cache",
            "Using HuggingFace cache",
            86,
        )
        return _search_cache[cache_key]

    try:
        headers = {}
        if token:
            headers["Authorization"] = f"Bearer {token}"

        log.info(f"HuggingFace start file={filename} exact={exact_only}")

        repos_to_check = []
        seen_repos = set()
        search_queries = _build_huggingface_search_queries(filename)

        if use_api_search:
            total_queries = max(1, len(search_queries))
            for query_index, search_query in enumerate(search_queries, start=1):
                _report_progress(
                    progress_callback,
                    "api_search",
                    f"HuggingFace API query {query_index}/{total_queries}",
                    32 + ((query_index - 1) / total_queries) * 10,
                    query=search_query,
                    query_index=query_index,
                    query_count=total_queries,
                )
                search_url = f"{HF_API_URL}/models?search={quote(search_query)}&limit=20"
                response = requests.get(search_url, headers=headers, timeout=10)
                if response.status_code != 200:
                    log.warn(
                        f"HuggingFace API status={response.status_code} query={search_query}"
                    )
                    continue

                repos = response.json()
                log.info(f"HuggingFace API query={search_query} repos={len(repos)}")

                for repo in repos:
                    repo_id = repo.get("id", "")
                    if repo_id and repo_id not in seen_repos:
                        seen_repos.add(repo_id)
                        repos_to_check.append(repo_id)

                _report_progress(
                    progress_callback,
                    "api_candidates",
                    f"HuggingFace API candidates: {len(repos_to_check)} repos",
                    44,
                    candidate_count=len(repos_to_check),
                )
        else:
            _report_progress(
                progress_callback,
                "api_skip",
                "Skipping HuggingFace API search",
                44,
            )

        repo_count = len(repos_to_check)
        for repo_index, repo_id in enumerate(repos_to_check, start=1):
            _report_progress(
                progress_callback,
                "repo_files",
                f"Checking HF repo {repo_index}/{repo_count}",
                48 + (repo_index / max(1, repo_count)) * 14,
                repo=repo_id,
                candidate_index=repo_index,
                candidate_count=repo_count,
            )
            log.debug(f"HuggingFace check repo={repo_id} file={filename}")
            files = _get_repo_tree(repo_id, headers=headers)
            if not files:
                continue

            log.debug(f"HuggingFace tree repo={repo_id} files={len(files)}")
            result = _find_matching_file_in_repo(
                repo_id,
                files,
                filename,
                exact_only=exact_only,
                headers=headers,
            )
            if result:
                _search_cache[cache_key] = result
                _report_progress(
                    progress_callback,
                    "found",
                    "Found HuggingFace match",
                    92,
                    repo=repo_id,
                    match_type=result.get("match_type"),
                )
                log.info(
                    f"HuggingFace found match={result['match_type']} file={filename} repo={repo_id} path={result['path']}"
                )
                return result

        author_fallback_repos = []
        author_fallback_repo_count = 0
        if use_comfy_org_fallback:
            total_authors = max(1, len(HF_AUTHOR_FALLBACKS))
            for author_index, author in enumerate(HF_AUTHOR_FALLBACKS, start=1):
                _report_progress(
                    progress_callback,
                    "author_index",
                    f"Checking HF fallback index {author_index}/{total_authors}",
                    64 + ((author_index - 1) / total_authors) * 6,
                    author=author,
                    author_index=author_index,
                    author_count=total_authors,
                )
                index = _get_author_index(
                    author,
                    headers={},
                    force_refresh=force_refresh,
                )
                if index:
                    result = _find_matching_file_in_author_index(
                        index,
                        filename,
                        exact_only=exact_only,
                        headers={},
                    )
                    author_fallback_repo_count += int(index.get("repo_count") or 0)
                    if result:
                        _search_cache[cache_key] = result
                        _report_progress(
                            progress_callback,
                            "found",
                            "Found HuggingFace fallback match",
                            92,
                            repo=result.get("repo_id"),
                            match_type=result.get("match_type"),
                        )
                        log.info(
                            f"HuggingFace found source=author_index match={result['match_type']} file={filename} repo={result['repo_id']} path={result['path']}"
                        )
                        return result

                    log.debug(
                        f"HuggingFace author_index miss author={author} file={filename} repos={index.get('repo_count', 0)} files={index.get('file_count', 0)}"
                    )
                    continue

                repos = _get_repos_by_author(author, headers=headers)
                log.info(f"HuggingFace author_fallback author={author} repos={len(repos)}")
                for repo_id in repos:
                    if repo_id not in seen_repos:
                        seen_repos.add(repo_id)
                        author_fallback_repos.append(repo_id)
                        author_fallback_repo_count += 1
        else:
            _report_progress(
                progress_callback,
                "author_skip",
                "Skipping HuggingFace fallback index",
                70,
            )

        fallback_repo_count = len(author_fallback_repos)
        for repo_index, repo_id in enumerate(author_fallback_repos, start=1):
            _report_progress(
                progress_callback,
                "author_repo",
                f"Checking HF fallback repo {repo_index}/{fallback_repo_count}",
                70 + (repo_index / max(1, fallback_repo_count)) * 8,
                repo=repo_id,
                candidate_index=repo_index,
                candidate_count=fallback_repo_count,
            )
            log.debug(
                f"HuggingFace author_fallback check repo={repo_id} file={filename}"
            )
            files = _get_repo_tree(repo_id, headers=headers)
            if not files:
                continue

            result = _find_matching_file_in_repo(
                repo_id,
                files,
                filename,
                exact_only=exact_only,
                headers=headers,
            )
            if result:
                _search_cache[cache_key] = result
                _report_progress(
                    progress_callback,
                    "found",
                    "Found HuggingFace fallback match",
                    92,
                    repo=repo_id,
                    match_type=result.get("match_type"),
                )
                log.info(
                    f"HuggingFace found source=author_fallback match={result['match_type']} file={filename} repo={repo_id} path={result['path']}"
                )
                return result

        if use_brave_fallback:
            _report_progress(
                progress_callback,
                "brave_search",
                "Searching Brave fallback",
                80,
            )
        brave_candidates = (
            _search_brave_for_huggingface_candidates(filename, brave_api_key or "")
            if use_brave_fallback
            else []
        )
        log.info(f"HuggingFace brave_fallback file={filename} candidates={len(brave_candidates)}")
        _report_progress(
            progress_callback,
            "brave_candidates",
            f"Brave candidates: {len(brave_candidates)}",
            84,
            candidate_count=len(brave_candidates),
        )

        brave_count = len(brave_candidates)
        for candidate_index, candidate in enumerate(brave_candidates, start=1):
            repo_id = candidate.get("repo", "")
            branch = candidate.get("branch", "main")
            expected_file = candidate.get("filename", filename)
            _report_progress(
                progress_callback,
                "brave_repo",
                f"Checking Brave candidate {candidate_index}/{brave_count}",
                84 + (candidate_index / max(1, brave_count)) * 6,
                repo=repo_id,
                candidate_index=candidate_index,
                candidate_count=brave_count,
            )
            files = _get_repo_tree(repo_id, headers=headers, branch=branch)
            if not files:
                continue

            result = _find_matching_file_in_repo(
                repo_id,
                files,
                expected_file,
                exact_only=True,
                headers=headers,
            )
            if result and result.get("filename", "").lower() == filename.lower():
                _search_cache[cache_key] = result
                _report_progress(
                    progress_callback,
                    "found",
                    "Found HuggingFace Brave match",
                    92,
                    repo=repo_id,
                    match_type=result.get("match_type"),
                )
                log.info(
                    f"HuggingFace found source=brave_fallback match=exact file={filename} repo={repo_id} path={result['path']}"
                )
                return result

        # Not found
        _search_cache[cache_key] = None
        _report_progress(
            progress_callback,
            "done",
            "HuggingFace checked",
            92,
            candidate_count=len(repos_to_check),
            author_repo_count=author_fallback_repo_count,
            brave_candidate_count=len(brave_candidates),
        )
        log.info(
            f"HuggingFace miss file={filename} repos={len(repos_to_check)} author_repos={author_fallback_repo_count} brave={len(brave_candidates)}"
        )
        return None

    except Exception as e:
        log.exception(f"HuggingFace search error for {filename}: {e}")
        _report_progress(
            progress_callback,
            "error",
            f"HuggingFace search error: {e}",
            100,
            status="error",
        )
        return None


def search_huggingface(
    query: str,
    model_type: Optional[str] = None,
    limit: int = 10,
    token: Optional[str] = None,
) -> List[Dict[str, Any]]:
    """
    Search HuggingFace Hub for models (general search).
    Returns repos that might be relevant, not guaranteed to have exact file.
    """
    results = []

    try:
        headers = {}
        if token:
            headers["Authorization"] = f"Bearer {token}"

        params = {"search": query, "limit": limit, "full": "true"}

        response = requests.get(
            f"{HF_API_URL}/models", params=params, headers=headers, timeout=15
        )

        if response.status_code == 200:
            models = response.json()

            for model in models:
                repo_id = model.get("id", "")

                results.append(
                    {
                        "source": "huggingface",
                        "repo": repo_id,
                        "name": model.get("modelId", repo_id),
                        "downloads": model.get("downloads", 0),
                        "likes": model.get("likes", 0),
                        "url": f"https://huggingface.co/{repo_id}",
                    }
                )

    except Exception as e:
        log.error(f"HuggingFace search error: {e}")

    return results


def get_repo_files(repo: str, token: Optional[str] = None) -> List[Dict[str, Any]]:
    """Get list of model files in a HuggingFace repo."""
    files = []

    try:
        headers = {}
        if token:
            headers["Authorization"] = f"Bearer {token}"

        items = _get_repo_tree(repo, headers=headers)
        if items:
            for item in items:
                path = item.get("path", "")
                if path.endswith(
                    (".safetensors", ".ckpt", ".pt", ".bin", ".pth", ".onnx", ".gguf")
                ):
                    url = get_huggingface_download_url(repo, path)
                    size = _extract_file_info_size(item)
                    if not size:
                        size = _fetch_remote_file_size_bytes(url, headers=headers)
                    files.append(
                        {
                            "filename": os.path.basename(path),
                            "path": path,
                            "url": url,
                            "size": size,
                        }
                    )

    except Exception as e:
        log.error(f"Error getting repo files: {e}")

    return files