import torch
import time
import os
from transformers import AutoModelForCausalLM, AutoProcessor
from peft import PeftModel
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}...")
        self.load_start_time = time.time()
        try:
            self.model = AutoModelForCausalLM.from_pretrained(
                MODEL_PATH,
                device_map="auto",
                torch_dtype=torch.bfloat16,
            )
            self.processor = AutoProcessor.from_pretrained(MODEL_PATH)
            self.model.eval()
            
            # 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: {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")

        adapter_path = os.path.join(ADAPTERS_BASE, persona_id, f"persona-{persona_id}.lora")
        
        # In some cases the directory structure might be different based on the artifacts
        # Checking if persona-persona_id.lora directory exists, or if it's just the adapter safely
        if not os.path.exists(adapter_path):
            # Fallback for common structure: artifacts/stage2/persona_id/adapter_model.safetensors
            # Wait, the instruction says: /artifacts/stage2/ananya/persona-ananya.lora/adapter_model.safetensors
            adapter_path = os.path.join(ADAPTERS_BASE, persona_id, f"persona-{persona_id}.lora")

        print(f"Switching to persona adapter: {persona_id} from {adapter_path}")
        start_time = time.time()
        
        try:
            # If an adapter is already loaded, we need to handle it.
            # PEFT models allow multiple adapters, but for simplicity we will unload/reload or just switch.
            if self.active_persona:
                # Merge and unload current adapter before loading new one to keep memory clean
                # or just use set_adapter if already loaded.
                # However, the instructions imply loading from disk.
                # PeftModel.from_pretrained on an already wrapped model adds a new adapter.
                if isinstance(self.model, PeftModel):
                    self.model = self.model.unload() # Unload back to base model

            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):
        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()