"""
Google Gemini adapter for the LLM abstraction layer.

This module provides an adapter for Google's Gemini models
that implements the LLMInterface.

Includes retry logic with exponential backoff for transient errors
(timeouts, rate limits, empty responses).

Retry configuration: 3 retries with 4s exponential backoff.
"""

import asyncio
import base64
import json
import re
import logging
import random
from typing import List, Optional, Type, Union

from google import genai
from google.genai import types
from pydantic import BaseModel

from inceptbench_new.llm.base import LLMImage, LLMInterface, LLMMessage
from inceptbench_new.utils.failure_tracker import AttemptError, FailureTracker

logger = logging.getLogger(__name__)

# Retry configuration: 3 retries with 4s exponential backoff
GEMINI_MAX_RETRIES = 3
GEMINI_BASE_DELAY = 4.0

# ---------------------------------------------------------------------------
# Generalized prompts for two-pass agentic vision (all educational subjects)
# ---------------------------------------------------------------------------

_VISION_PASS_SUFFIX = """

Do NOT write code. Use only your visual understanding.

Return ONLY a JSON object (no markdown fences):
{"description": "what the image shows", "needs_measurement": true/false, "items": [{"label": "A", "answer": "...", "explanation": "..."}, ...]}

Set needs_measurement to true if the answer requires: counting evenly-spaced divisions; measuring fractions, ratios, or lengths; reading exact values from charts/graphs/maps; extracting scale readings; or any precision measurement (math, science, history, ELA). Set to false for counting discrete objects, identifying shapes/labels, or reading clear text/equations."""

_MEASUREMENT_PROMPT = """You previously analyzed this image visually and found:
{vision_result}

Now verify and refine the answers using Python code execution for precise measurement. Load the image with OpenCV/PIL and measure pixel distances, ratios, or other quantities that require precision.

Subject-specific guidance:
- Math: fractions (denominator = number of PARTS, not dividers), number lines, bar/pie chart values (read y-axis labels first for scale), geometry angles/lengths/areas.
- Science: diagram scales, experimental data from charts, axis values.
- History/Social Studies: map distances, timeline positions, statistical chart values.
- ELA: chart/graph values in informational contexts.

For bar charts: read the ACTUAL axis tick labels to establish scale, then measure bar positions. For divided regions: N vertical dividers create N+1 segments; N marks including endpoints create N-1 segments. Do NOT re-count discrete objects — trust the visual counts. Only use code for measurement precision.

Your code must end by printing the final JSON so it is captured: print(json.dumps(result)). The array must be of the form [{"label": "CategoryName", "answer": "value", "explanation": "..."}, ...]."""


def _check_needs_measurement(text: str) -> bool:
    """Parse Pass 1 JSON response for needs_measurement flag."""
    if not text or not text.strip():
        return False
    clean = re.sub(r"^```\w*\n(.*)\n```$", r"\1", text.strip(), flags=re.DOTALL).strip()
    try:
        data = json.loads(clean)
        if isinstance(data, dict):
            return bool(data.get("needs_measurement", False))
        return False
    except (json.JSONDecodeError, Exception):
        return '"needs_measurement": true' in text.lower() or '"needs_measurement":true' in text.lower()


def _clean_schema_for_gemini(schema: dict) -> dict:
    """
    Recursively strip fields that Gemini's API doesn't support from a
    JSON Schema dict (e.g. ``additionalProperties``, ``$defs``, ``title``,
    ``$ref``).

    Pydantic v2 emits these by default in ``model_json_schema()`` and the
    Google GenAI SDK converts them to snake_case (``additional_properties``)
    which Gemini rejects with ``INVALID_ARGUMENT``.
    """
    STRIP_KEYS = {"additionalProperties", "title", "$defs", "$ref"}
    if isinstance(schema, dict):
        cleaned = {}
        for k, v in schema.items():
            if k in STRIP_KEYS:
                continue
            cleaned[k] = _clean_schema_for_gemini(v)
        return cleaned
    if isinstance(schema, list):
        return [_clean_schema_for_gemini(item) for item in schema]
    return schema


