import os
import torch
from transformers import AutoImageProcessor
from transformers import AutoModelForImageClassification
from transformers import AutoModelForSequenceClassification
from transformers import AutoTokenizer

from app.config.logging import logger
from app.config.settings import settings


class ModelManager:

    def __init__(self):
        self.device = None
        self.text_model = None
        self.text_tokenizer = None
        self.image_model = None
        self.image_processor = None

    def _setup_hf_token(self):
        if settings.HF_TOKEN:
            os.environ["HF_TOKEN"] = settings.HF_TOKEN
            os.environ["HUGGING_FACE_HUB_TOKEN"] = settings.HF_TOKEN

    def load_device(self):
        if settings.DEVICE == "auto":
            self.device = torch.device(
                "cuda" if torch.cuda.is_available() else "cpu"
            )
        else:
            self.device = torch.device(settings.DEVICE)

    def load_tokenizer(self):
        logger.info("Loading tokenizer...")
        self.text_tokenizer = AutoTokenizer.from_pretrained(
            settings.TEXT_MODEL,
            cache_dir=settings.MODEL_CACHE_DIR
        )
        logger.info("Tokenizer loaded.")

    def load_model(self):
        logger.info("Loading model...")
        self.text_model = AutoModelForSequenceClassification.from_pretrained(
            settings.TEXT_MODEL,
            cache_dir=settings.MODEL_CACHE_DIR
        )
        self.text_model.to(self.device)
        self.text_model.eval()
        logger.debug("id2label: %s", self.text_model.config.id2label)
        logger.info("Text model loaded.")

    def load_image_model(self):
        logger.info("Loading image model...")
        self.image_processor = AutoImageProcessor.from_pretrained(
            settings.IMAGE_MODEL,
            cache_dir=settings.MODEL_CACHE_DIR
        )
        self.image_model = AutoModelForImageClassification.from_pretrained(
            settings.IMAGE_MODEL,
            cache_dir=settings.MODEL_CACHE_DIR
        )
        self.image_model.to(self.device)
        self.image_model.eval()
        logger.info("Image model loaded.")

    def initialize(self):
        self.load_device()
        logger.info(f"Using device: {self.device}")
        self.load_tokenizer()
        self.load_model()
        self.load_image_model()


model_manager = ModelManager()