from transformers import AutoTokenizer
from transformers import AutoModelForSequenceClassification

from app.config.settings import settings
from app.config.logging import logger
from transformers import AutoImageProcessor
from transformers import AutoModelForImageClassification

import torch


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 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):
    print("Loading tokenizer...")

    self.text_tokenizer = AutoTokenizer.from_pretrained(
        settings.TEXT_MODEL,
        cache_dir=settings.MODEL_CACHE_DIR
    )

    print("Tokenizer loaded.")

  def load_model(self):
    print("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()
    print(self.text_model.config.id2label)

    print("Text model loaded.")

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

  def load_image_model(self):

    print("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()


    print("Image model loaded.")


model_manager = ModelManager()