"""
Image utility functions for educational content evaluation.

This module provides utilities for extracting image URLs from content,
downloading images, encoding them for API use, and handling inline SVG content.
"""

import base64
import logging
import re
from typing import Dict, List, Optional

import requests

logger = logging.getLogger(__name__)


def _clean_image_url(url: str) -> str:
    """
    Clean an extracted image URL by removing any trailing description or garbage.
    
    Handles formats like:
    - 'https://example.com/image.png "description"' -> 'https://example.com/image.png'
    - 'https://example.com/image.png \\"description\\"' -> 'https://example.com/image.png'
    - 'https://example.com/image.png)' -> 'https://example.com/image.png'
    """
    if not url:
        return url
    
    # Strip leading/trailing whitespace
    url = url.strip()
    
    # Find where the actual URL ends - look for image extension followed by optional query string
    # Pattern: .ext or .ext?query but nothing after that except valid URL chars
    extension_match = re.search(
        r'(\.(?:jpg|jpeg|png|gif|bmp|svg|webp))(\?[^\s"\'\\]*)?',
        url, 
        re.IGNORECASE
    )
    
    if extension_match:
        # URL ends at the extension (plus any query string)
        end_pos = extension_match.end()
        url = url[:end_pos]
    
    # Also handle Supabase URLs that might not have extensions visible
    # Stop at common "end of URL" indicators
    for stop_char in [' "', ' \\"', ' \\\"', '\t', ' \'', '\\n', '\n']:
        if stop_char in url:
            url = url.split(stop_char)[0]
    
    # Remove trailing punctuation that's not part of URLs
    url = url.rstrip('"\'\\ \t\n,;:)>')
    
    # URL-encode spaces (common issue with S3 URLs that have spaces in filenames)
    url = url.replace(' ', '%20')
    
    return url


def extract_image_urls(content: str) -> List[str]:
    """
    Extract image URLs from content.
    
    Args:
        content: The content to extract URLs from
        
    Returns:
        List of unique image URLs found in the content
    """
    # Pattern to match common image URL formats
    # Matches URLs ending in common image extensions
    # Note: [^\s<>"\\] excludes backslash to handle escaped quotes like \"
    url_pattern = r'https?://[^\s<>"\\]+?\.(?:jpg|jpeg|png|gif|bmp|svg|webp)(?:\?[^\s<>"\\]*)?'
    
    # Also match Supabase storage URLs (common in this codebase)
    # Be more restrictive - stop at spaces, quotes, or backslashes
    supabase_pattern = r'https://[^\s<>"\\]*supabase[^\s<>"\\]*storage[^\s<>"\\]*\.(?:jpg|jpeg|png|gif|bmp|svg|webp)'
    
    urls = []
    urls.extend(re.findall(url_pattern, content, re.IGNORECASE))
    urls.extend(re.findall(supabase_pattern, content, re.IGNORECASE))

    # Extract markdown image syntax: ![alt](url)
    markdown_image_matches = re.findall(r'!\[.*?\]\((.*?)\)', content)
    urls.extend(markdown_image_matches)

    # Extract HTML image tags
    html_image_matches = re.findall(r'<img[^>]+src=["\']([^"\']+)["\']', content, re.IGNORECASE)
    urls.extend(html_image_matches)

    # Clean all extracted URLs and remove any that are now invalid
    cleaned_urls = []
    for url in urls:
        cleaned = _clean_image_url(url)
        if cleaned and cleaned.startswith('http'):
            cleaned_urls.append(cleaned)
        else:
            logger.debug(f"Discarded invalid URL after cleaning: {url} -> {cleaned}")

    # Remove duplicates while preserving order
    seen = set()
    unique_urls = []
    for url in cleaned_urls:
        if url not in seen:
            seen.add(url)
            unique_urls.append(url)
    
    logger.debug(f"Extracted {len(unique_urls)} unique image URLs from content")
    return unique_urls