def _pydantic_to_gemini_schema(model_class: Type[BaseModel]) -> dict:
    """
    Convert a Pydantic model class to a Gemini-compatible JSON schema dict.

    Handles Pydantic v2's ``anyOf`` pattern for ``Optional`` fields by
    inlining the non-null type and marking the field ``nullable``.
    """
    raw = model_class.model_json_schema()

    # Resolve $defs references first
    defs = raw.pop("$defs", {})

    def _resolve_refs(obj):
        if isinstance(obj, dict):
            if "$ref" in obj:
                ref_path = obj["$ref"]  # e.g. "#/$defs/SomeModel"
                ref_name = ref_path.split("/")[-1]
                if ref_name in defs:
                    return _resolve_refs(defs[ref_name])
                logger.warning(
                    "Unresolvable $ref %r (available defs: %s)",
                    ref_path,
                    list(defs.keys()),
                )
                remaining = {k: v for k, v in obj.items() if k != "$ref"}
                return _resolve_refs(remaining) if remaining else {}
            return {k: _resolve_refs(v) for k, v in obj.items()}
        if isinstance(obj, list):
            return [_resolve_refs(item) for item in obj]
        return obj

    resolved = _resolve_refs(raw)

    # Convert anyOf[type, null] patterns to nullable
    def _fix_optional(obj):
        if isinstance(obj, dict):
            if "anyOf" in obj:
                variants = obj["anyOf"]
                non_null = [v for v in variants if v.get("type") != "null"]
                has_null = any(v.get("type") == "null" for v in variants)
                if has_null and len(non_null) == 1:
                    merged = {**obj, **non_null[0]}
                    merged.pop("anyOf", None)
                    merged["nullable"] = True
                    return {k: _fix_optional(v) for k, v in merged.items()}
            return {k: _fix_optional(v) for k, v in obj.items()}
        if isinstance(obj, list):
            return [_fix_optional(item) for item in obj]
        return obj

    fixed = _fix_optional(resolved)
    return _clean_schema_for_gemini(fixed)


