"""Model loading with singleton pattern for efficient inference.

Supports loading from:
1. Local model directory (development)
2. HuggingFace Hub (production with auto-download)

Configuration via environment variables:
- DENES_USE_HF_HUB: Enable HF Hub (default: true)
- DENES_MODEL_REPO: HF repository ID (default: mlima001/denes-router)
- HUGGINGFACE_TOKEN: HF access token (optional, for private models)
"""

from pathlib import Path
from typing import Optional
import logging
import os

import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

logger = logging.getLogger(__name__)

# Singleton instances
_model: Optional[AutoModelForSequenceClassification] = None
_tokenizer: Optional[AutoTokenizer] = None
_loaded = False

# ⚠️ CONFIGURATION - Update with your HuggingFace repository ID after upload
HF_REPO_ID = os.getenv("DENES_MODEL_REPO", "mlima001/denes-router")
USE_HF_HUB = os.getenv("DENES_USE_HF_HUB", "true").lower() == "true"
HF_TOKEN = os.getenv("HUGGINGFACE_TOKEN")  # Optional: HF token from environment variable


def get_model_path() -> Path:
    """Get the path to the trained model directory.

    Returns:
        Path to models/domain-classifier/ (local or HF cache)
    """
    # Local path (development mode)
    return Path(__file__).parent.parent.parent / "models" / "domain-classifier"


def get_model() -> AutoModelForSequenceClassification:
    """Get loaded model (singleton pattern).

    Loads the model on first call and caches it for subsequent calls.
    Supports loading from:
    1. Local directory (if available)
    2. HuggingFace Hub (if enabled and local not available)

    Model is set to evaluation mode and runs on CPU.

    Returns:
        Loaded BERT model for sequence classification

    Raises:
        FileNotFoundError: If model not found locally and HF Hub disabled
        RuntimeError: If model loading fails

    Environment Variables:
        DENES_USE_HF_HUB: Enable HuggingFace Hub download (default: "true")
        DENES_MODEL_REPO: HuggingFace repository ID (e.g., "username/model-name")
        HUGGINGFACE_TOKEN: HF access token (optional, for private models)
    """
    global _model, _loaded

    if not _loaded:
        model_path = get_model_path()

        # Strategy 1: Try local model first (development mode)
        if model_path.exists() and (model_path / "model.safetensors").exists():
            logger.info(f"Loading model from local path: {model_path}")
            try:
                _model = AutoModelForSequenceClassification.from_pretrained(str(model_path))
                _model.eval()
                _model.to("cpu")
                _loaded = True
                logger.info("✓ Model loaded successfully from local path")
                return _model
            except Exception as e:
                logger.warning(f"Failed to load local model: {e}")
                # Continue to try HuggingFace Hub

        # Strategy 2: Download from HuggingFace Hub (production mode)
        if USE_HF_HUB:
            if "SEU_USERNAME" in HF_REPO_ID:
                raise RuntimeError(
                    "HuggingFace repository not configured! "
                    f"Update HF_REPO_ID in model_loader.py or set DENES_MODEL_REPO env var. "
                    f"Current value: {HF_REPO_ID}"
                )

            logger.info(f"Downloading model from HuggingFace Hub: {HF_REPO_ID}")
            logger.info("(This may take 1-2 minutes on first download, then cached)")

            try:
                _model = AutoModelForSequenceClassification.from_pretrained(
                    HF_REPO_ID,
                    token=HF_TOKEN,  # Uses token from HUGGINGFACE_TOKEN env var
                    cache_dir=str(model_path.parent / ".hf_cache"),
                )
                _model.eval()
                _model.to("cpu")
                _loaded = True
                logger.info("✓ Model loaded successfully from HuggingFace Hub")
                logger.info(f"   Cached at: {model_path.parent / '.hf_cache'}")
                return _model

            except Exception as e:
                raise RuntimeError(
                    f"Failed to download/load model from HuggingFace Hub: {e}\n\n"
                    "Common fixes:\n"
                    "  1. Set HUGGINGFACE_TOKEN in .env (recommended)\n"
                    "  2. Or login to HuggingFace: huggingface-cli login\n"
                    "  3. Check repository ID is correct: {HF_REPO_ID}\n"
                    "  4. Verify you have access to the repository (if private)\n"
                    "  5. Or set DENES_USE_HF_HUB=false and use local model only"
                ) from e

        # Strategy 3: Neither local nor HF Hub available
        raise FileNotFoundError(
            f"Model not found!\n\n"
            f"Tried:\n"
            f"  1. Local path: {model_path} (not found)\n"
            f"  2. HuggingFace Hub: {HF_REPO_ID} (disabled or failed)\n\n"
            f"Solutions:\n"
            f"  1. Train the model: cd .. && uv run python cli.py train\n"
            f"  2. Enable HF Hub: set DENES_USE_HF_HUB=true\n"
            f"  3. Configure HF repo: set DENES_MODEL_REPO=your-username/model-name"
        )

    return _model


