from typing import Any, List, Optional
from veridex.core.signal import BaseSignal, DetectionResult

class CLIPSignal(BaseSignal):
    """
    Detects AI images using CLIP (Contrastive Language-Image Pre-Training) Zero-Shot Classification.

    Based on the research "Towards Universal Fake Image Detectors that Generalize Across Generative Models" (UnivFD)
    and "Raising the Bar of AI-generated Image Detection with CLIP".

    Methodology:
        1. Encodes the input image using a pre-trained CLIP vision encoder.
        2. Encodes a set of text prompts representing "Real" and "Fake" classes.
        3. Computes cosine similarity between the image embedding and text embeddings.
        4. Returns the probability of the "Fake" class (softmax of similarities).

    Attributes:
        name (str): 'clip_zeroshot'
        dtype (str): 'image'
        model_id (str): HuggingFace CLIP model ID (default: "openai/clip-vit-base-patch32").
    """

    def __init__(
        self,
        model_id: str = "openai/clip-vit-base-patch32",
        device: str = "cpu",
        real_prompts: Optional[List[str]] = None,
        fake_prompts: Optional[List[str]] = None
    ):
        """
        Initialize the CLIP signal.

        Args:
            model_id (str): The CLIP model to use.
            device (str): Computation device ('cpu' or 'cuda').
            real_prompts (List[str], optional): Custom prompts for the 'Real' class.
            fake_prompts (List[str], optional): Custom prompts for the 'Fake' class.
        """
        self.model_id = model_id
        self.device = device
        self._model = None
        self._processor = None

        # Default prompts if not provided
        self.real_prompts = real_prompts or [
            "a photo",
            "a real photo",
            "a photograph",
            "a photo of a real object",
            "an authentic image"
        ]
        self.fake_prompts = fake_prompts or [
            "an ai generated image",
            "a synthetic image",
            "a deepfake",
            "an image generated by stable diffusion",
            "an image generated by dall-e",
            "an image generated by midjourney",
            "artificial intelligence art"
        ]

    @property
    def name(self) -> str:
        return "clip_zeroshot"

    @property
    def dtype(self) -> str:
        return "image"

    def check_dependencies(self) -> None:
        try:
            import torch
            import transformers
        except ImportError as e:
            raise ImportError(
                "CLIPSignal requires 'torch' and 'transformers'. "
                "Install with `pip install veridex[image]` (if available) or install them directly."
            ) from e

    def _load_model(self):
        if self._model is not None and self._processor is not None:
            return self._model, self._processor

        self.check_dependencies()
        from transformers import CLIPProcessor, CLIPModel

        try:
            self._processor = CLIPProcessor.from_pretrained(self.model_id)
            self._model = CLIPModel.from_pretrained(self.model_id).to(self.device)
            self._model.eval()
        except Exception as e:
            raise RuntimeError(f"Failed to load CLIP model '{self.model_id}': {e}")

        return self._model, self._processor

    def run(self, input_data: Any) -> DetectionResult:
        try:
            from PIL import Image
            import torch
            import numpy as np
        except ImportError:
            self.check_dependencies()
            raise

        # 1. Prepare Input Image
        image = None
        if isinstance(input_data, str):
            try:
                image = Image.open(input_data).convert("RGB")
            except Exception as e:
                return DetectionResult(
                    score=0.0, confidence=0.0, metadata={},
                    error=f"Could not open image: {e}"
                )
        elif isinstance(input_data, Image.Image):
            image = input_data.convert("RGB")
        elif isinstance(input_data, np.ndarray):
             image = Image.fromarray(input_data).convert("RGB")
        else:
             return DetectionResult(
                score=0.0, confidence=0.0, metadata={},
                error="Input must be file path, PIL Image, or numpy array."
            )

        try:
            model, processor = self._load_model()

            # 2. Prepare Prompts
            all_prompts = self.real_prompts + self.fake_prompts

            # 3. Process Inputs
            inputs = processor(
                text=all_prompts,
                images=image,
                return_tensors="pt",
                padding=True
            ).to(self.device)

            # 4. Inference
            with torch.no_grad():
                outputs = model(**inputs)
                logits_per_image = outputs.logits_per_image  # image-text similarity score
                probs = logits_per_image.softmax(dim=1) # Shape: (1, num_prompts)

            # 5. Aggregate Scores
            # Sum probabilities of all "fake" prompts vs "real" prompts
            # probs[0] contains probabilities for [real_prompts..., fake_prompts...]

            num_real = len(self.real_prompts)

            prob_real = probs[0, :num_real].sum().item()
            prob_fake = probs[0, num_real:].sum().item()

            # Normalize just in case they don't sum to exactly 1.0 (though softmax ensures they do)
            total_prob = prob_real + prob_fake
            final_fake_score = prob_fake / total_prob if total_prob > 0 else 0.0

            # Confidence could be the max probability of the winning class or the margin
            # Here we treat the final score itself as the probability of being AI.
            # Confidence metric: How far is the score from 0.5 (uncertainty)?
            # Map 0.5 -> 0.0 confidence, 1.0/0.0 -> 1.0 confidence
            confidence = abs(final_fake_score - 0.5) * 2

            return DetectionResult(
                score=final_fake_score,
                confidence=confidence,
                metadata={
                    "model": self.model_id,
                    "prob_real": prob_real,
                    "prob_fake": prob_fake,
                    "top_prompt": all_prompts[probs[0].argmax().item()]
                }
            )

        except Exception as e:
            return DetectionResult(
                score=0.0,
                confidence=0.0,
                metadata={},
                error=f"CLIP execution failed: {e}"
            )
