"""
Unified image analysis using Gemini Flash with agentic vision (code execution).

This module provides comprehensive image analysis using only LLM capabilities:
- Object counting with code execution for precision
- Geometric shape analysis and classification
- Visual property detection and measurement
- Spatial relationship understanding

Replaces the previous OpenCV + LLM hybrid approach with a pure LLM solution
that leverages Gemini's code execution capabilities for accurate analysis.
"""

import asyncio
import base64
import logging
import random
import time
from collections import OrderedDict
from typing import Dict, List, Optional

import requests
from pydantic import BaseModel, Field

from inceptbench_new.llm import LLMFactory, LLMImage, LLMMessage
from inceptbench_new.tools.image_utils import _clean_image_url
from inceptbench_new.utils.failure_tracker import AttemptError, FailureTracker

logger = logging.getLogger(__name__)

# Retry configuration for image downloads
IMAGE_DOWNLOAD_MAX_RETRIES = 3
IMAGE_DOWNLOAD_BASE_DELAY = 4.0

# LRU cache configuration
MAX_CACHE_SIZE = 1000  # Maximum number of images to cache

# Module-level LRU cache for image analysis results
# Uses OrderedDict for LRU eviction - most recently used items are at the end
# Key: image URL (cleaned), Value: ImageAnalysisData
_image_analysis_cache: OrderedDict[str, "ImageAnalysisData"] = OrderedDict()


def _cache_get(key: str) -> Optional["ImageAnalysisData"]:
    """
    Get item from cache, moving it to end (most recently used).
    
    Returns:
        Cached ImageAnalysisData or None if not found
    """
    if key in _image_analysis_cache:
        # Move to end (most recently used)
        _image_analysis_cache.move_to_end(key)
        return _image_analysis_cache[key]
    return None


def _cache_set(key: str, value: "ImageAnalysisData") -> None:
    """
    Set item in cache, evicting oldest if at capacity.
    """
    if key in _image_analysis_cache:
        # Update existing and move to end
        _image_analysis_cache.move_to_end(key)
        _image_analysis_cache[key] = value
    else:
        # Add new item
        _image_analysis_cache[key] = value
        # Evict oldest if over capacity
        while len(_image_analysis_cache) > MAX_CACHE_SIZE:
            evicted_key, _ = _image_analysis_cache.popitem(last=False)
            logger.debug(f"Cache evicted oldest entry: {evicted_key[:60]}...")


def get_cache_stats() -> Dict[str, int]:
    """
    Get statistics about the image analysis cache.
    
    Returns:
        Dictionary with cache statistics (size, max_size, etc.)
    """
    return {
        "cached_images": len(_image_analysis_cache),
        "max_cache_size": MAX_CACHE_SIZE,
    }


def clear_cache() -> int:
    """
    Clear the image analysis cache.
    
    Returns:
        Number of entries that were cleared
    """
    global _image_analysis_cache
    count = len(_image_analysis_cache)
    _image_analysis_cache = OrderedDict()
    logger.info(f"Cleared image analysis cache ({count} entries)")
    return count


# =============================================================================
# System Prompt for Unified Image Analysis
# =============================================================================

