from PIL import Image

from app.config.image_policy import FORBIDDEN_IMAGE_LABELS
from app.config.settings import settings
from app.ml.image.clip_classifier import ClipClassifier
from app.ml.image.image_classifier import ImageClassifier
from app.ml.image.image_prediction_result import ImagePredictionResult
from app.ml.image.weapon_detector import WeaponDetector
from app.moderation.image.validator import ImageValidator

class ImageModerationService:

  def __init__(
      self,
      classifier: ImageClassifier,
      validator: ImageValidator,
      clip_classifier: ClipClassifier,
      weapon_detector: WeaponDetector
  ):

    self.classifier = classifier
    self.clip_classifier = clip_classifier
    self.weapon_detector = weapon_detector
    self.validator = validator

  def moderate(
      self,
      image: Image.Image
  ) -> ImagePredictionResult:

    self.validator.validate(image)

    # =========================
    # 1. NSFW MODEL
    # =========================

    nsfw_prediction = self.classifier.predict(image)

    if (
        nsfw_prediction.label.lower() == "nsfw"
        and nsfw_prediction.score >= settings.NSFW_THRESHOLD
    ):
      nsfw_prediction.approved = False

      nsfw_prediction.reason = "NSFW"

      nsfw_prediction.detected_labels = [
        nsfw_prediction.label
      ]

      return nsfw_prediction

    # =========================
    # 2. WEAPON DETECTOR
    # =========================

    weapon_prediction = (
      self.weapon_detector.predict(image)
    )

    if weapon_prediction.detected_labels:
      return ImagePredictionResult(

          label=weapon_prediction.label,

          score=weapon_prediction.score,

          approved=False,

          reason="FORBIDDEN_CONTENT",

          detected_labels=weapon_prediction.detected_labels

      )

    # =========================
    # 3. CLIP MODEL
    # =========================

    clip_prediction = (
      self.clip_classifier.predict(image)
    )

    detected_forbidden = [

      label

      for label in clip_prediction.detected_labels

      if label.lower()
         in FORBIDDEN_IMAGE_LABELS

    ]

    if detected_forbidden:
      return ImagePredictionResult(

          label=clip_prediction.label,

          score=clip_prediction.score,

          approved=False,

          reason="FORBIDDEN_CONTENT",

          detected_labels=detected_forbidden

      )

    # =========================
    # 4. NORMAL IMAGE
    # =========================

    return ImagePredictionResult(

        label="normal",

        score=1.0,

        approved=True,

        reason="OK",

        detected_labels=[]

    )