"""HuggingFace Hub client for model and dataset info."""

from __future__ import annotations

import logging
from dataclasses import dataclass

logger = logging.getLogger(__name__)


@dataclass
class ModelInfo:
    """Basic info about a HuggingFace model."""

    model_id: str
    size_bytes: int | None = None
    size_gb: float | None = None
    architecture: str | None = None
    license: str | None = None
    downloads: int | None = None


class HubClient:
    """Lightweight client for HuggingFace Hub API."""

    def __init__(self) -> None:
        self._api = None

    def _get_api(self):
        """Lazy-load huggingface_hub API."""
        if self._api is None:
            try:
                from huggingface_hub import HfApi

                self._api = HfApi()
            except ImportError:
                raise ImportError(
                    "huggingface_hub is required for Hub features. "
                    "Install with: pip install huggingface-hub"
                ) from None
        return self._api

    def get_model_info(self, model_id: str) -> ModelInfo | None:
        """Fetch model info from HuggingFace Hub."""
        try:
            api = self._get_api()
            info = api.model_info(model_id)

            size_bytes = None
            size_gb = None
            if info.safetensors:
                size_bytes = sum(
                    v.get("total", 0) if isinstance(v, dict) else 0
                    for v in info.safetensors.values()
                    if isinstance(v, dict)
                )
            elif hasattr(info, "siblings") and info.siblings:
                size_bytes = sum(
                    s.size
                    for s in info.siblings
                    if s.size and s.rfilename.endswith((".bin", ".safetensors", ".pt"))
                )

            if size_bytes and size_bytes > 0:
                size_gb = size_bytes / (1024**3)

            arch = None
            if info.config and "architectures" in info.config:
                arch = info.config["architectures"][0]

            return ModelInfo(
                model_id=model_id,
                size_bytes=size_bytes,
                size_gb=size_gb,
                architecture=arch,
                license=info.card_data.get("license") if info.card_data else None,
                downloads=info.downloads,
            )
        except ImportError:
            raise
        except Exception:
            logger.debug("Failed to fetch model info for %s", model_id, exc_info=True)
            return None

    def validate_model_exists(self, model_id: str) -> bool:
        """Check if a model exists on HuggingFace Hub."""
        try:
            api = self._get_api()
            api.model_info(model_id)
            return True
        except ImportError:
            raise
        except Exception:
            return False

    def validate_dataset_exists(self, dataset_id: str) -> bool:
        """Check if a dataset exists on HuggingFace Hub."""
        try:
            api = self._get_api()
            api.dataset_info(dataset_id)
            return True
        except ImportError:
            raise
        except Exception:
            return False

    def estimate_model_vram(self, model_id: str, quantization: str = "none") -> float | None:
        """Estimate VRAM needed to load a model.

        Returns estimated GB needed, or None if info not available.
        """
        info = self.get_model_info(model_id)
        if info is None or info.size_gb is None:
            return None

        quant_factor = {"4bit": 0.25, "8bit": 0.5, "none": 1.0}.get(quantization, 1.0)

        # Model weights + ~20% overhead for activations/optimizer
        return info.size_gb * quant_factor * 1.2
