"""Structured Gemini Vision quality assessment for robot demonstration datasets.

Sends sampled frames from episodes to Gemini Flash with specific quality questions.
Only called when the dataset has image observations. Skips gracefully for state-only datasets.

Cost: ~$0.003 per analysis (9 images to Gemini Flash).
"""

from __future__ import annotations

import json
import logging
import os
from dataclasses import dataclass, field
from typing import Any

import numpy as np

logger = logging.getLogger(__name__)


QUALITY_ASSESSMENT_PROMPT = """You are analyzing camera frames from a robot manipulation dataset.
I will show you frames from 3 different episodes at 3 timepoints each (start, middle, end).

For each episode, answer these questions with YES/NO and a confidence score (0.0-1.0):

1. TASK_VISIBLE: Can you see the robot performing a clear manipulation task?
2. TASK_COMPLETE: Does the final frame show what appears to be a completed task (object moved, assembled, etc)?
3. VISUAL_QUALITY: Are the camera images clear (not blurry, not occluded, good lighting)?
4. ANOMALY: Does any episode look clearly different or problematic compared to the others?

Also answer this cross-episode question:
5. CONSISTENT_TASK: Do all 3 episodes appear to be attempting the SAME task?

Respond ONLY in this JSON format:
{
  "episodes": [
    {
      "episode_index": 0,
      "task_visible": {"answer": true, "confidence": 0.9},
      "task_complete": {"answer": true, "confidence": 0.7},
      "visual_quality": {"answer": true, "confidence": 0.95},
      "anomaly": {"answer": false, "confidence": 0.85}
    }
  ],
  "consistent_task": {"answer": true, "confidence": 0.9},
  "task_description": "Robot arm picking up a red cube and placing it in a bin",
  "concerns": ["Episode 2 appears to end before task completion"]
}
"""


@dataclass
class VisionQualityResult:
    """Result of structured vision quality assessment."""

    task_visible_all: bool  # All episodes show visible task
    task_complete_fraction: float  # Fraction of episodes with completed task
    consistent_task: bool  # Episodes show same task
    visual_quality_all: bool  # All episodes have good visual quality
    anomaly_episodes: list[int]  # Episodes flagged as anomalous
    task_description: str  # VLM's description of the task
    concerns: list[str]  # Specific concerns raised
    grade_adjustments: dict[str, int]  # penalty name -> points to subtract
    total_penalty: int  # Sum of all penalties
    raw_response: dict = field(default_factory=dict)
    model_used: str = "gemini"
    error: str | None = None  # Set if assessment failed


def _has_image_observations(dataset_info: Any) -> bool:
    """Check if dataset has image observations."""
    if hasattr(dataset_info, "observation_keys"):
        for key in (dataset_info.observation_keys or []):
            if "image" in key.lower():
                return True
    if hasattr(dataset_info, "camera_names"):
        if dataset_info.camera_names:
            return True
    if hasattr(dataset_info, "features"):
        features = dataset_info.features or {}
        for key in features:
            if "image" in key.lower():
                return True
    return False


