import json
import torch

from pathlib import Path
from typing import Dict, Any

from src.app.model_utils.factory import ModelFactory


class ModelManager:
    """
    Manages model loading and inference operations
    
    Args:
        model_dir: Directory containing model artifacts
    """
    
    def __init__(self, model_dir: str = "../pretrained") -> None:
        self.model_dir = Path(model_dir)
        self.loaded_models: Dict[str, Any] = {}
        self._load_model_artifacts()
        

    def _load_model_artifacts(self) -> None:
        """
        Load model configuration and vocabulary
        """

        with open(self.model_dir / "config.json", "r") as f:
            self.config = json.load(f)
        
        with open(self.model_dir / "vocab.json", "r") as f:
            self.vocab = json.load(f)
        
        self.idx_to_label = {0: "Negative", 1: "Positive"}
        

    def get_model(self) -> torch.nn.Module:
        """
        Get the loaded model (cached for performance)
        
        Returns:
            Loaded PyTorch model in evaluation mode
        """
        
        model_type = self.config["model_type"]
        
        if model_type not in self.loaded_models:
            model = ModelFactory.create_model(
                model_type=model_type,
                model_params=self.config["model_params"],
                state_dict_path=self.model_dir / "best_model.pth"
            )
            self.loaded_models[model_type] = model
        
        return self.loaded_models[model_type]
    

    def get_vocab(self) -> Dict[str, int]:
        """
        Get vocabulary mapping
        """

        return self.vocab
    

    def get_config(self) -> Dict[str, Any]:
        """
        Get model configuration
        """
        
        return self.config