import os
import json
from numbers import Number
from huggingface_hub import HfApi, hf_hub_download
from huggingface_hub.utils import RepositoryNotFoundError, GatedRepoError, HfHubHTTPError, EntryNotFoundError

class ModelFetchError(Exception):
    """Custom exception for errors during standard model fetching."""
    pass

def get_model_config(repo_id: str) -> dict:
    """
    Attempts to download and parse the config.json file from the Hugging Face model repository.
    Returns the parsed JSON as a dictionary. Returns an empty dict if it fails.
    """
    try:
        token = os.environ.get("HUGGING_FACE_HUB_TOKEN", os.environ.get("HF_TOKEN"))
        config_path = hf_hub_download(repo_id=repo_id, filename="config.json", token=token)
        with open(config_path, 'r', encoding='utf-8') as f:
            return json.load(f)
    except (RepositoryNotFoundError, GatedRepoError, EntryNotFoundError, HfHubHTTPError, ValueError):
        # We silently fail configuration fetching if it's missing or gated
        # The parameter fetcher will throw the real error
        return {}

def get_model_parameters(repo_id: str) -> float:
    """
    Fetches the number of parameters for a Hugging Face model.
    Returns the parameter count in billions (B).
    """
    api = HfApi()
    
    try:
        # Fetch model info. We can pass a token from the environment if the user set HUGGING_FACE_HUB_TOKEN
        token = os.environ.get("HUGGING_FACE_HUB_TOKEN", os.environ.get("HF_TOKEN"))
        info = api.model_info(repo_id, token=token)
        
        # 1. First, check if safetensors metadata is available and gives total parameters
        # huggingface_hub > 0.19 exposes `safetensors` attribute if Safetensors is used.
        if getattr(info, "safetensors", None) is not None:
            if isinstance(info.safetensors, dict) and "total" in info.safetensors:
                total_params = info.safetensors["total"]
                if isinstance(total_params, Number) and total_params > 0:
                    return float(total_params) / 1e9
            elif hasattr(info.safetensors, "total") and getattr(info.safetensors, "total", 0) > 0:
                return float(info.safetensors.total) / 1e9

        # 2. Fallback: calculate from file sizes.
        # We look for model weight files (.safetensors, .bin, .pt)
        weight_files = []
        for sibling in info.siblings:
            fname = sibling.rfilename
            # Exclude fp8/int4 variations stored separately if possible, 
            # though usually the main ones are the default.
            if fname.endswith(".safetensors") or fname.endswith(".bin") or fname.endswith(".pt"):
                weight_files.append(fname)
        
        if not weight_files:
            raise ModelFetchError(f"No weight files (.safetensors, .bin, .pt) found in repository '{repo_id}'.")
            
        # Sum the sizes of the weight files
        # HfApi doesn't give size in `info.siblings`, we need to inspect tree or model_info with `files_metadata=True`
        info = api.model_info(repo_id, token=token, files_metadata=True)
        total_size_bytes = 0
        for sibling in info.siblings:
            fname = sibling.rfilename
            if fname in weight_files:
                total_size_bytes += getattr(sibling, "size", 0)
                
        if total_size_bytes == 0:
            raise ModelFetchError("Could not determine the size of the model weights.")
            
        # Approximation: assume files are stored in float16 (2 bytes per parameter)
        approx_params = total_size_bytes / 2.0
        return approx_params / 1e9

    except GatedRepoError:
        raise ModelFetchError(
            f"Repository '{repo_id}' is gated. You must accept the conditions on Hugging Face "
            "and set the HUGGING_FACE_HUB_TOKEN environment variable."
        )
    except RepositoryNotFoundError:
        raise ModelFetchError(f"Repository '{repo_id}' was not found on Hugging Face.")
    except HfHubHTTPError as e:
        raise ModelFetchError(f"HTTP Error while connecting to Hugging Face: {str(e)}")
    except Exception as e:
        raise ModelFetchError(f"An unexpected error occurred: {str(e)}")