def download_and_encode_image(image_url: str, timeout: int = 30) -> str:
    """
    Download an image and encode it as a base64 data URL.
    
    This helps avoid timeout issues when OpenAI's servers try to download
    external URLs directly. Falls back to the original URL if download fails.
    
    Args:
        image_url: The URL of the image to download and encode
        timeout: Request timeout in seconds (default: 30)
        
    Returns:
        Either a base64-encoded data URL (data:image/...;base64,...) or
        the original URL if download failed
    """
    # Clean the URL to remove any trailing description text
    cleaned_url = _clean_image_url(image_url)
    if cleaned_url != image_url:
        logger.debug(f"Cleaned image URL before download: {image_url} -> {cleaned_url}")
    
    try:
        response = requests.get(cleaned_url, timeout=timeout)
        response.raise_for_status()
        image_bytes = response.content
        
        if image_bytes:
            image_base64 = base64.b64encode(image_bytes).decode("utf-8")
            content_type = response.headers.get('content-type', 'image/png')
            
            # Determine the appropriate MIME type
            if 'jpeg' in content_type or 'jpg' in content_type:
                return f"data:image/jpeg;base64,{image_base64}"
            elif 'png' in content_type:
                return f"data:image/png;base64,{image_base64}"
            elif 'gif' in content_type:
                return f"data:image/gif;base64,{image_base64}"
            elif 'webp' in content_type:
                return f"data:image/webp;base64,{image_base64}"
            else:
                # Default to PNG for unknown types
                return f"data:image/png;base64,{image_base64}"
        else:
            logger.warning(
                f"Empty response when downloading {cleaned_url}, using direct URL"
            )
            return cleaned_url
            
    except Exception as e:
        logger.warning(
            f"Could not download image {cleaned_url}, using direct URL: {e}"
        )
        return cleaned_url


def prepare_images_for_api(image_urls: List[str]) -> List[dict]:
    """
    Prepare a list of image URLs for API consumption.
    
    Downloads and encodes each image, then returns them in the format
    expected by the OpenAI Vision API.
    
    Args:
        image_urls: List of image URLs to prepare
        
    Returns:
        List of dictionaries in the format:
        [{"type": "input_image", "image_url": "data:image/...;base64,..."}]
    """
    prepared_images = []
    
    for image_url in image_urls:
        encoded_url = download_and_encode_image(image_url)
        prepared_images.append({
            "type": "input_image",
            "image_url": encoded_url
        })
        
        if encoded_url.startswith("data:"):
            logger.info(f"Prepared image (base64 encoded): {image_url}")
        else:
            logger.info(f"Prepared image (direct URL): {image_url}")
    
    return prepared_images


def build_llm_content_with_images(text: str, image_urls: List[str], encode_base64: bool = False) -> List[dict]:
    """
    Build LLM message content array with text and images for vision-capable models.

    Args:
        text: The text content to include
        image_urls: List of image URLs to include
        encode_base64: If True, download and encode images as base64

    Returns:
        List of content blocks in OpenAI message format
    """
    content = [{"type": "text", "text": text}]

    for url in image_urls:
        # Clean the URL to remove any trailing description text
        cleaned_url = _clean_image_url(url)
        
        if encode_base64:
            encoded_url = download_and_encode_image(cleaned_url)
        else:
            encoded_url = cleaned_url
        content.append({
            "type": "image_url",
            "image_url": {"url": encoded_url}
        })

    return content


# =============================================================================
# SVG Extraction and Conversion
# =============================================================================

