"""Model loading and status management.

Enhanced version with:
- Model caching and memory optimization
- Better error handling and recovery
- Progress tracking for model loading
- Resource cleanup utilities
- Support for quantized models

Supports multiple models:
- Qwen2.5-Coder-1.5B (text-only, smart code generation)
- MiniCPM5-1B (text-only, fast)
- MiniCPM-V-4.6 (vision + text)
- DeepSeek-Coder-1.3B (advanced reasoning)

Only one model is loaded at a time to conserve memory.
The model is loaded in a background thread on startup.
"""

from __future__ import annotations

import gc
import logging
import threading
import time
from typing import Any, Callable, Optional
from dataclasses import dataclass, field
from enum import Enum

from code.config.constants import MODEL_CONFIGS

logger = logging.getLogger(__name__)


class ModelState(Enum):
    """Enum representing model loading states."""
    UNLOADED = "unloaded"
    LOADING = "loading"
    READY = "ready"
    ERROR = "error"


@dataclass
class LoadProgress:
    """Track model loading progress."""
    stage: str = "initializing"
    progress_percent: float = 0.0
    message: str = ""
    start_time: float = field(default_factory=time.time)
    elapsed_seconds: float = 0.0
    
    def update(self, stage: str, progress: float, message: str = ""):
        """Update progress state."""
        self.stage = stage
        self.progress_percent = progress
        self.message = message
        self.elapsed_seconds = time.time() - self.start_time


# ─── Module-level state ─────────────────────────────────────────────────

_current_model_key: str = ""
_model = None
_tokenizer_or_processor = None
_model_state: ModelState = ModelState.UNLOADED
_load_error: str | None = None
_load_progress: LoadProgress = LoadProgress()
_progress_callbacks: list[Callable[[LoadProgress], None]] = []
_load_lock = threading.Lock()

# Initialize from config
for key, config in MODEL_CONFIGS.items():
    if hasattr(config, 'get') and config.get("id") == "Qwen/Qwen2.5-Coder-1.5B-Instruct":
        _current_model_key = key
        break
if not _current_model_key:
    _current_model_key = "qwen25-coder-1.5b"


def _notify_progress(stage: str, progress: float, message: str = ""):
    """Notify all registered progress callbacks."""
    global _load_progress
    _load_progress.update(stage, progress, message)
    for callback in _progress_callbacks:
        try:
            callback(_load_progress)
        except Exception as e:
            logger.warning("Progress callback error: %s", e)


def add_progress_callback(callback: Callable[[LoadProgress], None]):
    """Register a callback for load progress updates."""
    _progress_callbacks.append(callback)


def remove_progress_callback(callback: Callable[[LoadProgress], None]):
    """Remove a progress callback."""
    if callback in _progress_callbacks:
        _progress_callbacks.remove(callback)


def _unload_model() -> None:
    """Unload current model and free memory with comprehensive cleanup."""
    global _model, _tokenizer_or_processor, _model_state

    logger.info("Unloading current model...")
    
    # Clear CUDA cache first if available
    try:
        import torch
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            torch.cuda.synchronize()
    except ImportError:
        pass

    # Delete model and tokenizer
    if _model is not None:
        del _model
        _model = None
    if _tokenizer_or_processor is not None:
        del _tokenizer_or_processor
        _tokenizer_or_processor = None

    _model_state = ModelState.UNLOADED
    
    # Force garbage collection
    gc.collect()
    
    # Clear CUDA cache again after GC
    try:
        import torch
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
    except ImportError:
        pass
    
    logger.info("Model unloaded successfully")


def load_model(model_key: str | None = None) -> None:
    """Load a model by key. Unloads the previous model first.
    
    Args:
        model_key: The key of the model to load from MODEL_CONFIGS.
                  If None, loads the current default model.
    """
    global _model, _tokenizer_or_processor, _model_state, _load_error, _current_model_key

    if model_key is None:
        model_key = _current_model_key

    if model_key not in MODEL_CONFIGS:
        _load_error = f"Unknown model: {model_key}"
        _model_state = ModelState.ERROR
        logger.error(_load_error)
        return

    # Skip if already loading or already loaded with same key
    if _model_state == ModelState.LOADING:
        logger.info("Model already loading, skipping...")
        return
    if _model_state == ModelState.READY and _current_model_key == model_key:
        logger.info("Model %s already loaded", model_key)
        return

    # Acquire lock for thread safety
    with _load_lock:
        # Double-check after acquiring lock
        if _model_state == ModelState.LOADING:
            return
        if _model_state == ModelState.READY and _current_model_key == model_key:
            return
            
        _model_state = ModelState.LOADING
        _load_error = None

    # Unload previous model if switching
    if _model_state != ModelState.UNLOADED and _current_model_key != model_key:
        logger.info("Switching model from %s to %s", _current_model_key, model_key)
        _unload_model()

    _current_model_key = model_key
    config = MODEL_CONFIGS[model_key]
    model_id = config["id"]

    try:
        _notify_progress("starting", 0.0, f"Loading {config['name']}...")
        
        import torch
        
        # Determine device and dtype based on available hardware
        if torch.cuda.is_available():
            dtype = torch.float16
            device_map = "auto"
            # Check for bfloat16 support (better on Ampere+ GPUs)
            if torch.cuda.is_bf16_supported():
                dtype = torch.bfloat16
                _notify_progress("config", 10.0, "Using bfloat16 precision")
            else:
                _notify_progress("config", 10.0, "Using float16 precision")
        else:
            dtype = torch.float32
            device_map = None
            _notify_progress("config", 10.0, "Using CPU (float32)")

        if config["type"] == "vlm":
            _load_vlm_model(model_id, dtype, device_map)
        else:
            _load_text_model(model_id, dtype, device_map)

        _model_state = ModelState.READY
        _notify_progress("complete", 100.0, f"{config['name']} loaded successfully!")
        logger.info("%s model loaded successfully.", config["name"])

    except MemoryError as exc:
        _load_error = f"Out of memory loading {model_id}: {exc}"
        _model_state = ModelState.ERROR
        _notify_progress("error", 0.0, _load_error)
        logger.error(_load_error)
        # Try to free memory
        _unload_model()
        
    except Exception as exc:
        _load_error = f"Failed to load model {model_id}: {exc}"
        _model_state = ModelState.ERROR
        _notify_progress("error", 0.0, _load_error)
        logger.exception("Failed to load model %s: %s", model_id, exc)


