"""
llm_guard/visual_caption_bridge.py — Visual Caption Bridge
===========================================================
Bridges text-only scoring with multi-modal agents by:
  1. Generating a text caption from an image using BLIP-2 (optional)
  2. Computing semantic grounding score between agent text and image description
     using cosine similarity of sentence embeddings

The caption bridge preserves compatibility with all existing SC signals:
existing MechanismFeatureExtractor, AgentGuard, and score_chain() work
unchanged — the bridge adds a new feature: visual_grounding score.

Two modes:
  image_path        : load image file → BLIP-2 caption → embed → cosine sim
  image_description : use text description directly (no BLIP-2 required)

Usage
-----
    from llm_guard.features.visual_caption_bridge import VisualCaptionBridge

    bridge = VisualCaptionBridge()

    # With real image (requires transformers + PIL):
    score = bridge.score_grounding("The answer is 42.", image_path="diagram.png")

    # With text description (no image required):
    score = bridge.score_grounding("The answer is 42.",
                                   image_description="A bar chart showing values.")
Purpose:    Generates text captions from images using BLIP-2 (optional) and computes visual grounding scores using cosine similarity of sentence embeddings. Adds a visual_grounding feature compatible with all existing SC signals.
Exports:    VisualCaptionBridge
Added:      v0.58.0
Status:     Experimental
"""
from __future__ import annotations

from typing import Optional
import numpy as np


class VisualCaptionBridge:
    """
    Visual grounding scorer using BLIP-2 captioning + sentence embedding similarity.

    Parameters
    ----------
    embedding_model : str
        Sentence-transformers model for text embedding. Default: all-MiniLM-L6-v2.
    blip2_model : str
        BLIP-2 model for image captioning. Default: Salesforce/blip2-opt-2.7b.
    device : str
        Device for BLIP-2 inference. Default: "cpu".
    """

    def __init__(
        self,
        embedding_model: str = "all-MiniLM-L6-v2",
        blip2_model: str = "Salesforce/blip2-opt-2.7b",
        device: str = "cpu",
    ) -> None:
        self.embedding_model = embedding_model
        self.blip2_model = blip2_model
        self.device = device
        self._sent_model = None
        self._blip2_processor = None
        self._blip2_model = None

    # ── Embedding ─────────────────────────────────────────────────────────────

    def _get_sent_model(self):
        """Lazy-load sentence-transformer model."""
        if self._sent_model is None:
            from sentence_transformers import SentenceTransformer
            self._sent_model = SentenceTransformer(self.embedding_model)
        return self._sent_model

    def _embed_text(self, text: str) -> np.ndarray:
        """Embed text using sentence-transformers. Returns unit-normalised vector."""
        model = self._get_sent_model()
        emb = model.encode([text], normalize_embeddings=True, show_progress_bar=False)
        return emb[0]

    # ── BLIP-2 captioning ─────────────────────────────────────────────────────

    def _load_blip2(self) -> None:
        """Lazy-load BLIP-2 processor and model."""
        if self._blip2_model is None:
            from transformers import Blip2Processor, Blip2ForConditionalGeneration
            import torch
            print(f"  Loading BLIP-2 ({self.blip2_model}) on {self.device}...")
            self._blip2_processor = Blip2Processor.from_pretrained(self.blip2_model)
            self._blip2_model = Blip2ForConditionalGeneration.from_pretrained(
                self.blip2_model,
                torch_dtype=torch.float16 if self.device != "cpu" else torch.float32,
            )
            if self.device != "cpu":
                self._blip2_model = self._blip2_model.to(self.device)
            self._blip2_model.eval()
            print("  BLIP-2 loaded.")

    def generate_caption(self, image_path: str) -> str:
        """
        Generate a text caption from an image file using BLIP-2.

        Parameters
        ----------
        image_path : str  Path to a JPEG/PNG image file.

        Returns
        -------
        str — generated caption text.
        """
        self._load_blip2()
        from PIL import Image
        import torch
        image = Image.open(image_path).convert("RGB")
        inputs = self._blip2_processor(images=image, return_tensors="pt")
        if self.device != "cpu":
            inputs = {k: v.to(self.device) for k, v in inputs.items()}
        with torch.no_grad():
            out = self._blip2_model.generate(**inputs, max_new_tokens=50)
        caption = self._blip2_processor.decode(out[0], skip_special_tokens=True)
        return caption.strip()

    # ── Grounding scoring ─────────────────────────────────────────────────────

    def score_grounding(
        self,
        agent_text: str,
        image_path: Optional[str] = None,
        image_description: Optional[str] = None,
    ) -> float:
        """
        Score visual grounding: how well does agent_text match the image content?

        Parameters
        ----------
        agent_text        : str   Agent's response or answer text.
        image_path        : str   Path to image file (uses BLIP-2 captioning).
        image_description : str   Text description of image (no BLIP-2 required).

        Exactly one of image_path or image_description must be provided.

        Returns
        -------
        float in [0, 1] — 1.0 = perfectly grounded, 0.0 = unrelated.

        Note: This is a MATCH score (higher = better). For failure-risk evaluation
        (AUROC where label=1 means wrong/ungrounded), use score_risk() instead.
        """
        if image_path is None and image_description is None:
            raise ValueError("Provide either image_path or image_description")

        if image_path is not None:
            description = self.generate_caption(image_path)
        else:
            description = image_description

        emb_agent = self._embed_text(agent_text)
        emb_image = self._embed_text(description)

        # Cosine similarity (embeddings are unit-normalised)
        sim = float(np.dot(emb_agent, emb_image))
        return float(np.clip(sim, 0.0, 1.0))

    def score_risk(
        self,
        agent_text: str,
        image_path: Optional[str] = None,
        image_description: Optional[str] = None,
    ) -> float:
        """
        Grounding failure risk: 1 - score_grounding().

        Returns a risk score in [0, 1] where 1.0 = completely ungrounded (wrong),
        0.0 = perfectly grounded (correct). Use this for AUROC evaluation where
        label=1 means the agent's response is not grounded in the image.

        Parameters
        ----------
        agent_text        : str   Agent's response or answer text.
        image_path        : str   Path to image file (uses BLIP-2 captioning).
        image_description : str   Text description of image (no BLIP-2 required).

        Returns
        -------
        float in [0, 1] — 1.0 = ungrounded/wrong, 0.0 = grounded/correct.
        """
        return 1.0 - self.score_grounding(
            agent_text,
            image_path=image_path,
            image_description=image_description,
        )