def extract_svg_content(content: str) -> List[Dict[str, any]]:
    """
    Extract inline SVG XML content from text.
    
    Handles both regular text and JSON-stringified content where quotes may be escaped.
    Looks for SVG tags with their complete content, typically marked with
    [SVG Image] or similar markers in the educational content.
    
    Args:
        content: The content to extract SVGs from (may be JSON-stringified)
        
    Returns:
        List of dictionaries containing SVG data:
        [{"index": 0, "content": "<svg>...</svg>", "marker": "[SVG Image]"}]
    """
    # First, try to parse as JSON to get unescaped content
    # This handles the case where content is json.dumps() output
    try:
        import json
        parsed = json.loads(content)
        # If it's a dict, recursively search all string values for SVGs
        if isinstance(parsed, dict):
            # Flatten all string values from the dict
            content_strings = []
            def collect_strings(obj):
                if isinstance(obj, str):
                    content_strings.append(obj)
                elif isinstance(obj, dict):
                    for value in obj.values():
                        collect_strings(value)
                elif isinstance(obj, list):
                    for item in obj:
                        collect_strings(item)
            
            collect_strings(parsed)
            # Process each string for SVGs
            all_svgs = []
            for text in content_strings:
                all_svgs.extend(_extract_svg_from_text(text))
            
            # Re-index to ensure sequential indices
            for i, svg in enumerate(all_svgs):
                svg['index'] = i
            
            return all_svgs
    except (json.JSONDecodeError, TypeError):
        # Not JSON, treat as regular text
        pass
    
    # Regular text extraction
    return _extract_svg_from_text(content)


def _extract_svg_from_text(text: str) -> List[Dict[str, any]]:
    """
    Helper function to extract SVGs from a plain text string.
    
    Handles nested SVG tags by counting opening/closing tags to find
    the complete outermost SVG element.
    
    Args:
        text: Plain text (not JSON-escaped) to extract SVGs from
        
    Returns:
        List of SVG dictionaries
    """
    svgs = []
    i = 0
    
    while i < len(text):
        # Find next opening <svg tag (case insensitive)
        match = re.search(r'<svg[\s>]', text[i:], re.IGNORECASE)
        if not match:
            break
        
        start_pos = i + match.start()
        
        # Find the end of the opening tag
        tag_end = text.find('>', start_pos)
        if tag_end == -1:
            break
        
        # Count nested svg tags to find the matching closing tag
        depth = 1
        pos = tag_end + 1
        
        while pos < len(text) and depth > 0:
            # Look for next svg tag (opening or closing)
            next_open = re.search(r'<svg[\s>]', text[pos:], re.IGNORECASE)
            next_close = re.search(r'</svg\s*>', text[pos:], re.IGNORECASE)
            
            if next_close is None:
                # No closing tag found
                break
            
            close_pos = pos + next_close.start()
            
            # Check if there's an opening tag before the closing tag
            if next_open and (pos + next_open.start()) < close_pos:
                # Found nested opening tag first
                depth += 1
                pos = pos + next_open.start() + 4  # Move past '<svg'
            else:
                # Found closing tag
                depth -= 1
                if depth == 0:
                    # Found matching closing tag
                    end_pos = close_pos + len(next_close.group(0))
                    svg_content = text[start_pos:end_pos]
                    
                    # Try to find any marker text before the SVG (like [SVG Image])
                    marker_pattern = r'\[SVG[^\]]*\]'
                    context_before = text[max(0, start_pos - 50):start_pos]
                    marker_match = re.search(marker_pattern, context_before, re.IGNORECASE)
                    marker = marker_match.group(0) if marker_match else None
                    
                    svgs.append({
                        "index": len(svgs),
                        "content": svg_content,
                        "marker": marker,
                        "position": start_pos
                    })
                    
                    i = end_pos
                    break
                else:
                    pos = close_pos + len(next_close.group(0))
        
        if depth > 0:
            # Couldn't find matching closing tag, skip this svg
            i = start_pos + 1
        
    if svgs:
        logger.info(f"Extracted {len(svgs)} inline SVG(s) from text")
    else:
        logger.debug("No inline SVGs found in text")
    
    return svgs


