"""VLM-powered task analyzer — uses vision-language models to understand task difficulty
from observation frames.

When an API key is available (ANTHROPIC_API_KEY or GOOGLE_API_KEY), samples frames
from episodes and asks a VLM structured questions about the task. Falls back to
heuristic-only analysis when no key is available.
"""

from __future__ import annotations

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

import numpy as np

from orbit.analyzer.task_context import TaskContext

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Data models
# ---------------------------------------------------------------------------


@dataclass
class VLMTaskContext:
    """VLM-enhanced task understanding."""

    task_description: str  # What the VLM thinks the task is
    task_type: str  # pick_and_place, insertion, wiping, etc.
    difficulty_tier: str  # simple, moderate, hard, very_hard
    num_objects: int  # Number of objects involved
    is_bimanual: bool  # VLM detected bimanual
    has_deformable: bool  # Deformable objects detected
    has_liquids: bool  # Liquids detected
    has_tools: bool  # Tool use detected
    requires_precision: bool  # Fine manipulation required
    failure_modes: list[str]  # Predicted failure modes
    strategy_count: int  # Number of distinct strategies observed
    vlm_complexity_score: float  # 0-1 from VLM reasoning
    vlm_confidence: float  # 0-1 how confident the VLM is
    model_used: str  # Which model was used
    tokens_used: int  # API tokens consumed
    cost_estimate: float  # Estimated cost in USD
    raw_response: dict = field(default_factory=dict)


# ---------------------------------------------------------------------------
# VLM prompt
# ---------------------------------------------------------------------------

_SYSTEM_PROMPT = (
    "You are an expert robotics researcher analyzing robot manipulation demonstrations. "
    "You will be shown frames from robot demonstrations and must analyze the task being performed. "
    "Respond ONLY with a JSON object matching this exact schema:\n"
    "{\n"
    '  "task_description": "<what the robot is doing>",\n'
    '  "task_type": "<pick_and_place|insertion|pushing|wiping'
    '|pouring|stacking|manipulation|transport>",\n'
    '  "difficulty_tier": "<simple|moderate|hard|very_hard>",\n'
    '  "num_objects": <int>,\n'
    '  "is_bimanual": <bool>,\n'
    '  "has_deformable": <bool>,\n'
    '  "has_liquids": <bool>,\n'
    '  "has_tools": <bool>,\n'
    '  "requires_precision": <bool>,\n'
    '  "failure_modes": ["<likely failure mode 1>", "<likely failure mode 2>"],\n'
    '  "strategy_count": <how many distinct strategies you see across episodes>,\n'
    '  "complexity_score": <0.0 to 1.0>,\n'
    '  "confidence": <0.0 to 1.0>\n'
    "}"
)


def _build_user_prompt(n_episodes: int, n_frames: int) -> str:
    """Build the user prompt describing what frames are shown."""
    return (
        f"I'm showing you {n_frames} frames from "
        f"{n_episodes} different robot demonstration episodes. "
        "For each episode, frames are sampled at the start, "
        "25%, 50%, 75%, and end of the episode.\n\n"
        "Please analyze:\n"
        "1. What manipulation task is being performed?\n"
        "2. How many objects are involved?\n"
        "3. Is this bimanual (two arms)? Does it require precise contact?\n"
        "4. Are there deformable objects? Liquids? Tools?\n"
        "5. What's the estimated difficulty tier?\n"
        "6. What are the likely failure modes?\n"
        "7. Do different episodes show different strategies for the same task?\n"
    )


# ---------------------------------------------------------------------------
# Frame sampling
# ---------------------------------------------------------------------------