SYSTEM_PROMPT = """You are an expert at comprehensive image analysis for educational content.
You have access to code execution capabilities that allow you to programmatically analyze images
with precision.

## Your Capabilities

You can analyze images to:
1. **Count objects** - Count any type of object (shapes, animals, items, etc.)
2. **Classify shapes** - Identify geometric shapes and their properties
3. **Measure properties** - Analyze angles, sizes, spatial relationships
4. **Detect patterns** - Identify arrangements, groupings, and structures
5. **Extract data** - Read charts, graphs, tables, and visual data

## Critical Analysis Approach

### STEP 1: UNDERSTAND THE IMAGE TYPE
- Is this a realistic scene, chart/graph, geometric diagram, or educational illustration?
- What is the primary educational purpose?
- Are there multiple distinct elements to analyze?

### STEP 2: USE CODE EXECUTION FOR PRECISION
When analyzing images, you should:
- Use code execution to load and analyze the image programmatically
- Apply computer vision techniques (edge detection, contour analysis, etc.) when needed
- Count objects systematically using code rather than visual estimation
- Measure geometric properties (angles, side lengths, areas) computationally
- Verify your visual observations with programmatic analysis

### STEP 3: PROVIDE COMPREHENSIVE ANALYSIS
For each image, provide:
1. **Overall description** - What the image shows
2. **Object counts** - Precise counts for all distinct object types
3. **Shape classifications** - Geometric shape identification with supporting evidence
4. **Measurements** - Angles, dimensions, ratios (when relevant)
5. **Spatial relationships** - Arrangements, groupings, patterns
6. **Confidence assessment** - How certain you are about your analysis

## Counting Guidelines

### PARTIAL OBJECTS
- **Charts/Graphs/Pictographs**: Partial symbols are meaningful (count as 0.5)
  Example: A pictograph showing "2.5 apples" uses half-symbols intentionally
- **Realistic Scenes**: Partial objects usually represent whole objects that are occluded
  Example: An apple half-hidden behind another counts as 1.0, not 0.5
- Use ONLY whole numbers or 0.5 increments (never 0.3, 0.7, etc.)

### GROUPED OBJECTS
When objects are in groups (rows, containers, sections):
- Count each group separately with clear identifiers
- Provide a breakdown: [{"group_identifier": "Row 1", "count": 3}, ...]
- Verify that sum of groups equals total count

### VISIBILITY RULES
- Charts/graphs: Count all visible symbols including partial ones
- Realistic scenes: Only count objects >80% visible as whole objects
- Objects cut off at image edges typically don't count (unless in charts)

## Shape Classification Rules

### TRIANGLES (3 sides)
- **Equilateral**: All sides equal (ratio > 0.95), all angles ~60°
- **Isosceles**: Exactly 2 sides equal
- **Scalene**: No sides equal
- Also classify by angles: **Acute** (all < 90°), **Right** (one = 90°), **Obtuse** (one > 90°)

### QUADRILATERALS (4 sides)
- **Square**: All 4 sides equal AND all 4 angles = 90°
- **Rectangle**: Opposite sides equal AND all 4 angles = 90°
- **Rhombus**: All 4 sides equal, angles NOT 90°
- **Parallelogram**: Opposite sides equal, opposite angles equal, angles NOT 90°
- **Trapezoid**: Exactly ONE pair of parallel sides
- **Kite**: Two pairs of ADJACENT sides equal (not opposite)
- **Irregular**: None of the above

### POLYGONS (5+ sides)
- **Regular**: All sides and angles equal
- **Irregular**: Varying sides or angles
- **Star**: Concave polygon with alternating acute/obtuse angles

## Output Requirements

Be thorough but concise. For each element analyzed:
1. Provide precise counts/measurements (use code execution)
2. Explain your classification reasoning
3. Note any ambiguities or limitations
4. Assess confidence level (HIGH/MEDIUM/LOW)

## Quality Control

- **Sanity check**: Do your counts/measurements make sense?
- **Consistency**: Do different analysis methods agree?
- **Educational suitability**: Could a student clearly understand this content?
- **Confidence**: Be honest about uncertainty

Remember: Use code execution to ensure accuracy. Don't just visually estimate - 
programmatically analyze the image whenever possible."""


ANALYSIS_USER_PROMPT = """Please analyze this image comprehensively for educational content evaluation.

Provide:

1. **Overall Description**: What does the image show? What is its educational purpose?

2. **Object Counts**: Count all distinct object types present. Use code execution for precision.
   - For grouped objects, list each group with its count
   - Use whole numbers or 0.5 increments only
   - Distinguish between chart symbols (0.5 for partial) and real objects (1.0 for occluded)

3. **Shape Classifications**: For any geometric shapes present:
   - Classify each shape (e.g., "square", "isosceles triangle", "rhombus")
   - Provide supporting measurements (angles, side ratios)
   - Explain your reasoning

4. **Visual Properties**: Note any relevant:
   - Measurements (angles, dimensions, areas)
   - Spatial arrangements
   - Patterns or structures
   - Text or labels

5. **Confidence Assessment**: Rate your overall confidence (HIGH/MEDIUM/LOW) and explain why.

Use code execution to analyze the image programmatically whenever possible for maximum accuracy.

Provide a clear, structured analysis."""


