import torch
from PIL import Image
from transformers import AutoProcessor, AutoModelForZeroShotObjectDetection
from api.config import settings

class GroundingDinoService:
    def __init__(self):
        self.model_id = settings.MODEL_ID
        print(f"Initializing Grounding DINO model '{self.model_id}'...")
        
        # Load processor and model
        self.processor = AutoProcessor.from_pretrained(self.model_id)
        self.model = AutoModelForZeroShotObjectDetection.from_pretrained(
            self.model_id,
            device_map="auto"
        )
        print("Grounding DINO model loaded successfully!")

    def detect(self, image: Image.Image) -> list[dict]:
        """
        Executes object detection on the PIL Image using the configured labels and thresholds.
        Returns a list of detections containing label, confidence score, and bounding boxes.
        """
        # Format labels as expected by the Grounding DINO processor:
        # A single lowercase string containing all labels separated by periods and ending with a period.
        text_prompt = ". ".join(settings.TEXT_LABELS).lower().strip()
        if not text_prompt.endswith("."):
            text_prompt += "."
        
        # Prepare inputs
        inputs = self.processor(
            images=image,
            text=text_prompt,
            return_tensors="pt"
        ).to(self.model.device)

        # Run inference
        with torch.no_grad():
            outputs = self.model(**inputs)

        # Post-process detections
        results = self.processor.post_process_grounded_object_detection(
            outputs,
            inputs.input_ids,
            threshold=settings.BOX_THRESHOLD,
            text_threshold=settings.TEXT_THRESHOLD,
            target_sizes=[image.size[::-1]]
        )

        detections = []
        result = results[0]

        # Iterate over output boxes, scores, and labels
        for box, score, label in zip(
            result["boxes"],
            result["scores"],
            result["labels"]
        ):
            confidence = round(score.item(), 3)
            
            # Filter detections by minimum confidence threshold
            if confidence < settings.MIN_CONFIDENCE:
                continue

            # Convert box coordinates to float list and round
            box_coords = [round(x, 2) for x in box.tolist()]

            detections.append({
                "label": label,
                "confidence": confidence,
                "box": box_coords
            })

        return detections