class GeminiAdapter(LLMInterface):
    """
    Adapter for Google Gemini API.
    
    This adapter handles all Gemini-specific details:
    - Authentication and API client setup
    - Request/response format conversion
    - Structured output via response schemas
    - Vision input handling
    - Error handling and retries
    """
    
    def __init__(
        self,
        model: str,
        api_key: str,
        timeout: float = 60.0,
        temperature: float = 0.0,
        max_tokens: int = 16384
    ):
        """
        Initialize Gemini adapter.
        
        Args:
            model: Model identifier (e.g., "gemini-3-pro-preview")
            api_key: Google AI API key
            timeout: Request timeout in seconds
            temperature: Sampling temperature (0.0 = deterministic)
            max_tokens: Maximum tokens in response
        """
        super().__init__(
            model=model,
            api_key=api_key,
            timeout=timeout,
            temperature=temperature,
            max_tokens=max_tokens
        )
        self.client = genai.Client(api_key=self.api_key)
        logger.debug(f"Gemini adapter initialized: {self.model}")

    @staticmethod
    def _safe_extract_text(response) -> str:
        """Extract text from a GenerateContentResponse without using .text property.

        The SDK's .text property can raise KeyError on responses whose JSON
        contains keys like "label". Manually walking the parts avoids this.
        """
        if response is None:
            return ""
        try:
            return (response.text or "").strip()
        except Exception:
            pass
        # Fallback: walk parts manually
        text_parts = []
        try:
            for candidate in (response.candidates or []):
                for part in (candidate.content.parts or []):
                    try:
                        t = part.text
                        if t:
                            text_parts.append(t)
                    except Exception:
                        pass
                    try:
                        cer = part.code_execution_result
                        if cer and cer.output:
                            text_parts.append(cer.output)
                    except Exception:
                        pass
        except Exception:
            pass
        return "".join(text_parts).strip()

    async def _stream_and_collect(
        self,
        contents: list,
        config: types.GenerateContentConfig,
    ) -> str:
        """
        Run Gemini generate_content_stream in a thread and collect full response text.
        Used for vision calls to avoid blocking/timeouts on long code-execution runs.
        Returns concatenated text from all parts (including code execution output).
        """
        def _run_stream():
            text_parts = []
            for chunk in self.client.models.generate_content_stream(
                model=self.model,
                contents=contents,
                config=config,
            ):
                if not chunk.candidates:
                    continue
                for candidate in chunk.candidates:
                    try:
                        parts = candidate.content.parts or []
                    except Exception:
                        continue
                    for part in parts:
                        try:
                            t = part.text
                            if t:
                                text_parts.append(t)
                        except Exception:
                            pass
                        try:
                            cer = part.code_execution_result
                            if cer and cer.output:
                                text_parts.append(cer.output)
                        except Exception:
                            pass
            return "".join(text_parts).strip()

        return await asyncio.to_thread(_run_stream)

    async def generate_text(
        self,
        messages: List[LLMMessage],
        **kwargs
    ) -> str:
        """
        Generate plain text using Gemini API.
        
        Includes retry logic with exponential backoff for transient errors.
        
        Args:
            messages: Conversation messages
            **kwargs: Override temperature, max_tokens, etc.
            
        Returns:
            Generated text string
        """
        # Extract system instruction if present
        system_instruction = next(
            (msg.content for msg in messages if msg.role == "system"),
            None
        )
        
        # Build contents from non-system messages
        contents = []
        for msg in messages:
            if msg.role != "system":
                role = "user" if msg.role == "user" else "model"
                contents.append(types.Content(
                    role=role,
                    parts=[types.Part.from_text(text=msg.content)]
                ))
        
        last_error = None
        attempt_errors: List[AttemptError] = []
        
        for attempt in range(GEMINI_MAX_RETRIES + 1):
            try:
                response = await self.client.aio.models.generate_content(
                    model=self.model,
                    contents=contents,
                    config=types.GenerateContentConfig(
                        system_instruction=system_instruction,
                        temperature=kwargs.get("temperature", self.temperature),
                        max_output_tokens=kwargs.get("max_tokens", self.max_tokens)
                    )
                )
                
                # Check for None response
                if response is None or response.text is None:
                    raise ValueError("Gemini returned empty response")
                
                # Record recovery if this wasn't the first attempt
                if attempt > 0:
                    FailureTracker.record_recovered(
                        component="gemini_adapter.generate_text",
                        message=f"Succeeded on attempt {attempt + 1}/{GEMINI_MAX_RETRIES + 1}",
                        context={"model": self.model},
                        attempt_errors=attempt_errors if attempt_errors else None
                    )
                
                return response.text.strip()
                
            except Exception as e:
                last_error = e
                error_str = str(e).lower()
                
                # Record this attempt's error
                attempt_errors.append(AttemptError(
                    attempt=attempt + 1,
                    error_message=str(e)
                ))
                
                # Check if error is retryable
                is_retryable = any(keyword in error_str for keyword in [
                    "timeout", "rate", "quota", "429", "503", "500",
                    "overloaded", "empty response", "none"
                ])
                
                if is_retryable and attempt < GEMINI_MAX_RETRIES:
                    delay = GEMINI_BASE_DELAY * (2 ** attempt) + random.uniform(0, 1)
                    logger.warning(
                        f"Gemini text generation error (attempt {attempt + 1}/"
                        f"{GEMINI_MAX_RETRIES + 1}): {e}. Retrying in {delay:.1f}s..."
                    )
                    await asyncio.sleep(delay)
                else:
                    logger.error(f"Gemini text generation failed after all retries: {e}")
                    FailureTracker.record_exhausted(
                        component="gemini_adapter.generate_text",
                        error_message=str(e),
                        context={"model": self.model, "attempts": attempt + 1},
                        attempt_errors=attempt_errors
                    )
                    raise
        
        # Should not reach here, but just in case
        raise last_error

    async def generate_structured(
        self,
        messages: List[LLMMessage],
        response_schema: Type[BaseModel],
        **kwargs
    ) -> BaseModel:
        """
        Generate structured output using Gemini's native JSON mode.
        
        Includes retry logic with exponential backoff for transient errors.
        
        Args:
            messages: Conversation messages
            response_schema: Pydantic model class for output structure
            **kwargs: Override temperature, max_tokens, etc.
            
        Returns:
            Instance of response_schema with model output
        """
        # Extract system instruction if present
        system_instruction = next(
            (msg.content for msg in messages if msg.role == "system"),
            None
        )
        
        # Build contents from non-system messages
        contents = []
        for msg in messages:
            if msg.role != "system":
                role = "user" if msg.role == "user" else "model"
                contents.append(types.Content(
                    role=role,
                    parts=[types.Part.from_text(text=msg.content)]
                ))
        
        last_error = None
        attempt_errors: List[AttemptError] = []
        
        # Convert Pydantic schema to Gemini-compatible dict
        gemini_schema = _pydantic_to_gemini_schema(response_schema)

        for attempt in range(GEMINI_MAX_RETRIES + 1):
            try:
                response = await self.client.aio.models.generate_content(
                    model=self.model,
                    contents=contents,
                    config=types.GenerateContentConfig(
                        system_instruction=system_instruction,
                        temperature=kwargs.get("temperature", self.temperature),
                        max_output_tokens=kwargs.get("max_tokens", self.max_tokens),
                        response_mime_type="application/json",
                        response_schema=gemini_schema,
                    )
                )
                
                # Check for None response
                if response is None or response.text is None:
                    raise ValueError("Gemini returned empty response (None)")
                
                # Parse the JSON response into the schema
                parsed_data = json.loads(response.text)
                
                # Record recovery if this wasn't the first attempt
                if attempt > 0:
                    FailureTracker.record_recovered(
                        component="gemini_adapter.generate_structured",
                        message=f"Succeeded on attempt {attempt + 1}/{GEMINI_MAX_RETRIES + 1}",
                        context={"model": self.model, "schema": response_schema.__name__},
                        attempt_errors=attempt_errors if attempt_errors else None
                    )
                
                return response_schema(**parsed_data)
                
            except Exception as e:
                last_error = e
                error_str = str(e).lower()
                
                # Record this attempt's error
                attempt_errors.append(AttemptError(
                    attempt=attempt + 1,
                    error_message=str(e)
                ))
                
                # Check if error is retryable
                is_retryable = any(keyword in error_str for keyword in [
                    "timeout", "rate", "quota", "429", "503", "500",
                    "overloaded", "empty response", "none", "nonetype"
                ])
                
                if is_retryable and attempt < GEMINI_MAX_RETRIES:
                    delay = GEMINI_BASE_DELAY * (2 ** attempt) + random.uniform(0, 1)
                    logger.warning(
                        f"Gemini structured output error (attempt {attempt + 1}/"
                        f"{GEMINI_MAX_RETRIES + 1}): {e}. Retrying in {delay:.1f}s..."
                    )
                    await asyncio.sleep(delay)
                else:
                    logger.error(
                        f"Gemini structured output generation failed for "
                        f"{response_schema.__name__} after all retries: {e}"
                    )
                    FailureTracker.record_exhausted(
                        component="gemini_adapter.generate_structured",
                        error_message=str(e),
                        context={"model": self.model, "schema": response_schema.__name__, "attempts": attempt + 1},
                        attempt_errors=attempt_errors
                    )
                    raise
        
        # Should not reach here, but just in case
        raise last_error

    async def generate_with_vision(
        self,
        messages: List[LLMMessage],
        images: List[LLMImage],
        response_schema: Optional[Type[BaseModel]] = None,
        enable_code_execution: bool = False,
        **kwargs
    ) -> Union[str, BaseModel]:
        """
        Generate response with image inputs (Gemini vision).
        
        Handles both plain text and structured output with vision.
        Images can be provided as URLs or base64-encoded data.
        Includes retry logic with exponential backoff for transient errors.
        
        Args:
            messages: Conversation messages
            images: List of images to analyze
            response_schema: Optional schema for structured output
            enable_code_execution: Enable code execution tool for agentic vision
            **kwargs: Override temperature, max_tokens, etc.
            
        Returns:
            Text string if no schema, otherwise instance of response_schema
        """
        if not self.supports_vision:
            raise NotImplementedError(
                f"Model {self.model} does not support vision inputs"
            )
        
        # Extract system instruction if present
        system_instruction = next(
            (msg.content for msg in messages if msg.role == "system"),
            None
        )
        
        # Build parts list with text and images (image_parts kept for two-pass reuse)
        text_parts = []
        for msg in messages:
            if msg.role == "user":
                text_parts.append(types.Part.from_text(text=msg.content))
        image_parts = []
        for img in images:
            if img.url:
                image_parts.append(types.Part.from_uri(
                    file_uri=img.url,
                    mime_type=img.media_type or "image/png"
                ))
            elif img.base64_data:
                image_bytes = base64.b64decode(img.base64_data)
                image_parts.append(types.Part.from_bytes(
                    data=image_bytes,
                    mime_type=img.media_type or "image/png"
                ))
        parts = text_parts + image_parts
        contents = [types.Content(role="user", parts=parts)]

        def _parse_final_response(response_text: str):
            """Parse final response into response_schema if requested."""
            if not response_schema:
                return response_text.strip()
            text = re.sub(r"^```\w*\n(.*)\n```$", r"\1", response_text.strip(), flags=re.DOTALL).strip()
            try:
                data = json.loads(text)
                return response_schema(**data)
            except json.JSONDecodeError:
                pass
            # Try to find JSON object/array in text (skip braces inside strings)
            for open_c, close_c in (("[", "]"), ("{", "}")):
                i = text.find(open_c)
                if i == -1:
                    continue
                depth = 0
                in_str = False
                escape = False
                for j in range(i, len(text)):
                    c = text[j]
                    if escape:
                        escape = False
                        continue
                    if c == "\\":
                        if in_str:
                            escape = True
                        continue
                    if c == '"':
                        in_str = not in_str
                        continue
                    if in_str:
                        continue
                    if c == open_c:
                        depth += 1
                    elif c == close_c:
                        depth -= 1
                        if depth == 0:
                            try:
                                return response_schema(**json.loads(text[i : j + 1]))
                            except (json.JSONDecodeError, Exception):
                                break
            return response_text.strip()

        last_error = None
        attempt_errors: List[AttemptError] = []

        for attempt in range(GEMINI_MAX_RETRIES + 1):
            try:
                if enable_code_execution:
                    # Two-pass agentic vision: Pass 1 visual only, Pass 2 code exec if needed
                    user_text = "\n\n".join(msg.content for msg in messages if msg.role == "user")
                    pass1_parts = [types.Part.from_text(text=user_text + _VISION_PASS_SUFFIX)] + image_parts
                    pass1_contents = [types.Content(role="user", parts=pass1_parts)]
                    config_pass1 = types.GenerateContentConfig(
                        system_instruction=system_instruction,
                        temperature=kwargs.get("temperature", self.temperature),
                        max_output_tokens=kwargs.get("max_tokens", self.max_tokens),
                    )
                    # Pass 1 uses non-streaming async API (no code exec, fast)
                    pass1_response = await self.client.aio.models.generate_content(
                        model=self.model,
                        contents=pass1_contents,
                        config=config_pass1,
                    )
                    pass1_text = self._safe_extract_text(pass1_response)
                    if not pass1_text:
                        raise ValueError("Gemini returned empty response (Pass 1)")

                    if not _check_needs_measurement(pass1_text):
                        if attempt > 0:
                            FailureTracker.record_recovered(
                                component="gemini_adapter.generate_with_vision",
                                message=f"Succeeded on attempt {attempt + 1}/{GEMINI_MAX_RETRIES + 1}",
                                context={"model": self.model, "image_count": len(images)},
                                attempt_errors=attempt_errors if attempt_errors else None,
                            )
                        return _parse_final_response(pass1_text)

                    # Pass 2: measurement with code execution
                    measure_text = _MEASUREMENT_PROMPT.replace("{vision_result}", pass1_text)
                    pass2_parts = [types.Part.from_text(text=measure_text)] + image_parts
                    pass2_contents = [types.Content(role="user", parts=pass2_parts)]
                    config_pass2 = types.GenerateContentConfig(
                        system_instruction=system_instruction,
                        temperature=kwargs.get("temperature", self.temperature),
                        max_output_tokens=kwargs.get("max_tokens", self.max_tokens),
                        tools=[types.Tool(code_execution=types.ToolCodeExecution())],
                    )
                    pass2_text = await self._stream_and_collect(pass2_contents, config_pass2)
                    if not pass2_text:
                        raise ValueError("Gemini returned empty response (Pass 2)")
                    response_text = pass2_text
                else:
                    # Single pass, streaming (no code exec)
                    config = types.GenerateContentConfig(
                        system_instruction=system_instruction,
                        temperature=kwargs.get("temperature", self.temperature),
                        max_output_tokens=kwargs.get("max_tokens", self.max_tokens),
                    )
                    if response_schema:
                        config.response_mime_type = "application/json"
                        config.response_schema = _pydantic_to_gemini_schema(response_schema)
                    response_text = await self._stream_and_collect(contents, config)
                    if not response_text:
                        raise ValueError("Gemini returned empty response (None)")

                if attempt > 0:
                    FailureTracker.record_recovered(
                        component="gemini_adapter.generate_with_vision",
                        message=f"Succeeded on attempt {attempt + 1}/{GEMINI_MAX_RETRIES + 1}",
                        context={"model": self.model, "image_count": len(images)},
                        attempt_errors=attempt_errors if attempt_errors else None,
                    )

                return _parse_final_response(response_text)

            except Exception as e:
                last_error = e
                error_str = str(e).lower()
                attempt_errors.append(AttemptError(attempt=attempt + 1, error_message=str(e)))
                is_retryable = any(keyword in error_str for keyword in [
                    "timeout", "rate", "quota", "429", "503", "500",
                    "overloaded", "empty response", "none", "nonetype"
                ])
                if is_retryable and attempt < GEMINI_MAX_RETRIES:
                    delay = GEMINI_BASE_DELAY * (2 ** attempt) + random.uniform(0, 1)
                    logger.warning(
                        f"Gemini vision error (attempt {attempt + 1}/{GEMINI_MAX_RETRIES + 1}): {e}. "
                        f"Retrying in {delay:.1f}s..."
                    )
                    await asyncio.sleep(delay)
                else:
                    logger.error(f"Gemini vision generation failed after all retries: {e}")
                    FailureTracker.record_exhausted(
                        component="gemini_adapter.generate_with_vision",
                        error_message=str(e),
                        context={"model": self.model, "image_count": len(images), "attempts": attempt + 1},
                        attempt_errors=attempt_errors,
                    )
                    raise

        raise last_error

    @property
    def supports_vision(self) -> bool:
        """Check if this model supports vision inputs."""
        # Most Gemini models support vision
        return "gemini" in self.model.lower()
    
    @property
    def supports_structured_output(self) -> bool:
        """Check if this model supports native structured output."""
        # Gemini supports JSON mode with response schemas
        return True