def get_tokenizer() -> AutoTokenizer:
    """Get loaded tokenizer (singleton pattern).

    Loads the tokenizer on first call and caches it for subsequent calls.
    Supports loading from:
    1. Local directory (if available)
    2. HuggingFace Hub (if enabled and local not available)

    Returns:
        Loaded tokenizer for the BERT model

    Raises:
        FileNotFoundError: If tokenizer not found locally and HF Hub disabled
        RuntimeError: If tokenizer loading fails

    Environment Variables:
        DENES_USE_HF_HUB: Enable HuggingFace Hub download (default: "true")
        DENES_MODEL_REPO: HuggingFace repository ID (e.g., "username/model-name")
        HUGGINGFACE_TOKEN: HF access token (optional, for private models)
    """
    global _tokenizer

    if _tokenizer is None:
        model_path = get_model_path()

        # Strategy 1: Try local tokenizer first
        if model_path.exists() and (model_path / "tokenizer.json").exists():
            logger.info(f"Loading tokenizer from local path: {model_path}")
            try:
                _tokenizer = AutoTokenizer.from_pretrained(str(model_path), use_fast=True)
                logger.info("✓ Tokenizer loaded successfully from local path")
                return _tokenizer
            except Exception as e:
                logger.warning(f"Failed to load local tokenizer: {e}")

        # Strategy 2: Download from HuggingFace Hub
        if USE_HF_HUB:
            if "SEU_USERNAME" in HF_REPO_ID:
                raise RuntimeError(
                    "HuggingFace repository not configured! "
                    f"Update HF_REPO_ID in model_loader.py or set DENES_MODEL_REPO env var."
                )

            logger.info(f"Downloading tokenizer from HuggingFace Hub: {HF_REPO_ID}")

            try:
                _tokenizer = AutoTokenizer.from_pretrained(
                    HF_REPO_ID,
                    token=HF_TOKEN,  # Uses token from HUGGINGFACE_TOKEN env var
                    cache_dir=str(model_path.parent / ".hf_cache"),
                    use_fast=True,
                )
                logger.info("✓ Tokenizer loaded successfully from HuggingFace Hub")
                return _tokenizer

            except Exception as e:
                raise RuntimeError(
                    f"Failed to download/load tokenizer from HuggingFace Hub: {e}\n\n"
                    "Common fixes:\n"
                    "  1. Set HUGGINGFACE_TOKEN in .env (recommended)\n"
                    "  2. Or login to HuggingFace: huggingface-cli login\n"
                    "  3. Check repository ID: {HF_REPO_ID}\n"
                    "  4. Or set DENES_USE_HF_HUB=false to disable auto-download"
                ) from e

        # Neither local nor HF Hub available
        raise FileNotFoundError(
            f"Tokenizer not found at {model_path} and HF Hub disabled/failed"
        )

    return _tokenizer


def is_model_loaded() -> bool:
    """Check if model and tokenizer are loaded.

    Returns:
        True if both model and tokenizer are loaded, False otherwise
    """
    return _loaded and _tokenizer is not None


def unload_model() -> None:
    """Unload model and tokenizer from memory.

    Useful for testing or freeing memory. Not typically needed in production.
    """
    global _model, _tokenizer, _loaded

    if _loaded or _tokenizer is not None:
        logger.info("Unloading model and tokenizer...")
        _model = None
        _tokenizer = None
        _loaded = False
        # Force garbage collection
        import gc

        gc.collect()
        logger.info("Model and tokenizer unloaded")


def get_model_info() -> dict:
    """Get information about current model configuration.

    Returns:
        Dict with model configuration details
    """
    model_path = get_model_path()
    local_exists = model_path.exists() and (model_path / "model.safetensors").exists()

    return {
        "hf_repo_id": HF_REPO_ID,
        "use_hf_hub": USE_HF_HUB,
        "local_path": str(model_path),
        "local_exists": local_exists,
        "model_loaded": _loaded,
        "tokenizer_loaded": _tokenizer is not None,
    }
