from typing import Dict, Any
import time

from huggingface_hub import hf_hub_download

from src.model import AIImageDetector
from src.models.model import Analysis, Attribute, ImageClassifierResult, Metadata
from src.helpers.image_utility_helper import load_image_resource


class AIImageClassifier2:
    def __init__(self):
        # Download model once during initialization
        self.model_identifier = "Bombek1/ai-image-detector-siglip-dinov2"
        self.model_path = hf_hub_download(
            repo_id=self.model_identifier,
            filename="pytorch_model.pt"
        )

        print(f"Model loaded from: {self.model_path}")

        # Initialize detector
        self.detector = AIImageDetector(self.model_path)

    # =====================================
    # Detect Function
    # =====================================
    def detect(self, image_path: str, image_type: str = None) -> Dict[str, Any]:
        start_time = time.time()

        # Load image using the separate function
        pil_image, raw_bytes = load_image_resource(image_path, image_type)
        if pil_image is None:
            raise ValueError("Failed to load image from input")

        # Predict
        result = self.detector.predict(pil_image)

        # Extract values
        is_ai = result["prediction"].lower().startswith("ai")
        ai_probability = float(result["probability"])
        ai_confidence = ai_probability  # Always reflect P(AI)

        end_time = time.time()
        processing_time = end_time - start_time

        # Format to Pydantic model
        formatted_result = ImageClassifierResult(
            analysis=Analysis(
                is_ai=is_ai,
                ai_confidence=ai_confidence
            ),
            attributes=[
                Attribute(
                    type="image_classifier_2",
                    weight=1.0,
                    ai_confidence=ai_confidence,
                    is_ai=is_ai,
                    parameters=[
                        {
                            "model": self.model_identifier,
                            "prediction": result["prediction"],
                            "confidence": float(result["confidence"]),
                            "probability_ai": ai_probability
                        }
                    ]
                )
            ],
            metadata=Metadata(start_time=start_time, end_time=end_time, processing_time=processing_time)
        )

        return formatted_result.model_dump()


if __name__ == "__main__":
    
    classifier = AIImageClassifier2()

    response = classifier.detect("/path/to/image.jpg")
    print(response)