"""
HuggingFace Hub integration for personality classifier models.

Provides functions to push trained models to and download from HuggingFace Hub.
"""

from __future__ import annotations

import json
import logging
from datetime import datetime
from pathlib import Path
from typing import Any

logger = logging.getLogger(__name__)


# Default HuggingFace repository for Lindr models
DEFAULT_HF_REPO = "lindr-ai/personality-classifier"


def push_to_hub(
    onnx_path: str | Path,
    version: str,
    training_config: dict | None = None,
    metrics: dict | None = None,
    repo_id: str = DEFAULT_HF_REPO,
    token: str | None = None,
    private: bool = True,
) -> str:
    """
    Push trained ONNX model to HuggingFace Hub.

    Args:
        onnx_path: Path to the ONNX model file
        version: Version string (e.g., 'v1.0.0')
        training_config: Training configuration to store as metadata
        metrics: Training metrics to store as metadata
        repo_id: HuggingFace repository ID (default: lindr-ai/personality-classifier)
        token: HuggingFace API token (or use HF_TOKEN env var)
        private: Whether the repository should be private

    Returns:
        URL to the uploaded model
    """
    try:
        from huggingface_hub import HfApi, create_repo, upload_file
    except ImportError:
        raise ImportError(
            "huggingface_hub is required for Hub integration. "
            "Install with: pip install huggingface_hub"
        )

    onnx_path = Path(onnx_path)
    if not onnx_path.exists():
        raise FileNotFoundError(f"ONNX model not found: {onnx_path}")

    api = HfApi(token=token)

    # Ensure repository exists
    try:
        api.repo_info(repo_id=repo_id)
        logger.info(f"Repository exists: {repo_id}")
    except Exception:
        logger.info(f"Creating repository: {repo_id}")
        create_repo(
            repo_id=repo_id,
            repo_type="model",
            private=private,
            token=token,
        )

    # Upload versioned model
    versioned_filename = f"personality-classifier-{version}.onnx"
    logger.info(f"Uploading {versioned_filename}...")

    upload_file(
        path_or_fileobj=str(onnx_path),
        path_in_repo=f"models/{versioned_filename}",
        repo_id=repo_id,
        token=token,
        commit_message=f"Add personality classifier {version}",
    )

    # Also upload as "latest"
    logger.info("Updating latest model...")
    upload_file(
        path_or_fileobj=str(onnx_path),
        path_in_repo="models/personality-classifier-latest.onnx",
        repo_id=repo_id,
        token=token,
        commit_message=f"Update latest to {version}",
    )

    # Create and upload metadata
    metadata = _create_metadata(version, training_config, metrics, onnx_path)
    metadata_json = json.dumps(metadata, indent=2)

    # Upload versioned metadata
    upload_file(
        path_or_fileobj=metadata_json.encode("utf-8"),
        path_in_repo=f"models/metadata-{version}.json",
        repo_id=repo_id,
        token=token,
        commit_message=f"Add metadata for {version}",
    )

    # Update latest metadata
    upload_file(
        path_or_fileobj=metadata_json.encode("utf-8"),
        path_in_repo="models/metadata-latest.json",
        repo_id=repo_id,
        token=token,
        commit_message=f"Update latest metadata to {version}",
    )

    model_url = f"https://huggingface.co/{repo_id}/blob/main/models/{versioned_filename}"
    logger.info(f"Model uploaded successfully: {model_url}")

    return model_url


def _create_metadata(
    version: str,
    training_config: dict | None,
    metrics: dict | None,
    onnx_path: Path,
) -> dict:
    """Create metadata dictionary for the model."""
    return {
        "version": version,
        "created_at": datetime.utcnow().isoformat() + "Z",
        "model_type": "personality-classifier",
        "architecture": "MLP",
        "input_dim": 384,
        "output_dim": 10,
        "embedding_model": "Xenova/all-MiniLM-L6-v2",
        "dimensions": [
            "openness",
            "conscientiousness",
            "extraversion",
            "agreeableness",
            "neuroticism",
            "assertiveness",
            "ambition",
            "resilience",
            "integrity",
            "curiosity",
        ],
        "file_size_bytes": onnx_path.stat().st_size,
        "training_config": training_config or {},
        "metrics": metrics or {},
    }


def download_model(
    version: str = "latest",
    cache_dir: str | Path | None = None,
    repo_id: str = DEFAULT_HF_REPO,
    token: str | None = None,
) -> Path:
    """
    Download ONNX model from HuggingFace Hub.

    Args:
        version: Model version ('latest' or specific version like 'v1.0.0')
        cache_dir: Directory to cache downloaded models
        repo_id: HuggingFace repository ID
        token: HuggingFace API token (for private repos)

    Returns:
        Path to the downloaded ONNX model
    """
    try:
        from huggingface_hub import hf_hub_download
    except ImportError:
        raise ImportError(
            "huggingface_hub is required for Hub integration. "
            "Install with: pip install huggingface_hub"
        )

    filename = f"personality-classifier-{version}.onnx"

    logger.info(f"Downloading model: {repo_id}/models/{filename}")

    model_path = hf_hub_download(
        repo_id=repo_id,
        filename=f"models/{filename}",
        cache_dir=cache_dir,
        token=token,
    )

    logger.info(f"Model downloaded: {model_path}")
    return Path(model_path)


def download_metadata(
    version: str = "latest",
    cache_dir: str | Path | None = None,
    repo_id: str = DEFAULT_HF_REPO,
    token: str | None = None,
) -> dict:
    """
    Download model metadata from HuggingFace Hub.

    Args:
        version: Model version
        cache_dir: Directory to cache downloaded files
        repo_id: HuggingFace repository ID
        token: HuggingFace API token

    Returns:
        Metadata dictionary
    """
    try:
        from huggingface_hub import hf_hub_download
    except ImportError:
        raise ImportError(
            "huggingface_hub is required for Hub integration. "
            "Install with: pip install huggingface_hub"
        )

    filename = f"metadata-{version}.json"

    metadata_path = hf_hub_download(
        repo_id=repo_id,
        filename=f"models/{filename}",
        cache_dir=cache_dir,
        token=token,
    )

    with open(metadata_path, "r", encoding="utf-8") as f:
        return json.load(f)


def list_versions(
    repo_id: str = DEFAULT_HF_REPO,
    token: str | None = None,
) -> list[str]:
    """
    List available model versions in the repository.

    Args:
        repo_id: HuggingFace repository ID
        token: HuggingFace API token

    Returns:
        List of version strings
    """
    try:
        from huggingface_hub import HfApi
    except ImportError:
        raise ImportError(
            "huggingface_hub is required for Hub integration. "
            "Install with: pip install huggingface_hub"
        )

    api = HfApi(token=token)

    try:
        files = api.list_repo_files(repo_id=repo_id)
    except Exception as e:
        logger.warning(f"Failed to list repository files: {e}")
        return []

    # Extract versions from model filenames
    versions = []
    for f in files:
        if f.startswith("models/personality-classifier-") and f.endswith(".onnx"):
            # Extract version from filename
            name = f.replace("models/personality-classifier-", "").replace(".onnx", "")
            if name != "latest":
                versions.append(name)

    return sorted(versions, reverse=True)
