import os
import gc
import time
import torch
import config
from huggingface_hub import hf_hub_download
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from model_registry import PLAIN_MODELS, CPU_GGUF_MODELS, validate_model_id

# Global state for loaded model & tokenizer
active_model = None
active_tokenizer = None
active_model_id = None
model_load_time_ms = 0.0
loaded_model_hash = ""
loaded_model_revision = "main"

def get_context_limit(model_id: str) -> int:
    if model_id == "plain-llama-3-8b":
        return config.MODEL_CONTEXT_LIMIT_LLAMA_3_8B
    elif model_id == "plain-llama-3.2-3b":
        return config.MODEL_CONTEXT_LIMIT_LLAMA_3_2_3B
    elif model_id == "plain-qwen-0.5b":
        return config.MODEL_CONTEXT_LIMIT_QWEN_0_5B
    return 2048

def unload_all_models():
    global active_model, active_tokenizer, active_model_id, loaded_model_hash, loaded_model_revision
    if active_model is not None:
        # If it is a GGUF model, we may want to call its destructor if needed (llama_cpp cleans up automatically on delete)
        del active_model
        active_model = None
    if active_tokenizer is not None:
        del active_tokenizer
        active_tokenizer = None
    active_model_id = None
    loaded_model_hash = ""
    loaded_model_revision = "main"
    
    # Run garbage collection
    gc.collect()
    # Clear PyTorch CUDA Cache if GPU is active
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

def load_tokenizer_only(model_id: str):
    validate_model_id(model_id)
    repo_id = PLAIN_MODELS[model_id]
    
    # Always load Hugging Face tokenizer for chat template rendering and token counting
    tokenizer = AutoTokenizer.from_pretrained(
        repo_id,
        token=config.HF_TOKEN if config.HF_TOKEN else None,
        trust_remote_code=True
    )
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    return tokenizer

def load_model_and_tokenizer(model_id: str):
    global active_model, active_tokenizer, active_model_id, model_load_time_ms, loaded_model_hash, loaded_model_revision
    
    validate_model_id(model_id)
    
    if active_model_id == model_id:
        return active_model, active_tokenizer
        
    unload_all_models()
    
    start_time = time.time()
    
    # Load Hugging Face tokenizer first (used for both GPU & CPU profiles)
    tokenizer = load_tokenizer_only(model_id)
    
    ctx_limit = get_context_limit(model_id)
    
    if config.INFERENCE_PROFILE == "gpu-transformers":
        # Load Hugging Face Model with bitsandbytes 4-bit quantization
        repo_id = PLAIN_MODELS[model_id]
        
        bnb_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_use_double_quant=True,
            bnb_4bit_compute_dtype=torch.float16
        )
        
        model = AutoModelForCausalLM.from_pretrained(
            repo_id,
            quantization_config=bnb_config,
            device_map="auto",
            token=config.HF_TOKEN if config.HF_TOKEN else None,
            trust_remote_code=True
        )
        loaded_model_revision = "main"
        loaded_model_hash = "gpu-transformers-hash-4bit"
        
    elif config.INFERENCE_PROFILE == "cpu-gguf":
        # Load GGUF model via llama-cpp-python
        import llama_cpp
        
        gguf_cfg = CPU_GGUF_MODELS[model_id]
        repo_id = gguf_cfg["repo_id"]
        filename = gguf_cfg["filename"]
        
        # Download file using huggingface_hub
        model_path = hf_hub_download(
            repo_id=repo_id,
            filename=filename,
            token=config.HF_TOKEN if config.HF_TOKEN else None
        )
        
        # Initialize llama.cpp model
        model = llama_cpp.Llama(
            model_path=model_path,
            n_ctx=ctx_limit,
            n_threads=config.MODEL_THREADS,
            n_batch=config.MODEL_BATCH_SIZE,
            verbose=False
        )
        loaded_model_revision = "main"
        # Extract file size/hash metadata
        import hashlib
        h = hashlib.sha256()
        with open(model_path, "rb") as f:
            # Read first 1MB for hashing to avoid blocking
            h.update(f.read(1024 * 1024))
        loaded_model_hash = h.hexdigest()
        
    else:
        raise ValueError(f"Unknown INFERENCE_PROFILE: {config.INFERENCE_PROFILE}")
        
    active_model = model
    active_tokenizer = tokenizer
    active_model_id = model_id
    model_load_time_ms = (time.time() - start_time) * 1000.0
    
    return active_model, active_tokenizer