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

from .base import BaseSource, ResolvedSource, ResolutionResult


class HuggingFaceSource(BaseSource):
    """HuggingFace Hub source for resolving model assets."""
    
    def __init__(self, token: str | None = None):
        """
        Initialize HuggingFace source.
        
        Args:
            token: Optional HuggingFace API token for accessing gated models
        """
        self._token = token
        self._api_base = "https://huggingface.co/api"
    
    @property
    def name(self) -> str:
        return "huggingface"
    
    def is_available(self, config: dict[str, Any]) -> bool:
        """
        HuggingFace search is always available (public search needs no token).
        
        Args:
            config: Configuration dictionary
            
        Returns:
            Always True
        """
        return True
    
    def search(self, filename: str, asset_type: str) -> list[ResolutionResult]:
        """
        Search HuggingFace Hub for an asset by filename.
        
        Args:
            filename: The asset filename to search for
            asset_type: The type of asset (e.g., 'checkpoints', 'loras')
            
        Returns:
            List of ResolutionResult objects, sorted by confidence (highest first)
        """
        results: list[ResolutionResult] = []
        
        try:
            # Search for the filename across HuggingFace Hub
            # We'll search for repos that might contain this file
            search_url = f"{self._api_base}/models"
            
            # Extract base name without extension for better search
            base_name = filename.rsplit(".", 1)[0] if "." in filename else filename
            
            params = {
                "search": base_name,
                "limit": 20,  # Get more results to filter through
            }
            
            headers = {}
            if self._token:
                headers["Authorization"] = f"Bearer {self._token}"
            
            response = requests.get(search_url, params=params, headers=headers, timeout=10)
            
            if response.status_code == 401:
                return [
                    ResolutionResult(
                        input_name=filename,
                        resolved_name=filename,
                        source=None,
                        confidence="unknown",
                        error="HuggingFace token invalid or expired",
                    )
                ]
            
            if response.status_code != 200:
                return [
                    ResolutionResult(
                        input_name=filename,
                        resolved_name=filename,
                        source=None,
                        confidence="unknown",
                        error=f"HuggingFace API error: {response.status_code}",
                    )
                ]
            
            repos_data = response.json()
            
            # For each repo, check if it contains the file
            for repo in repos_data[:10]:  # Limit to first 10 repos to avoid too many API calls
                repo_id = repo.get("id") or repo.get("modelId")
                if not repo_id:
                    continue
                
                # Try to get file list for this repo
                try:
                    files = self._list_repo_files(repo_id)
                    
                    # Check if our filename exists in this repo
                    for file_path in files:
                        file_name = file_path.split("/")[-1]  # Get just the filename
                        
                        # Calculate similarity
                        similarity = fuzz.ratio(filename.lower(), file_name.lower())
                        
                        # Determine confidence based on similarity
                        if filename.lower() == file_name.lower():
                            confidence = "high"
                        elif similarity >= 85:
                            confidence = "medium"
                        else:
                            confidence = "low"
                        
                        # Only include results with at least low confidence
                        if similarity >= 70:
                            # Get file metadata
                            metadata = self._get_file_metadata_safe(repo_id, file_path)
                            
                            source = ResolvedSource(
                                source_type="huggingface",
                                repo=repo_id,
                                file=file_path,
                                download_url=f"https://huggingface.co/{repo_id}/resolve/main/{file_path}",
                                sha256=metadata.get("sha256"),
                                file_size_bytes=metadata.get("size"),
                            )
                            
                            results.append(
                                ResolutionResult(
                                    input_name=filename,
                                    resolved_name=file_name,
                                    source=source,
                                    confidence=confidence,
                                    error=None,
                                )
                            )
                
                except Exception:
                    # Skip repos that fail to list
                    continue
            
            # Sort by confidence (high > medium > low) and then by similarity
            confidence_order = {"high": 3, "medium": 2, "low": 1, "unknown": 0}
            results.sort(
                key=lambda r: (
                    confidence_order.get(r.confidence, 0),
                    fuzz.ratio(filename.lower(), r.resolved_name.lower()),
                ),
                reverse=True,
            )
            
            # Return top 5 results
            return results[:5]
        
        except requests.exceptions.RequestException as e:
            return [
                ResolutionResult(
                    input_name=filename,
                    resolved_name=filename,
                    source=None,
                    confidence="unknown",
                    error=f"Network error: {str(e)}",
                )
            ]
        except Exception as e:
            return [
                ResolutionResult(
                    input_name=filename,
                    resolved_name=filename,
                    source=None,
                    confidence="unknown",
                    error=f"Unexpected error: {str(e)}",
                )
            ]
    
    def _list_repo_files(self, repo_id: str) -> list[str]:
        """
        List all files in a HuggingFace repository.
        
        Args:
            repo_id: Repository ID (e.g., 'black-forest-labs/FLUX.1-dev')
            
        Returns:
            List of file paths in the repository
        """
        url = f"{self._api_base}/models/{repo_id}/tree/main"
        
        headers = {}
        if self._token:
            headers["Authorization"] = f"Bearer {self._token}"
        
        response = requests.get(url, headers=headers, timeout=10)
        
        if response.status_code == 403:
            # Gated model - provide helpful error
            raise Exception(
                f"Model {repo_id} is gated. Please visit https://huggingface.co/{repo_id} "
                "to request access and ensure your token has the required permissions."
            )
        
        response.raise_for_status()
        
        files = []
        for item in response.json():
            if item.get("type") == "file":
                files.append(item.get("path", ""))
        
        return files
    
    def _get_file_metadata_safe(self, repo_id: str, filename: str) -> dict[str, Any]:
        """
        Safely get file metadata without raising exceptions.
        
        Args:
            repo_id: Repository ID
            filename: File path within the repository
            
        Returns:
            Dictionary with 'sha256' and 'size' keys if available
        """
        try:
            return self.get_file_metadata(repo_id, filename)
        except Exception:
            return {}
    
    def get_file_metadata(self, repo_id: str, filename: str) -> dict[str, Any]:
        """
        Get metadata for a specific file in a HuggingFace repository.
        
        Args:
            repo_id: Repository ID (e.g., 'runwayml/stable-diffusion-v1-5')
            filename: File path within the repository
            
        Returns:
            Dictionary containing 'sha256' and 'size' keys
            
        Raises:
            requests.HTTPError: If the API request fails
            Exception: For gated models with helpful message
        """
        # Use the HF Hub API to get file metadata
        url = f"https://huggingface.co/{repo_id}/resolve/main/{filename}"
        
        headers = {}
        if self._token:
            headers["Authorization"] = f"Bearer {self._token}"
        
        # Make a HEAD request to get metadata without downloading
        response = requests.head(url, headers=headers, allow_redirects=True, timeout=10)
        
        if response.status_code == 401:
            raise Exception("HuggingFace token invalid or expired")
        
        if response.status_code == 403:
            raise Exception(
                f"Model {repo_id} is gated. Please visit https://huggingface.co/{repo_id} "
                "to request access and ensure your token has the required permissions."
            )
        
        response.raise_for_status()
        
        metadata: dict[str, Any] = {}
        
        # Extract SHA-256 from headers if available
        # HuggingFace includes this in the X-Repo-Commit header or ETag
        if "X-Linked-Etag" in response.headers:
            etag = response.headers["X-Linked-Etag"].strip('"')
            # HF etags are often the SHA256 hash
            if len(etag) == 64:  # SHA256 is 64 hex chars
                metadata["sha256"] = etag
        elif "ETag" in response.headers:
            etag = response.headers["ETag"].strip('"')
            if len(etag) == 64:
                metadata["sha256"] = etag
        
        # Get file size from Content-Length header
        if "Content-Length" in response.headers:
            try:
                metadata["size"] = int(response.headers["Content-Length"])
            except ValueError:
                pass
        
        # Alternative: use the API endpoint for file info
        if not metadata.get("sha256"):
            try:
                api_url = f"{self._api_base}/models/{repo_id}/tree/main/{filename}"
                api_headers = {}
                if self._token:
                    api_headers["Authorization"] = f"Bearer {self._token}"
                
                api_response = requests.get(api_url, headers=api_headers, timeout=10)
                if api_response.status_code == 200:
                    file_info = api_response.json()
                    if "lfs" in file_info and "sha256" in file_info["lfs"]:
                        metadata["sha256"] = file_info["lfs"]["sha256"]
                    if "size" in file_info:
                        metadata["size"] = file_info["size"]
            except Exception:
                pass  # Fallback failed, return what we have
        
        return metadata