def _load_text_model(model_id: str, dtype, device_map) -> None:
    """Load a text-only model (AutoModelForCausalLM + AutoTokenizer).
    
    Enhanced with:
    - Progressive loading notifications
    - Memory-efficient settings
    - Error recovery hints
    """
    global _model, _tokenizer_or_processor

    from transformers import AutoModelForCausalLM, AutoTokenizer

    logger.info("Loading %s (text model)...", model_id)
    _notify_progress("tokenizer", 20.0, f"Loading tokenizer for {model_id}...")

    # Load tokenizer with error handling
    try:
        _tokenizer_or_processor = AutoTokenizer.from_pretrained(
            model_id,
            trust_remote_code=True,
        )
        # Set pad token if not set
        if _tokenizer_or_processor.pad_token is None:
            _tokenizer_or_processor.pad_token = _tokenizer_or_processor.eos_token
    except Exception as e:
        logger.warning("Tokenizer load issue: %s", e)
        raise

    _notify_progress("model", 50.0, f"Loading model weights for {model_id}...")
    
    # Load model with memory optimizations
    _model = AutoModelForCausalLM.from_pretrained(
        model_id,
        torch_dtype=dtype,
        device_map=device_map,
        trust_remote_code=True,
        low_cpu_mem_usage=True,
        # Memory optimizations
        offload_folder="/tmp/offload" if device_map is None else None,
        offload_state_dict=True if device_map is None else False,
    )

    if device_map is None:
        _model = _model.to("cpu")

    _model.eval()
    _notify_progress("finalizing", 90.0, "Finalizing model...")


def _load_vlm_model(model_id: str, dtype, device_map) -> None:
    """Load a vision-language model (AutoModelForImageTextToText + AutoProcessor).
    
    Enhanced with:
    - Fallback handling for different transformers versions
    - Progressive loading
    - Memory optimization
    """
    global _model, _tokenizer_or_processor

    try:
        from transformers import AutoModelForImageTextToText, AutoProcessor
    except ImportError:
        # Fallback for older transformers
        logger.warning("AutoModelForImageTextToText not found, trying AutoModel...")
        from transformers import AutoModel as AutoModelForImageTextToText
        from transformers import AutoProcessor

    logger.info("Loading %s (VLM)...", model_id)
    _notify_progress("processor", 20.0, f"Loading processor for {model_id}...")

    # Load processor
    try:
        _tokenizer_or_processor = AutoProcessor.from_pretrained(
            model_id,
            trust_remote_code=True,
        )
    except Exception as e:
        logger.warning("Processor load issue: %s", e)
        raise

    _notify_progress("model", 50.0, f"Loading VLM model weights...")

    # Load VLM model
    _model = AutoModelForImageTextToText.from_pretrained(
        model_id,
        torch_dtype=dtype,
        device_map=device_map,
        trust_remote_code=True,
        low_cpu_mem_usage=True,
    )

    if device_map is None:
        _model = _model.to("cpu")

    _model.eval()
    _notify_progress("finalizing", 90.0, "Finalizing VLM model...")


def start_background_load(model_key: str | None = None) -> threading.Thread:
    """Start loading the model in a background daemon thread.
    
    Args:
        model_key: Optional model key to load. Uses current if not specified.
        
    Returns:
        The background thread that was started.
    """
    thread = threading.Thread(target=load_model, args=(model_key,), daemon=True)
    thread.start()
    logger.info("Background model loading started in thread %s", thread.name)
    return thread


