"""
Model management functions for nexaai library.

This module provides functions for downloading, listing, and managing models
from HuggingFace Hub and the local store.
"""

import logging
import os
from typing import Callable, List, Optional

from .internal.model_hub import ensure_model_downloaded
from .internal.store import Store
from .internal.types import DownloadProgressInfo, ModelInfo


def download_model(
    repo_id: str,
    quant_spec: Optional[str] = None,
    token: Optional[str] = os.getenv('HF_TOKEN', None),
    progress_callback: Optional[Callable[[DownloadProgressInfo], None]] = None,
) -> None:
    """
    Download a model from HuggingFace Hub.

    Args:
        repo_id: HuggingFace repository ID (e.g., "microsoft/phi-3-mini-4k-instruct")
        quant_spec: Optional quant name (e.g., "Q4_K_M") to download specific quantization
        token: Optional HuggingFace token for private repositories, default to environment variable HF_TOKEN
        progress_callback: Optional callback function for download progress.
            The callback receives a DownloadProgressInfo dict mapping filename to FileProgressInfo.
            Structure: {filename: {downloaded_bytes: int, total_bytes: int}}

    Example:
        >>> def progress(info):
        ...     for filename, file_info in info.items():
        ...         pct = (file_info['downloaded_bytes'] / file_info['total_bytes'] * 100
        ...                if file_info['total_bytes'] > 0 else 0)
        ...         print(f"{filename}: {pct:.1f}%")
        >>> nexaai.download_model("microsoft/phi-3-mini-4k-instruct", progress_callback=progress)
    """
    logger = logging.getLogger('nexaai')

    # Use unified internal function to ensure consistent behavior
    manifest = ensure_model_downloaded(
        repo_id=repo_id,
        quant_spec=quant_spec,
        token=token,
        progress_callback=progress_callback,
    )

    if manifest is not None:
        logger.info(f'Model {repo_id} is ready (already existed or downloaded successfully)')


def list_models() -> List[ModelInfo]:
    """
    List all models downloaded in the local store.

    Returns:
        List of ModelInfo objects containing information about each model

    Example:
        >>> models = nexaai.list_models()
        >>> for model in models:
        ...     print(f"{model.repo_id}: {model.model_type} ({model.size / 1024**3:.2f} GB)")
    """
    store = Store.get()
    manifests = store.list_models()
    return [ModelInfo.from_manifest(manifest) for manifest in manifests]


def remove_model(repo_id: str) -> bool:
    """
    Remove a model from the local store.

    Args:
        repo_id: HuggingFace repository ID (e.g., "microsoft/phi-3-mini-4k-instruct")

    Returns:
        True if the model was removed successfully, False otherwise

    Example:
        >>> nexaai.remove_model("microsoft/phi-3-mini-4k-instruct")
        True
    """
    store = Store.get()
    return store.remove_model(repo_id)