# =============================================================================
# Pydantic Models for Structured Output
# =============================================================================

class GroupCount(BaseModel):
    """Count data for a specific group within grouped objects."""
    group_identifier: str = Field(description="Clear name for this group (e.g., 'Row 1', 'Blue section')")
    count: float = Field(description="Count for this group (whole numbers or 0.5 increments only)")
    
    class Config:
        extra = "forbid"


class ObjectCount(BaseModel):
    """Count data for a specific object type in the image."""
    object_type: str = Field(description="Type of object (e.g., 'apples', 'triangles', 'stars')")
    total_count: float = Field(description="Total count (whole numbers or 0.5 increments only)")
    counting_method: str = Field(description="How the count was determined (e.g., 'code execution', 'visual grid scan')")
    confidence: str = Field(description="Confidence level: HIGH, MEDIUM, or LOW")
    group_breakdown: List[GroupCount] = Field(
        default_factory=list,
        description="Per-group counts if objects are grouped (rows, sections, etc.)"
    )
    notes: str = Field(
        default="",
        description="Any relevant notes about counting (partial objects, occlusion, etc.)"
    )
    
    class Config:
        extra = "forbid"  # Explicitly forbid additional properties


class ShapeClassification(BaseModel):
    """Classification and measurements for a geometric shape."""
    shape_type: str = Field(description="Shape classification (e.g., 'square', 'isosceles right triangle', 'rhombus')")
    confidence: str = Field(description="Confidence level: HIGH, MEDIUM, or LOW")
    supporting_evidence: List[str] = Field(description="Key measurements or properties supporting this classification")
    measurements: str = Field(
        default="",
        description="Specific measurements as text (e.g., 'side_length=100px, area=10000px^2')"
    )
    
    class Config:
        extra = "forbid"


class VisualProperty(BaseModel):
    """A visual property or measurement from the image."""
    property_name: str = Field(description="Name of the property (e.g., 'angle_ABC', 'grid_dimensions')")
    value: str = Field(description="Value of the property (with units if applicable)")
    confidence: str = Field(description="Confidence level: HIGH, MEDIUM, or LOW")
    
    class Config:
        extra = "forbid"


class ImageAnalysisData(BaseModel):
    """Complete analysis data for a single image."""
    image_url: str
    image_index: int
    
    # Overall description
    overall_description: str = Field(description="General description of what the image shows")
    image_type: str = Field(description="Type of image: realistic_scene, chart, geometric_diagram, educational_illustration, or mixed")
    educational_purpose: str = Field(description="Inferred educational purpose or learning objective")
    
    # Object counts
    object_counts: List[ObjectCount] = Field(
        default_factory=list,
        description="Counts for each distinct object type found"
    )
    
    # Shape classifications
    shape_classifications: List[ShapeClassification] = Field(
        default_factory=list,
        description="Classifications for geometric shapes present"
    )
    
    # Visual properties and measurements
    visual_properties: List[VisualProperty] = Field(
        default_factory=list,
        description="Other relevant measurements or properties"
    )
    
    # Quality assessment
    overall_confidence: str = Field(description="Overall confidence: HIGH, MEDIUM, or LOW")
    quality_notes: str = Field(description="Assessment of image quality and suitability for educational use")
    limitations: str = Field(
        default="",
        description="Any limitations or ambiguities in the analysis"
    )
    
    class Config:
        extra = "forbid"


class ImageAnalysisToolResult(BaseModel):
    """Tool response format for image analysis."""
    image_analyses: List[ImageAnalysisData]
    
    class Config:
        extra = "forbid"


class ImageAnalysisResult(BaseModel):
    """Result from analyzing one or more images."""
    success: bool
    image_analyses: List[ImageAnalysisData]
    error_message: str = ""
    
    class Config:
        extra = "forbid"


# =============================================================================
# Per-Image Analysis Helper
# =============================================================================