def convert_svg_to_png(svg_content: str, width: int = 800, height: int = 600) -> Optional[bytes]:
    """
    Convert SVG XML string to PNG image bytes.
    
    Tries multiple conversion methods in order:
    1. cairosvg (requires system cairo libraries)
    2. svglib + reportlab (pure Python, but also needs cairo)
    
    Args:
        svg_content: SVG XML string (complete <svg>...</svg> tag)
        width: Output PNG width in pixels (default: 800)
        height: Output PNG height in pixels (default: 600)
        
    Returns:
        PNG image as bytes, or None if conversion fails
    """
    # Method 1: Try cairosvg first (if available)
    try:
        import cairosvg
        
        png_bytes = cairosvg.svg2png(
            bytestring=svg_content.encode('utf-8'),
            output_width=width,
            output_height=height
        )
        logger.debug(f"Successfully converted SVG to PNG using cairosvg ({len(png_bytes)} bytes)")
        return png_bytes
        
    except (ImportError, OSError) as e:
        logger.debug(f"cairosvg not available: {e}")
    except Exception as e:
        logger.error(f"cairosvg conversion failed: {e}")
        logger.error(f"SVG content length: {len(svg_content)}")
        logger.error(f"SVG first 200 chars: {repr(svg_content[:200])}")
    
    # Method 2: Try svglib + reportlab
    try:
        from svglib.svglib import svg2rlg
        from reportlab.graphics import renderPM
        from io import BytesIO, StringIO
        
        # Convert SVG string to drawing
        drawing = svg2rlg(StringIO(svg_content))
        
        if drawing:
            # Scale to desired dimensions
            scale_x = width / drawing.width if drawing.width > 0 else 1
            scale_y = height / drawing.height if drawing.height > 0 else 1
            scale = min(scale_x, scale_y)
            
            drawing.width = drawing.width * scale
            drawing.height = drawing.height * scale
            drawing.scale(scale, scale)
            
            # Render to PNG
            png_data = BytesIO()
            renderPM.drawToFile(drawing, png_data, fmt='PNG')
            png_bytes = png_data.getvalue()
            
            logger.debug(f"Successfully converted SVG to PNG using svglib ({len(png_bytes)} bytes)")
            return png_bytes
            
    except ImportError:
        logger.debug("svglib/reportlab not available")
    except Exception as e:
        logger.debug(f"svglib conversion failed: {e}")
    
    # No conversion method available
    logger.warning(
        "No SVG conversion library available. SVG images will be skipped.\n"
        "To enable SVG support, you need cairo system libraries:\n"
        "  macOS: brew install cairo pango gdk-pixbuf libffi\n"
        "  Then: pip uninstall cairosvg cairocffi && pip install cairosvg\n"
        "\nFor now, only external image URLs will be analyzed."
    )
    return None


def prepare_svgs_for_analysis(svgs: List[Dict[str, any]], width: int = 800, height: int = 600) -> List[Dict[str, any]]:
    """
    Convert SVG content to base64-encoded PNG images for analysis.
    
    Args:
        svgs: List of SVG dictionaries from extract_svg_content()
        width: PNG width in pixels (default: 800)
        height: PNG height in pixels (default: 600)
        
    Returns:
        List of dictionaries in format:
        [{
            "type": "input_image",
            "image_url": "data:image/png;base64,...",
            "source": "svg",
            "original_index": 0,
            "marker": "[SVG Image]"
        }]
    """
    prepared = []
    
    for svg_item in svgs:
        try:
            png_bytes = convert_svg_to_png(svg_item['content'], width, height)
            
            if png_bytes:
                png_base64 = base64.b64encode(png_bytes).decode('utf-8')
                prepared.append({
                    "type": "input_image",
                    "image_url": f"data:image/png;base64,{png_base64}",
                    "source": "svg",
                    "original_index": svg_item['index'],
                    "marker": svg_item.get('marker'),
                    "position": svg_item.get('position')
                })
                logger.info(f"Prepared SVG {svg_item['index']} for analysis")
            else:
                logger.warning(f"Skipping SVG {svg_item['index']}: conversion returned None")
                
        except Exception as e:
            logger.warning(f"Failed to prepare SVG {svg_item['index']}: {e}")
            continue
    
    return prepared


