from dotenv import load_dotenv
from transformers import pipeline
import os

load_dotenv()

class Config:
    EMOTION_MODEL_NAME = os.getenv("EMOTION_MODEL_NAME", "j-hartmann/emotion-english-distilroberta-base")
    SENTIMENT_MODEL_NAME = os.getenv("SENTIMENT_MODEL_NAME", "distilbert-base-uncased-finetuned-sst-2-english")
    DEVICE = int(os.getenv("MODEL_DEVICE", "-1"))

class ModelRegistry:
    def __init__(self):
        self.models = {
            "emotion": None,
            "sentiment": None
        }

    def load_models(self):
        print("Loading pretrained Hugging Face models...")
        self.models["emotion"] = pipeline("text-classification", model=Config.EMOTION_MODEL_NAME, device=Config.DEVICE)
        self.models["sentiment"] = pipeline("sentiment-analysis", model=Config.SENTIMENT_MODEL_NAME, device=Config.DEVICE)
        print("Models loaded successfully and ready for inference.")

    def initialize(self):
        if not all(self.models.values()):
            self.load_models()

    def get(self, name: str):
        return self.models.get(name)

model_registry = ModelRegistry()