import torch
import time
import os
from unsloth import FastVisionModel
from .config import MODEL_PATH, ADAPTERS_BASE

class ModelLoader:
    def __init__(self):
        self.model = None
        self.processor = None
        self.active_persona = None
        self.active_adapter = None
        self.load_start_time = None
        self.chat_template = None

    def load_base_model(self):
        print(f"Loading base model from {MODEL_PATH} using Unsloth...")
        self.load_start_time = time.time()
        try:
            # Gemma 4 needs FastVisionModel for correct architecture patching
            self.model, self.processor = FastVisionModel.from_pretrained(
                MODEL_PATH,
                load_in_4bit=True, # Recommended for H100 efficiency
                use_gradient_checkpointing="unsloth",
            )
            
            # Load chat template if exists
            chat_template_path = os.path.join(MODEL_PATH, "chat_template.jinja")
            if os.path.exists(chat_template_path):
                with open(chat_template_path, 'r') as f:
                    self.chat_template = f.read()
            
            print(f"Base model loaded in {time.time() - self.load_start_time:.2f}s")
            print(f"GPU memory allocated: {self.get_gpu_memory_gb():.2f} GB")
        except Exception as e:
            print(f"Error loading model with Unsloth: {e}")
            self.model = None
            raise

    def load_persona_adapter(self, persona_id):
        if persona_id == self.active_persona:
            return

        if not self.model:
            raise RuntimeError("Base model not loaded")

        # Instructions say: /artifacts/stage2/ananya/persona-ananya.lora/
        adapter_path = os.path.join(ADAPTERS_BASE, persona_id, f"persona-{persona_id}.lora")
        
        if not os.path.exists(adapter_path):
            print(f"Warning: Adapter path {adapter_path} not found. Falling back to base model for persona: {persona_id}")
            self.unload_adapter()
            self.active_persona = persona_id
            return

        print(f"Switching to persona adapter: {persona_id} from {adapter_path}")
        start_time = time.time()
        
        try:
            # Unsloth handle PEFT models differently, but for inference:
            # We can use standard PEFT methods or Unsloth's if available
            from peft import PeftModel
            
            # If already a PeftModel, unload first
            if isinstance(self.model, PeftModel):
                self.model = self.model.unload()

            self.model = PeftModel.from_pretrained(
                self.model,
                adapter_path,
                adapter_name=persona_id
            )
            self.active_persona = persona_id
            print(f"Adapter {persona_id} loaded in {time.time() - start_time:.2f}s")
        except Exception as e:
            print(f"Error loading adapter: {e}")
            raise

    def unload_adapter(self):
        from peft import PeftModel
        if isinstance(self.model, PeftModel):
            self.model = self.model.unload()
            self.active_persona = None

    def get_gpu_memory_gb(self):
        if torch.cuda.is_available():
            return torch.cuda.memory_allocated() / 1e9
        return 0.0

    def is_ready(self):
        return self.model is not None

# Singleton instance
model_loader = ModelLoader()