def _select_episodes(
    episode_qualities: dict[int, float] | None,
    total_episodes: int,
) -> list[int]:
    """Select 3 episodes: best quality, worst quality, and random middle.

    Args:
        episode_qualities: Mapping of episode index to quality score.
            If None, picks indices 0, total//2, total-1.
        total_episodes: Total number of episodes.

    Returns:
        List of 3 episode indices.
    """
    if total_episodes <= 0:
        return []
    if total_episodes <= 3:
        return list(range(total_episodes))

    if episode_qualities and len(episode_qualities) >= 3:
        sorted_eps = sorted(episode_qualities.items(), key=lambda x: x[1])
        worst = sorted_eps[0][0]
        best = sorted_eps[-1][0]
        # Pick a middle episode (not best or worst)
        middle_candidates = [idx for idx, _ in sorted_eps[1:-1]]
        if middle_candidates:
            rng = np.random.RandomState(42)
            middle = middle_candidates[rng.randint(len(middle_candidates))]
        else:
            middle = sorted_eps[len(sorted_eps) // 2][0]
        return [best, middle, worst]

    # Fallback: evenly spaced
    return [0, total_episodes // 2, total_episodes - 1]


def _extract_frame_indices(episode_length: int) -> list[int]:
    """Extract frame indices at 10%, 50%, 90% of episode length."""
    if episode_length <= 0:
        return []
    if episode_length <= 3:
        return list(range(episode_length))
    return [
        int(episode_length * 0.10),
        int(episode_length * 0.50),
        int(episode_length * 0.90),
    ]


def _parse_response(response_text: str) -> dict:
    """Parse JSON response from Gemini, handling markdown code blocks."""
    text = response_text.strip()
    # Strip markdown code block if present
    if text.startswith("```"):
        lines = text.split("\n")
        # Remove first and last lines (```json and ```)
        lines = [l for l in lines if not l.strip().startswith("```")]
        text = "\n".join(lines)
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        logger.warning("Failed to parse Gemini vision quality response as JSON")
        return {}


def _compute_adjustments(parsed: dict) -> dict[str, int]:
    """Map parsed VLM response to grade adjustments."""
    adjustments: dict[str, int] = {}

    episodes = parsed.get("episodes", [])
    if not episodes:
        return adjustments

    # task_visible: if any episode has task_visible=False, penalize
    invisible = [e for e in episodes if not e.get("task_visible", {}).get("answer", True)]
    if invisible:
        adjustments["camera_misconfigured"] = -5

    # task_complete: if >50% episodes incomplete
    incomplete = [e for e in episodes if not e.get("task_complete", {}).get("answer", True)]
    if len(incomplete) > len(episodes) / 2:
        adjustments["incomplete_tasks"] = -10

    # consistent_task
    consistent = parsed.get("consistent_task", {})
    if not consistent.get("answer", True):
        adjustments["inconsistent_tasks"] = -8

    # visual_quality
    poor_quality = [e for e in episodes if not e.get("visual_quality", {}).get("answer", True)]
    if poor_quality:
        adjustments["poor_visual_quality"] = -3

    return adjustments


def assess_vision_quality(
    dataset_info: Any,
    frame_loader: Any | None = None,
    episode_qualities: dict[int, float] | None = None,
) -> VisionQualityResult:
    """Run structured Gemini vision quality assessment.

    Args:
        dataset_info: DatasetInfo with metadata about the dataset.
        frame_loader: Callable(episode_idx, frame_idx) -> PIL.Image or np.ndarray.
            If None, assessment is skipped.
        episode_qualities: Optional per-episode quality scores for episode selection.

    Returns:
        VisionQualityResult. If no images or no API key, returns a no-op result with zero penalty.
    """
    # Check if dataset has images
    if not _has_image_observations(dataset_info):
        return VisionQualityResult(
            task_visible_all=True,
            task_complete_fraction=1.0,
            consistent_task=True,
            visual_quality_all=True,
            anomaly_episodes=[],
            task_description="State-only dataset — no visual assessment",
            concerns=[],
            grade_adjustments={},
            total_penalty=0,
            error="No image observations in dataset",
        )

    if frame_loader is None:
        return VisionQualityResult(
            task_visible_all=True,
            task_complete_fraction=1.0,
            consistent_task=True,
            visual_quality_all=True,
            anomaly_episodes=[],
            task_description="Frame loader not available",
            concerns=[],
            grade_adjustments={},
            total_penalty=0,
            error="No frame loader provided",
        )

    api_key = os.environ.get("GOOGLE_API_KEY")
    if not api_key:
        return VisionQualityResult(
            task_visible_all=True,
            task_complete_fraction=1.0,
            consistent_task=True,
            visual_quality_all=True,
            anomaly_episodes=[],
            task_description="No API key available",
            concerns=[],
            grade_adjustments={},
            total_penalty=0,
            error="GOOGLE_API_KEY not set",
        )

    num_episodes = getattr(dataset_info, "num_episodes", 0)
    selected = _select_episodes(episode_qualities, num_episodes)
    if not selected:
        return VisionQualityResult(
            task_visible_all=True,
            task_complete_fraction=1.0,
            consistent_task=True,
            visual_quality_all=True,
            anomaly_episodes=[],
            task_description="No episodes to assess",
            concerns=[],
            grade_adjustments={},
            total_penalty=0,
            error="No episodes available",
        )

    # Load frames
    import base64
    import io

    frames_data = []
    for ep_idx in selected:
        ep_length = getattr(dataset_info, "avg_episode_length", 100)
        frame_indices = _extract_frame_indices(int(ep_length))
        for f_idx in frame_indices:
            try:
                frame = frame_loader(ep_idx, f_idx)
                if frame is not None:
                    # Convert to base64 for API
                    if hasattr(frame, "save"):
                        # PIL Image
                        buf = io.BytesIO()
                        frame.save(buf, format="JPEG")
                        frames_data.append(base64.b64encode(buf.getvalue()).decode())
                    elif isinstance(frame, np.ndarray):
                        from PIL import Image
                        img = Image.fromarray(frame.astype(np.uint8))
                        buf = io.BytesIO()
                        img.save(buf, format="JPEG")
                        frames_data.append(base64.b64encode(buf.getvalue()).decode())
            except Exception as e:
                logger.debug(f"Failed to load frame ep={ep_idx} f={f_idx}: {e}")

    if not frames_data:
        return VisionQualityResult(
            task_visible_all=True,
            task_complete_fraction=1.0,
            consistent_task=True,
            visual_quality_all=True,
            anomaly_episodes=[],
            task_description="Could not load any frames",
            concerns=[],
            grade_adjustments={},
            total_penalty=0,
            error="Frame loading failed",
        )

    # Call Gemini Flash
    try:
        import google.generativeai as genai

        genai.configure(api_key=api_key)
        model = genai.GenerativeModel("gemini-2.0-flash")

        # Build content with images
        content_parts = [QUALITY_ASSESSMENT_PROMPT]
        for i, b64 in enumerate(frames_data):
            ep_num = i // 3
            frame_num = i % 3
            frame_label = ["start", "middle", "end"][frame_num] if frame_num < 3 else f"frame_{frame_num}"
            content_parts.append(f"\nEpisode {ep_num + 1}, {frame_label} frame:")
            content_parts.append({
                "mime_type": "image/jpeg",
                "data": b64,
            })

        response = model.generate_content(content_parts)
        parsed = _parse_response(response.text)

    except Exception as e:
        logger.warning(f"Gemini vision quality assessment failed: {e}")
        return VisionQualityResult(
            task_visible_all=True,
            task_complete_fraction=1.0,
            consistent_task=True,
            visual_quality_all=True,
            anomaly_episodes=[],
            task_description="Assessment failed",
            concerns=[],
            grade_adjustments={},
            total_penalty=0,
            error=str(e),
        )

    # Parse results
    episodes = parsed.get("episodes", [])
    adjustments = _compute_adjustments(parsed)
    total_penalty = sum(adjustments.values())

    # Extract per-field summaries
    task_visible_all = all(
        e.get("task_visible", {}).get("answer", True) for e in episodes
    )
    complete_count = sum(
        1 for e in episodes if e.get("task_complete", {}).get("answer", True)
    )
    task_complete_fraction = complete_count / max(len(episodes), 1)
    consistent = parsed.get("consistent_task", {}).get("answer", True)
    visual_quality_all = all(
        e.get("visual_quality", {}).get("answer", True) for e in episodes
    )
    anomaly_eps = [
        selected[i] for i, e in enumerate(episodes)
        if e.get("anomaly", {}).get("answer", False) and i < len(selected)
    ]

    return VisionQualityResult(
        task_visible_all=task_visible_all,
        task_complete_fraction=task_complete_fraction,
        consistent_task=consistent,
        visual_quality_all=visual_quality_all,
        anomaly_episodes=anomaly_eps,
        task_description=parsed.get("task_description", "Unknown"),
        concerns=parsed.get("concerns", []),
        grade_adjustments=adjustments,
        total_penalty=total_penalty,
        raw_response=parsed,
        model_used="gemini-2.0-flash",
    )