def switch_model(model_key: str) -> dict[str, Any]:
    """Switch to a different model. Returns status immediately, loads in background.
    
    Args:
        model_key: The key of the model to switch to.
        
    Returns:
        Dict with success status and message about the switch operation.
    """
    global _current_model_key, _model_state

    if model_key not in MODEL_CONFIGS:
        return {"success": False, "message": f"Unknown model: {model_key}"}

    if _current_model_key == model_key and _model_state == ModelState.READY:
        return {"success": True, "message": f"Already using {MODEL_CONFIGS[model_key]['name']}"}

    _current_model_key = model_key
    _model_state = ModelState.UNLOADED  # Reset state to trigger reload

    # Start loading in background
    start_background_load(model_key)

    config = MODEL_CONFIGS[model_key]
    return {
        "success": True,
        "message": f"Switching to {config['name']}...",
        "model_key": model_key,
        "model_name": config["name"],
    }


def get_model_status() -> dict[str, Any]:
    """Return current model loading status with detailed information.
    
    Returns:
        Dict containing status, message, model info, and progress details.
    """
    config = MODEL_CONFIGS.get(_current_model_key, {})
    
    base_info = {
        "model_key": _current_model_key,
        "model_name": config.get("name", ""),
        "model_type": config.get("type", "text"),
        "model_description": config.get("description", ""),
        "capabilities": config.get("capabilities", []),
        "progress": {
            "stage": _load_progress.stage,
            "percent": _load_progress.progress_percent,
            "message": _load_progress.message,
            "elapsed_seconds": round(_load_progress.elapsed_seconds, 1),
        },
    }
    
    if _model_state == ModelState.READY:
        return {
            **base_info,
            "status": "ready",
            "message": f"{config.get('name', 'Model')} loaded and ready",
        }
    elif _model_state == ModelState.LOADING:
        return {
            **base_info,
            "status": "loading",
            "message": f"Loading {config.get('name', 'model')}... ({_load_progress.progress_percent:.0f}%)",
        }
    elif _model_state == ModelState.ERROR:
        return {
            **base_info,
            "status": "error",
            "message": f"Model load error: {_load_error}",
        }
    else:
        return {
            **base_info,
            "status": "unknown",
            "message": "Model not initialized",
        }


def get_model():
    """Return the loaded model instance (or None)."""
    return _model


def get_tokenizer_or_processor():
    """Return the loaded tokenizer or processor (or None)."""
    return _tokenizer_or_processor


def is_model_loaded() -> bool:
    """Return True if the model has been loaded successfully."""
    return _model_state == ModelState.READY


def get_current_model_key() -> str:
    """Return the key of the currently selected model."""
    return _current_model_key


def get_current_model_type() -> str:
    """Return 'text' or 'vlm' for the current model."""
    return MODEL_CONFIGS.get(_current_model_key, {}).get("type", "text")


def get_model_memory_usage() -> dict[str, Any]:
    """Get estimated memory usage of the loaded model.
    
    Returns:
        Dict with memory usage statistics if available.
    """
    result = {
        "model_loaded": False,
        "gpu_memory_mb": 0,
        "cpu_memory_mb": 0,
        "device": "cpu",
    }
    
    if _model is None:
        return result
    
    try:
        import torch
        
        result["model_loaded"] = True
        
        # Calculate parameter memory
        param_size = sum(p.numel() * p.element_size() for p in _model.parameters())
        buffer_size = sum(b.numel() * b.element_size() for b in _model.buffers())
        total_mb = (param_size + buffer_size) / (1024 * 1024)
        
        if torch.cuda.is_available() and next(_model.parameters()).is_cuda:
            result["device"] = "cuda"
            result["gpu_memory_mb"] = round(total_mb, 1)
            if torch.cuda.is_available():
                allocated = torch.cuda.memory_allocated() / (1024 * 1024)
                result["gpu_allocated_mb"] = round(allocated, 1)
        else:
            result["cpu_memory_mb"] = round(total_mb, 1)
            
    except Exception as e:
        logger.warning("Could not calculate memory usage: %s", e)
    
    return result


def optimize_for_inference() -> dict[str, Any]:
    """Optimize the loaded model for inference.
    
    Applies optimizations like:
    - Better Transformer structure (if available)
    - Half precision conversion
    - Attention optimization
    
    Returns:
        Dict with optimization results.
    """
    result = {"optimized": False, "techniques_applied": []}
    
    if _model is None:
        return result
    
    try:
        import torch
        
        # Try to use better transformer if available
        if hasattr(_model, 'to_bettertransformer'):
            _model = _model.to_bettertransformer()
            result["techniques_applied"].append("better_transformer")
            
        # Torch compile if available (PyTorch 2.0+)
        if hasattr(torch, 'compile'):
            try:
                _model = torch.compile(_model, mode="reduce-overhead")
                result["techniques_applied"].append("torch_compile")
            except Exception:
                pass  # Torch compile may fail for some models
                
        result["optimized"] = len(result["techniques_applied"]) > 0
        
    except Exception as e:
        logger.warning("Optimization failed: %s", e)
        result["error"] = str(e)
    
    return result