def extract_all_images(content: str) -> Dict[str, any]:
    """
    Extract both image URLs and inline SVGs from content.
    
    This is the main entry point for comprehensive image extraction that
    handles both external image URLs and inline SVG content.
    
    Args:
        content: The content to extract images from
        
    Returns:
        Dictionary with:
        {
            'urls': List[str],              # Regular image URLs
            'svgs': List[Dict],             # SVG data from extract_svg_content()
            'svg_images': List[Dict]        # Prepared SVG images (base64 PNG)
        }
    """
    result = {
        'urls': extract_image_urls(content),
        'svgs': extract_svg_content(content),
        'svg_images': []
    }
    
    # Convert SVGs to images if any were found
    if result['svgs']:
        result['svg_images'] = prepare_svgs_for_analysis(result['svgs'])
    
    total_images = len(result['urls']) + len(result['svg_images'])
    logger.info(
        f"Extracted {total_images} total image(s): "
        f"{len(result['urls'])} URL(s), {len(result['svg_images'])} SVG(s)"
    )
    
    return result


def get_all_image_urls_for_analysis(content: str) -> List[str]:
    """
    Get all image URLs ready for analysis, including converted SVGs.
    
    This is a convenience function that returns a flat list of all image URLs
    (both external URLs and base64-encoded converted SVGs) ready to pass to
    the unified image analyzer.
    
    Args:
        content: The content to extract images from
        
    Returns:
        List of image URLs (external URLs + base64 data URLs from SVGs)
    """
    image_data = extract_all_images(content)
    
    # Start with external URLs
    all_urls = list(image_data['urls'])
    
    # Add converted SVG data URLs
    for svg_image in image_data['svg_images']:
        all_urls.append(svg_image['image_url'])
    
    return all_urls


def remove_svg_from_content(content: str) -> str:
    """
    Remove inline SVG XML code from content, leaving only text markers.
    
    Since SVGs are separately analyzed as converted images, the raw SVG code
    is redundant and consumes large amounts of tokens. This function removes
    the SVG XML while preserving any surrounding text markers like [SVG Image].
    
    Handles nested SVG tags by properly tracking depth and removing complete
    outermost SVG blocks.
    
    Args:
        content: The content potentially containing inline SVG tags
        
    Returns:
        Content with SVG XML code removed
    """
    result = []
    i = 0
    
    while i < len(content):
        # Find next opening <svg tag (case insensitive)
        match = re.search(r'<svg[\s>]', content[i:], re.IGNORECASE)
        if not match:
            # No more SVGs, append rest of content
            result.append(content[i:])
            break
        
        start_pos = i + match.start()
        
        # Append content before this SVG
        result.append(content[i:start_pos])
        
        # Find the end of the opening tag
        tag_end = content.find('>', start_pos)
        if tag_end == -1:
            # Malformed SVG, skip it
            result.append(content[start_pos:start_pos + 4])  # Just append '<svg'
            i = start_pos + 4
            continue
        
        # Count nested svg tags to find the matching closing tag
        depth = 1
        pos = tag_end + 1
        
        while pos < len(content) and depth > 0:
            # Look for next svg tag (opening or closing)
            next_open = re.search(r'<svg[\s>]', content[pos:], re.IGNORECASE)
            next_close = re.search(r'</svg\s*>', content[pos:], re.IGNORECASE)
            
            if next_close is None:
                # No closing tag found, malformed SVG
                break
            
            close_pos = pos + next_close.start()
            
            # Check if there's an opening tag before the closing tag
            if next_open and (pos + next_open.start()) < close_pos:
                # Found nested opening tag first
                depth += 1
                pos = pos + next_open.start() + 4  # Move past '<svg'
            else:
                # Found closing tag
                depth -= 1
                if depth == 0:
                    # Found matching closing tag for the outermost SVG
                    end_pos = close_pos + len(next_close.group(0))
                    
                    # Replace entire SVG block with placeholder
                    result.append('[SVG converted to image for analysis]')
                    
                    i = end_pos
                    break
                else:
                    pos = close_pos + len(next_close.group(0))
        
        if depth > 0:
            # Couldn't find matching closing tag, treat as malformed
            # Just append the opening tag and continue
            result.append(content[start_pos:start_pos + 4])
            i = start_pos + 4
    
    cleaned = ''.join(result)
    logger.debug(f"Removed SVG code from content (original: {len(content)} chars, cleaned: {len(cleaned)} chars)")
    
    return cleaned