async def _analyze_single_image(
    llm,
    image: LLMImage,
    image_url: str,
    image_index: int,
) -> ImageAnalysisData:
    """
    Analyze a single image using Gemini Flash with code execution.
    
    Each image gets its own LLM call so the analysis is specific to that image,
    avoiding the duplication problem of sharing a single response across images.
    
    Parameters
    ----------
    llm : LLMInterface
        The LLM instance to use for analysis
    image : LLMImage
        The prepared image (base64 or URL)
    image_url : str
        The cleaned image URL (used for logging and as the stored URL)
    image_index : int
        1-based index of this image in the analysis batch
        
    Returns
    -------
    ImageAnalysisData
        Analysis data specific to this single image
    """
    response_text = await llm.generate_with_vision(
        messages=[
            LLMMessage(role="system", content=SYSTEM_PROMPT),
            LLMMessage(role="user", content=ANALYSIS_USER_PROMPT),
        ],
        images=[image],
        response_schema=None,  # Text mode — Gemini code-execution requires it
        enable_code_execution=True,
    )
    
    return ImageAnalysisData(
        image_url=image_url,
        image_index=image_index,
        overall_description=response_text[:500] if response_text else "Analysis completed",
        image_type="mixed",
        educational_purpose="Educational content analysis",
        object_counts=[],
        shape_classifications=[],
        visual_properties=[],
        overall_confidence="MEDIUM",
        quality_notes=response_text if response_text else "Analysis completed with code execution",
        limitations="",
    )


# =============================================================================
# Main Analysis Function
# =============================================================================