def sample_frames(
    observation_frames: list[list[np.ndarray]],
    frames_per_episode: int = 5,
    max_episodes: int = 3,
) -> list[np.ndarray]:
    """Sample representative frames from episodes.

    Args:
        observation_frames: List of episodes, each a list of frames (H, W, C).
        frames_per_episode: Number of frames to sample per episode.
        max_episodes: Maximum episodes to sample from.

    Returns:
        Flat list of sampled frames.
    """
    sampled = []
    # Sample evenly spaced episodes
    if len(observation_frames) <= max_episodes:
        episode_indices = list(range(len(observation_frames)))
    else:
        episode_indices = (
            np.linspace(0, len(observation_frames) - 1, max_episodes).astype(int).tolist()
        )

    for ep_idx in episode_indices:
        episode = observation_frames[ep_idx]
        if not episode:
            continue

        n_frames = len(episode)
        if n_frames <= frames_per_episode:
            sampled.extend(episode)
        else:
            # Sample at 0%, 25%, 50%, 75%, 100%
            positions = np.linspace(0, n_frames - 1, frames_per_episode).astype(int)
            for pos in positions:
                sampled.append(episode[pos])

    return sampled


def frames_to_base64(frames: list[np.ndarray]) -> list[str]:
    """Convert numpy frames to base64-encoded JPEG strings."""
    encoded = []
    try:
        import cv2

        for frame in frames:
            if frame is None or frame.size == 0:
                continue
            # Ensure uint8
            if frame.dtype != np.uint8:
                frame = (np.clip(frame, 0, 255)).astype(np.uint8)
            # Convert RGB to BGR for cv2
            if frame.ndim == 3 and frame.shape[2] == 3:
                frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
            _, buf = cv2.imencode(".jpg", frame, [cv2.IMWRITE_JPEG_QUALITY, 85])
            encoded.append(base64.b64encode(buf.tobytes()).decode("ascii"))
    except ImportError:
        logger.warning("OpenCV not available for frame encoding")
    return encoded


# ---------------------------------------------------------------------------
# VLM API calls
# ---------------------------------------------------------------------------


def _call_anthropic(
    frames_b64: list[str],
    user_prompt: str,
    api_key: str,
    timeout: float = 30.0,
) -> tuple[dict, int, float]:
    """Call Anthropic Claude for task analysis."""
    import anthropic

    client = anthropic.Anthropic(api_key=api_key, timeout=timeout)

    content: list[dict] = []
    for frame_b64 in frames_b64[:12]:  # Max 12 frames
        content.append(
            {
                "type": "image",
                "source": {
                    "type": "base64",
                    "media_type": "image/jpeg",
                    "data": frame_b64,
                },
            }
        )
    content.append({"type": "text", "text": user_prompt})

    response = client.messages.create(
        model="claude-sonnet-4-20250514",
        max_tokens=1024,
        system=_SYSTEM_PROMPT,
        messages=[{"role": "user", "content": content}],
    )

    text = response.content[0].text.strip()
    parsed = _parse_json(text)

    input_tokens = response.usage.input_tokens
    output_tokens = response.usage.output_tokens
    total_tokens = input_tokens + output_tokens
    cost = (input_tokens * 3 + output_tokens * 15) / 1_000_000

    return parsed, total_tokens, cost


def _call_gemini(
    frames_b64: list[str],
    user_prompt: str,
    api_key: str,
) -> tuple[dict, int, float]:
    """Call Google Gemini for task analysis."""
    import google.generativeai as genai

    genai.configure(api_key=api_key)
    model = genai.GenerativeModel("gemini-2.5-flash", system_instruction=_SYSTEM_PROMPT)

    parts: list[Any] = []
    for frame_b64 in frames_b64[:12]:
        parts.append(
            {
                "inline_data": {
                    "mime_type": "image/jpeg",
                    "data": frame_b64,
                }
            }
        )
    parts.append(user_prompt)

    response = model.generate_content(parts)
    text = response.text.strip()
    parsed = _parse_json(text)

    tokens_meta = getattr(response, "usage_metadata", None)
    token_count = getattr(tokens_meta, "total_token_count", 0) if tokens_meta else 0
    cost = max(0.001, token_count * 0.000001)

    return parsed, token_count, cost


