#!/usr/bin/env python3
"""
Cookbook 09 — Judge Images and Audio with Your LLM

SCENARIO:
    Your AI generates product descriptions for an e-commerce catalog.
    Each product has a photo and a generated text description. You need
    to verify that descriptions actually match what's in the image —
    not hallucinate features, colors, or details.

    You also have a customer service bot that transcribes voicemails.
    You need to verify transcription accuracy against the audio.

    This cookbook shows how to:
    1. Pass image URLs to the LLM judge alongside text
    2. Auto-generate grading criteria from a short description
    3. Evaluate image-text alignment with Gemini's native vision
    4. Combine multimodal with the feedback loop

Usage:
    export GOOGLE_API_KEY=...
    cd python && uv run python -m examples.09_multimodal_judge
"""

import os
import sys

# ---------------------------------------------------------------------------
# Check for API key
# ---------------------------------------------------------------------------
if not os.environ.get("GOOGLE_API_KEY"):
    print("This cookbook requires GOOGLE_API_KEY for Gemini vision.")
    print("export GOOGLE_API_KEY=your-key")
    sys.exit(1)

MODEL = "gemini/gemini-2.5-flash"

from fi.evals import evaluate


# ===================================================================
# SECTION 1: Image-Text Alignment — Does the description match?
# ===================================================================
print("=" * 70)
print("SECTION 1: Image-Text Alignment")
print("=" * 70)

# Google Cloud sample images (publicly accessible, no auth needed)
FLOWER_IMAGE = "https://storage.googleapis.com/cloud-samples-data/ai-platform/flowers/daisy/100080576_f52e8ee070_n.jpg"

# Good description — matches the image
print("\n--- Test 1: Accurate description ---")
result = evaluate(
    prompt="""Rate how accurately the text description matches the image.
    Score 1.0 if every detail in the description is visible in the image.
    Score 0.5 if the description is partially correct but has some inaccuracies.
    Score 0.0 if the description is completely wrong or describes something else.""",
    output="A white daisy flower with a yellow center, growing in a garden.",
    image_url=FLOWER_IMAGE,
    engine="llm",
    model=MODEL,
)
print(f"  Score: {result.score}")
print(f"  Reason: {result.reason[:120]}...")

# Bad description — hallucinated details
print("\n--- Test 2: Hallucinated description ---")
result = evaluate(
    prompt="""Rate how accurately the text description matches the image.
    Score 1.0 if every detail in the description is visible in the image.
    Score 0.5 if the description is partially correct but has some inaccuracies.
    Score 0.0 if the description is completely wrong or describes something else.""",
    output="A golden retriever puppy playing fetch on a sandy beach.",
    image_url=FLOWER_IMAGE,
    engine="llm",
    model=MODEL,
)
print(f"  Score: {result.score}")
print(f"  Reason: {result.reason[:120]}...")


# ===================================================================
# SECTION 2: Auto-Generate Grading Criteria (generate_prompt=True)
# ===================================================================
print("\n" + "=" * 70)
print("SECTION 2: Auto-Generate Grading Criteria")
print("=" * 70)

TULIP_IMAGE = "https://storage.googleapis.com/cloud-samples-data/ai-platform/flowers/tulips/10791227_7168491604.jpg"

# Instead of writing a detailed rubric, just describe what you want
print("\n--- Test 3: Generated prompt for flower identification ---")
result = evaluate(
    prompt="flower species identification accuracy from photos",
    output="This image shows a bright red tulip in full bloom with green stems.",
    image_url=TULIP_IMAGE,
    engine="llm",
    model=MODEL,
    generate_prompt=True,
)
print(f"  Score: {result.score}")
print(f"  Reason: {result.reason[:120]}...")

# You can also see what prompt was generated by calling the generator directly
from fi.evals.core.prompt_generator import generate_grading_criteria

criteria = generate_grading_criteria(
    "product photo quality assessment for e-commerce listings",
    MODEL,
    {"image_url": "...", "output": "..."},
)
print(f"\n  Generated criteria preview:\n  {criteria[:200]}...")


# ===================================================================
# SECTION 3: Multiple Images — Compare Input vs Output
# ===================================================================
print("\n" + "=" * 70)
print("SECTION 3: Comparing Images")
print("=" * 70)

ROSE_IMAGE = "https://storage.googleapis.com/cloud-samples-data/ai-platform/flowers/roses/12240303_80d87f77a3_n.jpg"

print("\n--- Test 4: Does the output description match the correct image? ---")
result = evaluate(
    prompt="""You are given two images and a text description.
    input_image_url is the reference product photo.
    output_image_url is what the AI selected as matching.
    Rate whether the text description matches the input_image_url (1.0) or
    actually describes the output_image_url instead (0.0).""",
    output="A pink rose flower in a garden.",
    input_image_url=FLOWER_IMAGE,
    output_image_url=ROSE_IMAGE,
    engine="llm",
    model=MODEL,
)
print(f"  Score: {result.score}")
print(f"  Reason: {result.reason[:120]}...")


# ===================================================================
# SECTION 4: Audio Transcription Check
# ===================================================================
print("\n" + "=" * 70)
print("SECTION 4: Audio Evaluation")
print("=" * 70)

# Google Cloud public audio sample
AUDIO_URL = "https://storage.googleapis.com/cloud-samples-data/speech/brooklyn_bridge.flac"

print("\n--- Test 5: Audio transcription accuracy ---")
result = evaluate(
    prompt="""Rate how accurately the transcription captures the audio content.
    Score 1.0 if the transcription is accurate and complete.
    Score 0.5 if partially correct. Score 0.0 if completely wrong.""",
    output="How old is the Brooklyn Bridge?",
    audio_url=AUDIO_URL,
    engine="llm",
    model=MODEL,
)
print(f"  Score: {result.score}")
print(f"  Reason: {result.reason[:120]}...")


# ===================================================================
# SECTION 5: Text-Only Still Works (Backwards Compatible)
# ===================================================================
print("\n" + "=" * 70)
print("SECTION 5: Text-Only Still Works")
print("=" * 70)

print("\n--- Test 6: Regular text judge (no images) ---")
result = evaluate(
    prompt="Rate the factual accuracy of the response given the context.",
    output="The Eiffel Tower is 330 meters tall and was built in 1889.",
    context="The Eiffel Tower, completed in 1889, stands at 330 metres.",
    engine="llm",
    model=MODEL,
)
print(f"  Score: {result.score}")
print(f"  Reason: {result.reason[:120]}...")

print("\n" + "=" * 70)
print("ALL TESTS COMPLETE")
print("=" * 70)