async def analyze_images(image_urls: List[str]) -> ImageAnalysisResult:
    """
    Analyze images using Gemini Flash with code execution (agentic vision).
    
    This unified approach replaces the previous OpenCV + LLM hybrid with a pure
    LLM solution that leverages code execution for precision.
    
    Parameters
    ----------
    image_urls : List[str]
        List of image URLs to analyze
        
    Returns
    -------
    ImageAnalysisResult
        Comprehensive analysis results including counts, classifications, and measurements
    """
    if not image_urls:
        return ImageAnalysisResult(
            success=True,
            image_analyses=[],
            error_message=""
        )
    
    start_time = time.time()
    
    # Clean all URLs first
    cleaned_urls = []
    for image_url in image_urls:
        cleaned_url = _clean_image_url(image_url)
        if cleaned_url != image_url:
            logger.debug(f"Cleaned image URL: {image_url} -> {cleaned_url}")
        cleaned_urls.append(cleaned_url)
    
    # Check cache for each image
    cached_results: Dict[str, ImageAnalysisData] = {}
    uncached_urls: List[str] = []
    uncached_indices: List[int] = []
    
    for i, cleaned_url in enumerate(cleaned_urls):
        cached_result = _cache_get(cleaned_url)
        if cached_result is not None:
            cached_results[cleaned_url] = cached_result
            logger.debug(f"Cache hit for image: {cleaned_url[:60]}...")
        else:
            uncached_urls.append(cleaned_url)
            uncached_indices.append(i)
    
    cache_hits = len(cached_results)
    cache_misses = len(uncached_urls)
    
    if cache_hits > 0:
        logger.info(f"Image analysis cache: {cache_hits} hit(s), {cache_misses} miss(es)")
    
    # If all images are cached, return cached results
    if not uncached_urls:
        logger.info(f"All {cache_hits} image(s) found in cache, skipping LLM call")
        # Build result list in original order
        all_results = []
        for i, cleaned_url in enumerate(cleaned_urls):
            result = cached_results[cleaned_url]
            # Update image_index to match current request order
            result_copy = ImageAnalysisData(
                image_url=result.image_url,
                image_index=i + 1,
                overall_description=result.overall_description,
                image_type=result.image_type,
                educational_purpose=result.educational_purpose,
                object_counts=result.object_counts,
                shape_classifications=result.shape_classifications,
                visual_properties=result.visual_properties,
                overall_confidence=result.overall_confidence,
                quality_notes=result.quality_notes,
                limitations=result.limitations
            )
            all_results.append(result_copy)
        
        return ImageAnalysisResult(
            success=True,
            image_analyses=all_results,
            error_message=""
        )
    
    # Process uncached images
    try:
        logger.info(f"Analyzing {len(uncached_urls)} uncached image(s) with Gemini Flash (agentic vision)...")
        
        llm = LLMFactory.create("unified_image_analyzer")
        
        # Download and prepare uncached images with retry logic
        images = []
        for i, cleaned_url in enumerate(uncached_urls):
            image_bytes = None
            content_type = 'image/png'
            download_success = False
            attempt_errors: List[AttemptError] = []
            
            for attempt in range(IMAGE_DOWNLOAD_MAX_RETRIES + 1):
                try:
                    await asyncio.sleep(0.2)  # Small delay between requests
                    image_response = requests.get(cleaned_url, timeout=30)
                    image_response.raise_for_status()
                    image_bytes = image_response.content
                    content_type = image_response.headers.get('content-type', 'image/png')
                    download_success = True
                    
                    # Record recovery if this wasn't the first attempt
                    if attempt > 0:
                        FailureTracker.record_recovered(
                            component="unified_image_analyzer.download",
                            message=f"Succeeded on attempt {attempt + 1}/{IMAGE_DOWNLOAD_MAX_RETRIES + 1}",
                            context={"image_index": i + 1, "url": cleaned_url[:100]},
                            attempt_errors=attempt_errors if attempt_errors else None
                        )
                        logger.info(f"Image {i+1} download succeeded on attempt {attempt + 1}")
                    break
                    
                except Exception as e:
                    attempt_errors.append(AttemptError(
                        attempt=attempt + 1,
                        error_message=str(e)
                    ))
                    
                    if attempt < IMAGE_DOWNLOAD_MAX_RETRIES:
                        delay = IMAGE_DOWNLOAD_BASE_DELAY * (2 ** attempt) + random.uniform(0, 1)
                        logger.warning(
                            f"Image {i+1} download failed (attempt {attempt + 1}/"
                            f"{IMAGE_DOWNLOAD_MAX_RETRIES + 1}): {e}. Retrying in {delay:.1f}s..."
                        )
                        await asyncio.sleep(delay)
                    else:
                        # All retries exhausted
                        logger.error(f"Image {i+1} download failed after all retries: {e}")
                        FailureTracker.record_exhausted(
                            component="unified_image_analyzer.download",
                            error_message=str(e),
                            context={"image_index": i + 1, "url": cleaned_url[:100], "attempts": attempt + 1},
                            attempt_errors=attempt_errors
                        )
            
            if download_success and image_bytes:
                image_base64 = base64.b64encode(image_bytes).decode("utf-8")
                media_type = "image/jpeg" if 'jpeg' in content_type or 'jpg' in content_type \
                    else "image/png"
                
                images.append(LLMImage(
                    base64_data=image_base64,
                    media_type=media_type
                ))
            else:
                # Download failed - use URL fallback
                logger.warning(f"Image {i+1}: Using URL fallback after download failure")
                images.append(LLMImage(url=cleaned_url))
        
        # Analyze each image individually so every image gets its own
        # dedicated LLM response (avoids duplicating one response across all).
        # Calls are parallelised with asyncio.gather for throughput.
        analysis_tasks = [
            _analyze_single_image(
                llm=llm,
                image=images[i],
                image_url=uncached_urls[i],
                image_index=i + 1,
            )
            for i in range(len(uncached_urls))
        ]
        image_analyses: List[ImageAnalysisData] = list(
            await asyncio.gather(*analysis_tasks)
        )
        
        logger.info(
            f"Received analysis responses for {len(image_analyses)} image(s)"
        )
        
        tool_result = ImageAnalysisToolResult(image_analyses=image_analyses)
        
        logger.info(
            f"Successfully analyzed {len(tool_result.image_analyses)} image(s) "
            f"in {time.time() - start_time:.2f}s"
        )
        
        # Cache the new results (with LRU eviction if at capacity)
        for result in tool_result.image_analyses:
            # Use the cleaned URL as cache key
            idx = result.image_index - 1  # Convert to 0-based
            if idx < len(uncached_urls):
                cache_key = uncached_urls[idx]
                _cache_set(cache_key, result)
                logger.debug(f"Cached analysis for: {cache_key[:60]}...")
        
        # Merge cached and new results in original order
        new_results_by_url = {
            uncached_urls[r.image_index - 1]: r for r in tool_result.image_analyses
        }
        
        all_results = []
        for i, cleaned_url in enumerate(cleaned_urls):
            if cleaned_url in cached_results:
                result = cached_results[cleaned_url]
            elif cleaned_url in new_results_by_url:
                result = new_results_by_url[cleaned_url]
            else:
                # Shouldn't happen, but handle gracefully
                logger.warning(f"Missing result for image: {cleaned_url[:60]}...")
                continue
            
            # Update image_index to match current request order
            result_copy = ImageAnalysisData(
                image_url=result.image_url,
                image_index=i + 1,
                overall_description=result.overall_description,
                image_type=result.image_type,
                educational_purpose=result.educational_purpose,
                object_counts=result.object_counts,
                shape_classifications=result.shape_classifications,
                visual_properties=result.visual_properties,
                overall_confidence=result.overall_confidence,
                quality_notes=result.quality_notes,
                limitations=result.limitations
            )
            all_results.append(result_copy)
        
        return ImageAnalysisResult(
            success=True,
            image_analyses=all_results,
            error_message=""
        )
        
    except Exception as e:
        logger.error(f"Error analyzing images after {time.time() - start_time:.2f}s: {e}")
        return ImageAnalysisResult(
            success=False,
            image_analyses=[],
            error_message=str(e)
        )