def _parse_json(text: str) -> dict:
    """Robustly parse JSON from VLM response."""
    # Strip markdown code blocks
    if "```" in text:
        lines = text.split("\n")
        in_block = False
        json_lines = []
        for line in lines:
            if line.strip().startswith("```"):
                if in_block:
                    break
                in_block = True
                continue
            if in_block:
                json_lines.append(line)
        if json_lines:
            text = "\n".join(json_lines)

    try:
        return json.loads(text)
    except json.JSONDecodeError:
        pass

    start = text.find("{")
    end = text.rfind("}")
    if start != -1 and end != -1 and end > start:
        try:
            return json.loads(text[start : end + 1])
        except json.JSONDecodeError:
            pass

    logger.warning("Failed to parse VLM task analysis response")
    return {}


# ---------------------------------------------------------------------------
# Difficulty tier mapping
# ---------------------------------------------------------------------------

_TIER_TO_SCORE = {
    "simple": 0.20,
    "moderate": 0.45,
    "hard": 0.70,
    "very_hard": 0.90,
}


# ---------------------------------------------------------------------------
# Main analyzer
# ---------------------------------------------------------------------------


class VLMTaskAnalyzer:
    """Uses vision-language models to understand task difficulty from observation frames."""

    def analyze(
        self,
        observation_frames: list[list[np.ndarray]] | None = None,
        frames_base64: list[str] | None = None,
        api_key: str | None = None,
        anthropic_key: str | None = None,
    ) -> VLMTaskContext | None:
        """Analyze task from observation frames using a VLM.

        Args:
            observation_frames: List of episodes, each a list of frames (H, W, C).
            frames_base64: Pre-encoded frames as base64 strings.
            api_key: Google API key (falls back to GOOGLE_API_KEY env).
            anthropic_key: Anthropic API key (falls back to ANTHROPIC_API_KEY env).

        Returns:
            VLMTaskContext if analysis succeeds, None if no API key or failure.
        """
        # Resolve API keys
        gemini_key = api_key or os.environ.get("GOOGLE_API_KEY")
        claude_key = anthropic_key or os.environ.get("ANTHROPIC_API_KEY")

        if not gemini_key and not claude_key:
            logger.info("No API keys available for VLM task analysis")
            return None

        # Prepare frames
        if frames_base64 is None and observation_frames is not None:
            sampled = sample_frames(observation_frames)
            frames_base64 = frames_to_base64(sampled)

        if not frames_base64:
            logger.info("No frames available for VLM task analysis")
            return None

        # Check for empty/black frames
        # (simple heuristic: very short base64 strings are likely empty)
        valid_frames = [f for f in frames_base64 if len(f) > 100]
        if len(valid_frames) < 2:
            logger.warning("Most frames appear empty/black, skipping VLM analysis")
            return None

        n_episodes = min(3, len(observation_frames)) if observation_frames else 1
        user_prompt = _build_user_prompt(n_episodes, len(valid_frames))

        # Try Gemini first (cheaper)
        if gemini_key:
            try:
                parsed, tokens, cost = _call_gemini(valid_frames, user_prompt, gemini_key)
                return self._build_result(parsed, "gemini-2.5-flash", tokens, cost)
            except ImportError:
                logger.warning("google-generativeai not installed")
            except Exception as e:
                logger.warning("Gemini VLM task analysis failed: %s", e)

        # Try Claude
        if claude_key:
            try:
                parsed, tokens, cost = _call_anthropic(valid_frames, user_prompt, claude_key)
                return self._build_result(parsed, "claude-sonnet", tokens, cost)
            except ImportError:
                logger.warning("anthropic SDK not installed")
            except Exception as e:
                logger.warning("Claude VLM task analysis failed: %s", e)

        logger.info("All VLM task analysis calls failed")
        return None

    def _build_result(
        self,
        parsed: dict,
        model: str,
        tokens: int,
        cost: float,
    ) -> VLMTaskContext:
        """Build VLMTaskContext from parsed VLM JSON response."""
        tier = parsed.get("difficulty_tier", "moderate")
        complexity = _TIER_TO_SCORE.get(tier, 0.45)

        return VLMTaskContext(
            task_description=parsed.get("task_description", "Unknown task"),
            task_type=parsed.get("task_type", "manipulation"),
            difficulty_tier=tier,
            num_objects=int(parsed.get("num_objects", 1)),
            is_bimanual=bool(parsed.get("is_bimanual", False)),
            has_deformable=bool(parsed.get("has_deformable", False)),
            has_liquids=bool(parsed.get("has_liquids", False)),
            has_tools=bool(parsed.get("has_tools", False)),
            requires_precision=bool(parsed.get("requires_precision", False)),
            failure_modes=parsed.get("failure_modes", []),
            strategy_count=int(parsed.get("strategy_count", 1)),
            vlm_complexity_score=float(parsed.get("complexity_score", complexity)),
            vlm_confidence=float(parsed.get("confidence", 0.6)),
            model_used=model,
            tokens_used=tokens,
            cost_estimate=cost,
            raw_response=parsed,
        )

    def refine_task_context(
        self,
        heuristic_context: TaskContext,
        vlm_context: VLMTaskContext,
    ) -> TaskContext:
        """Refine a heuristic TaskContext with VLM insights.

        The VLM can see things heuristics can't (deformable objects, specific task types).
        When VLM is confident, we weight its assessment more heavily.

        Args:
            heuristic_context: TaskContext from heuristic analysis.
            vlm_context: VLMTaskContext from VLM analysis.

        Returns:
            Refined TaskContext combining both sources.
        """
        vlm_weight = min(0.6, vlm_context.vlm_confidence * 0.7)
        heuristic_weight = 1.0 - vlm_weight

        # Blend complexity scores
        blended_complexity = (
            heuristic_weight * heuristic_context.task_complexity_score
            + vlm_weight * vlm_context.vlm_complexity_score
        )

        # VLM overrides for specific properties
        is_bimanual = vlm_context.is_bimanual or heuristic_context.is_bimanual

        # Deformable/liquid/tool bonuses (these are hard to detect heuristically)
        complexity_bonus = 0.0
        if vlm_context.has_deformable:
            complexity_bonus += 0.10
        if vlm_context.has_liquids:
            complexity_bonus += 0.08
        if vlm_context.has_tools:
            complexity_bonus += 0.05
        if vlm_context.requires_precision:
            complexity_bonus += 0.05

        blended_complexity = float(np.clip(blended_complexity + complexity_bonus, 0.0, 1.0))

        # Multimodality from VLM
        vlm_multimodality = min(1.0, vlm_context.strategy_count / 3.0) * 0.5
        blended_multimodality = (
            heuristic_weight * heuristic_context.trajectory_multimodality
            + vlm_weight * vlm_multimodality
        )

        from orbit.analyzer.task_context import _estimate_task_type

        return TaskContext(
            task_complexity_score=blended_complexity,
            dimension_complexity=heuristic_context.dimension_complexity,
            temporal_complexity=heuristic_context.temporal_complexity,
            coordination_complexity=heuristic_context.coordination_complexity,
            object_interaction_complexity=max(
                heuristic_context.object_interaction_complexity,
                vlm_context.num_objects / 5.0,  # VLM object count as proxy
            ),
            trajectory_multimodality=float(np.clip(blended_multimodality, 0.0, 1.0)),
            is_bimanual=is_bimanual,
            action_dim=heuristic_context.action_dim,
            mean_episode_length=heuristic_context.mean_episode_length,
            estimated_task_type=_estimate_task_type(blended_complexity),
            complexity_breakdown={
                **heuristic_context.complexity_breakdown,
                "vlm_complexity": vlm_context.vlm_complexity_score,
                "vlm_deformable": float(vlm_context.has_deformable),
                "vlm_precision": float(vlm_context.requires_precision),
            },
        )
