src/utils.py

"""
Utility functions for image processing and validation.
"""

import base64
import io
import os
from pathlib import Path

from PIL import Image


SUPPORTED_FORMATS = {".png", ".jpg", ".jpeg", ".gif", ".webp", ".bmp"}
MAX_IMAGE_SIZE_MB = 10
MAX_DIMENSION = 4096  # Max width/height for API submission


def validate_image(image_bytes: bytes, filename: str = "image.png") -> tuple[bool, str]:
    """Validate an image file for chart analysis.

    Returns:
        (is_valid, error_message)
    """
    ext = Path(filename).suffix.lower()
    if ext not in SUPPORTED_FORMATS:
        return False, f"Unsupported format: {ext}. Use: {', '.join(SUPPORTED_FORMATS)}"

    size_mb = len(image_bytes) / (1024 * 1024)
    if size_mb > MAX_IMAGE_SIZE_MB:
        return False, f"Image too large: {size_mb:.1f}MB (max {MAX_IMAGE_SIZE_MB}MB)"

    try:
        img = Image.open(io.BytesIO(image_bytes))
        img.verify()
    except Exception as e:
        return False, f"Invalid image file: {e}"

    return True, ""


def prepare_image_for_api(
    image_bytes: bytes,
    max_dimension: int = MAX_DIMENSION,
) -> tuple[bytes, str]:
    """Resize image if needed and return (processed_bytes, media_type).

    Ensures the image fits within API size limits while preserving quality.
    """
    img = Image.open(io.BytesIO(image_bytes))

    # Convert RGBA to RGB if needed (for JPEG output)
    if img.mode == "RGBA":
        background = Image.new("RGB", img.size, (255, 255, 255))
        background.paste(img, mask=img.split()[3])
        img = background
    elif img.mode != "RGB":
        img = img.convert("RGB")

    # Resize if too large
    w, h = img.size
    if w > max_dimension or h > max_dimension:
        ratio = min(max_dimension / w, max_dimension / h)
        new_size = (int(w * ratio), int(h * ratio))
        img = img.resize(new_size, Image.LANCZOS)

    # Save as PNG for best quality
    buffer = io.BytesIO()
    img.save(buffer, format="PNG", optimize=True)
    return buffer.getvalue(), "image/png"


def image_to_base64(image_bytes: bytes) -> str:
    """Encode image bytes to base64 string."""
    return base64.standard_b64encode(image_bytes).decode("utf-8")


def load_image_from_path(path: str) -> bytes:
    """Load image bytes from a file path."""
    with open(path, "rb") as f:
        return f.read()


def get_image_info(image_bytes: bytes) -> dict:
    """Get basic image metadata."""
    img = Image.open(io.BytesIO(image_bytes))
    return {
        "width": img.size[0],
        "height": img.size[1],
        "format": img.format or "unknown",
        "mode": img.mode,
        "size_kb": len(image_bytes) / 1024,
    }


def detect_ticker_from_image(image_bytes: bytes, provider: str, model: str, api_key: str) -> str:
    """Quick vision call to extract the stock ticker from a chart screenshot.

    Returns the ticker string (e.g., 'NVDA', 'AAPL', 'D05.SI') or empty string.
    """
    processed_bytes, media_type = prepare_image_for_api(image_bytes)
    b64 = image_to_base64(processed_bytes)

    prompt = (
        "Look at this stock chart image. What is the stock ticker symbol shown? "
        "Reply with ONLY the ticker symbol (e.g., NVDA, AAPL, D05.SI). "
        "If you cannot determine it, reply with UNKNOWN."
    )

    try:
        if provider == "google":
            from google import genai
            from google.genai import types
            client = genai.Client(api_key=api_key)
            image_part = types.Part.from_bytes(data=processed_bytes, mime_type=media_type)
            response = client.models.generate_content(model=model, contents=[image_part, prompt])
            result = response.text.strip()
        elif provider in ("huggingface", "openai", "openrouter"):
            import openai
            if provider == "huggingface":
                base_url = "https://router.huggingface.co/v1"
            elif provider == "openrouter":
                base_url = "https://openrouter.ai/api/v1"
            else:
                base_url = None
            client = openai.OpenAI(base_url=base_url, api_key=api_key) if base_url else openai.OpenAI(api_key=api_key)
            response = client.chat.completions.create(
                model=model,
                max_tokens=20,
                messages=[{
                    "role": "user",
                    "content": [
                        {"type": "image_url", "image_url": {"url": f"data:{media_type};base64,{b64}"}},
                        {"type": "text", "text": prompt},
                    ],
                }],
            )
            result = response.choices[0].message.content.strip()
        elif provider == "anthropic":
            import anthropic
            client = anthropic.Anthropic(api_key=api_key)
            response = client.messages.create(
                model=model, max_tokens=20,
                messages=[{"role": "user", "content": [
                    {"type": "image", "source": {"type": "base64", "media_type": media_type, "data": b64}},
                    {"type": "text", "text": prompt},
                ]}],
            )
            result = response.content[0].text.strip()
        else:
            return ""

        # Clean up — extract just the ticker from the response
        result = result.replace("*", "").replace("`", "").strip()
        if result and result != "UNKNOWN" and len(result) <= 15:
            return result
    except Exception:
        pass

    return ""