# =============================================================================
# Formatting Functions for Prompt Inclusion
# =============================================================================

def format_analysis_for_prompt(analysis_result: ImageAnalysisResult) -> str:
    """
    Format analysis results into a readable string for inclusion in evaluation prompts.
    
    Parameters
    ----------
    analysis_result : ImageAnalysisResult
        The analysis results to format
        
    Returns
    -------
    str
        Formatted string describing the analysis
    """
    if not analysis_result.success:
        return f"\n\nImage Analysis Failed: {analysis_result.error_message}"
    
    if not analysis_result.image_analyses:
        return ""
    
    output_parts = ["\n\n## COMPREHENSIVE IMAGE ANALYSIS (Gemini Flash + Code Execution)"]
    output_parts.append("\nProgrammatic image analysis with LLM reasoning and code execution.")
    output_parts.append("AUTHORITATIVE FOR: Object counts, shape classifications, measurements, spatial relationships.")
    output_parts.append("Use this data as ground truth for visual content verification.\n")
    
    for img_data in analysis_result.image_analyses:
        if len(analysis_result.image_analyses) > 1:
            output_parts.append(f"\n### Image {img_data.image_index}")
        
        output_parts.append(f"**Description**: {img_data.overall_description}")
        output_parts.append(f"**Type**: {img_data.image_type}")
        output_parts.append(f"**Educational Purpose**: {img_data.educational_purpose}")
        output_parts.append(f"**Overall Confidence**: {img_data.overall_confidence}")
        
        # Object counts
        if img_data.object_counts:
            output_parts.append("\n**Object Counts**:")
            for obj_count in img_data.object_counts:
                output_parts.append(f"- **{obj_count.object_type}**: {obj_count.total_count}")
                output_parts.append(f"  - Method: {obj_count.counting_method}")
                output_parts.append(f"  - Confidence: {obj_count.confidence}")
                
                if obj_count.group_breakdown:
                    output_parts.append("  - Group Breakdown:")
                    for group in obj_count.group_breakdown:
                        output_parts.append(f"    * {group.group_identifier}: {group.count}")
                
                if obj_count.notes:
                    output_parts.append(f"  - Notes: {obj_count.notes}")
        
        # Shape classifications
        if img_data.shape_classifications:
            output_parts.append("\n**Shape Classifications**:")
            for shape in img_data.shape_classifications:
                output_parts.append(f"- **{shape.shape_type}** (Confidence: {shape.confidence})")
                output_parts.append(f"  - Evidence: {', '.join(shape.supporting_evidence)}")
                
                if shape.measurements:
                    output_parts.append(f"  - Measurements: {shape.measurements}")
        
        # Visual properties
        if img_data.visual_properties:
            output_parts.append("\n**Visual Properties**:")
            for prop in img_data.visual_properties:
                output_parts.append(f"- {prop.property_name}: {prop.value} (Confidence: {prop.confidence})")
        
        # Quality assessment
        output_parts.append(f"\n**Quality Assessment**: {img_data.quality_notes}")
        if img_data.limitations:
            output_parts.append(f"**Limitations**: {img_data.limitations}")
    
    return "\n".join(output_parts)
