validator.py

"""LLM-based image validation."""
from __future__ import annotations

import os
import json
import base64
from io import BytesIO
from pathlib import Path
from datetime import datetime
from openai import OpenAI
import boto3

import config
import storage


def get_openai_client():
    api_key = os.environ.get("OPENAI_API_KEY") or config.OPENAI_API_KEY
    if not api_key:
        raise ValueError("OPENAI_API_KEY not set")
    return OpenAI(api_key=api_key)


_s3_client = None

def get_s3_client():
    """Get or create S3 client."""
    global _s3_client
    if _s3_client is None:
        _s3_client = boto3.client("s3", region_name=config.S3_REGION)
    return _s3_client


def encode_image_from_s3(s3_path: str) -> str:
    """Load image from S3 directly into memory and encode to base64 (no disk I/O)."""
    # Extract key from s3://bucket/key
    if s3_path.startswith("s3://"):
        parts = s3_path[5:].split("/", 1)
        s3_key = parts[1]
    else:
        s3_key = s3_path

    client = get_s3_client()
    response = client.get_object(Bucket=config.S3_BUCKET, Key=s3_key)
    image_bytes = response["Body"].read()
    return base64.b64encode(image_bytes).decode("utf-8")


def encode_image_from_local(filepath: str) -> str:
    """Encode local image to base64."""
    with open(filepath, "rb") as f:
        return base64.b64encode(f.read()).decode("utf-8")


def validate_images_batch(images: list, category: str, batch_size: int = 5, model: str = None) -> list:
    """Validate multiple images in a single LLM call."""
    prompt = config.VALIDATION_PROMPTS.get(category)
    if not prompt:
        raise ValueError(f"No validation prompt for category: {category}")

    # Use provided model or fall back to config
    model = model or getattr(config, "OPENAI_MODEL", "gpt-4o-mini")

    # Encode all images
    encoded_images = []
    valid_indices = []
    for i, img in enumerate(images):
        image_path = img.get("s3_processed_path", "")
        try:
            if image_path.startswith("s3://") or image_path.startswith("processed/") or image_path.startswith("raw/"):
                b64 = encode_image_from_s3(image_path)
            else:
                b64 = encode_image_from_local(image_path)
            encoded_images.append((i, img, b64))
            valid_indices.append(i)
        except Exception as e:
            print(f"✗ {img['id']}: Failed to load image - {e}")

    if not encoded_images:
        return [{"valid": False, "reason": "Failed to load image"} for _ in images]

    # Build batch prompt - use full prompt, just change the response format
    batch_prompt = f"""You will analyze {len(encoded_images)} images.

{prompt.split("Respond with JSON")[0]}

Respond with a JSON object containing a "results" array with one object per image, in the same order as provided:
{{
    "results": [
        {{"image": 1, "valid": true/false, "reason": "brief reason if invalid, null if valid"}},
        {{"image": 2, "valid": true/false, "reason": "..."}},
        ...
    ]
}}"""

    # Build content array with all images
    content = [{"type": "text", "text": batch_prompt}]
    for idx, (i, img, b64) in enumerate(encoded_images):
        content.append({"type": "text", "text": f"Image {idx + 1} ({img['id']}):"})
        content.append({
            "type": "image_url",
            "image_url": {"url": f"data:image/jpeg;base64,{b64}", "detail": "low"},
        })

    client = get_openai_client()

    try:
        if os.environ.get("DEBUG"):
            print(f"  [DEBUG] Calling {model} with {len(encoded_images)} images...")
        response = client.chat.completions.create(
            model=model,
            messages=[{"role": "user", "content": content}],
            max_completion_tokens=100 * len(encoded_images),
            response_format={"type": "json_object"},
        )
        result_text = response.choices[0].message.content
        if os.environ.get("DEBUG"):
            print(f"  [DEBUG] Raw response: {result_text[:500] if result_text else 'EMPTY'}")
        # Parse the result - handle both array and object with "results" key
        parsed = json.loads(result_text)
        if isinstance(parsed, dict) and "results" in parsed:
            results = parsed["results"]
        elif isinstance(parsed, list):
            results = parsed
        else:
            results = [parsed]
    except Exception as e:
        print(f"✗ Batch LLM error: {e}")
        if os.environ.get("DEBUG"):
            import traceback
            traceback.print_exc()
        return [{"valid": False, "reason": f"LLM error: {e}"} for _ in images]

    # Map results back to original images
    final_results = [{"valid": False, "reason": "Failed to load image"} for _ in images]
    for idx, (orig_idx, img, _) in enumerate(encoded_images):
        if idx < len(results):
            result = results[idx]
            final_results[orig_idx] = result
            status = "✓" if result.get("valid") else "✗"
            reason = result.get("reason", "")
            print(f"{status} {img['id']}: {reason if reason else 'Valid'}")

    return final_results


def validate_image(image_data: dict, category: str = None, model: str = None) -> dict:
    """Validate a single image using LLM.

    Args:
        image_data: Image dict with s3_processed_path or local path
        category: Validation prompt key (defaults to image's category)
        model: Model to use (defaults to config.OPENAI_MODEL)
    """
    category = category or image_data.get("category")
    prompt = config.VALIDATION_PROMPTS.get(category)

    if not prompt:
        raise ValueError(f"No validation prompt for category: {category}")

    # Get image path (S3 or local)
    image_path = image_data.get("s3_processed_path", "")

    # Encode image (direct memory load from S3, no temp files)
    try:
        if image_path.startswith("s3://") or image_path.startswith("processed/") or image_path.startswith("raw/"):
            image_b64 = encode_image_from_s3(image_path)
        else:
            image_b64 = encode_image_from_local(image_path)
    except Exception as e:
        print(f"✗ {image_data['id']}: Failed to load image - {e}")
        return {"valid": False, "reason": f"Failed to load image: {e}"}

    client = get_openai_client()
    model = model or getattr(config, "OPENAI_MODEL", "gpt-4o-mini")

    try:
        response = client.chat.completions.create(
            model=model,
            messages=[
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": prompt},
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/jpeg;base64,{image_b64}",
                                "detail": "low",
                            },
                        },
                    ],
                }
            ],
            max_completion_tokens=200,
            response_format={"type": "json_object"},
        )

        result_text = response.choices[0].message.content
        if os.environ.get("DEBUG"):
            print(f"  [DEBUG] Raw response: {result_text[:500] if result_text else 'EMPTY'}")
        result = json.loads(result_text)

    except Exception as e:
        print(f"✗ {image_data['id']}: LLM error - {e}")
        if os.environ.get("DEBUG"):
            import traceback
            traceback.print_exc()
        return {"valid": False, "reason": f"LLM error: {e}"}

    # Print result
    status = "✓" if result.get("valid") else "✗"
    reason = result.get("reason", "")
    print(f"{status} {image_data['id']}: {reason if reason else 'Valid'}")

    return result


def validate_s3_image(s3_key: str, category: str) -> dict:
    """Validate a single image directly from S3 key (no disk I/O)."""
    prompt = config.VALIDATION_PROMPTS.get(category)
    if not prompt:
        raise ValueError(f"No validation prompt for category: {category}")

    client = get_openai_client()
    image_b64 = encode_image_from_s3(s3_key)

    # Determine image type from key
    ext = Path(s3_key).suffix.lower()
    media_type = "image/jpeg" if ext in [".jpg", ".jpeg"] else "image/png"
    model = getattr(config, "OPENAI_MODEL", "gpt-4o-mini")

    try:
        response = client.chat.completions.create(
            model=model,
            messages=[
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": prompt},
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:{media_type};base64,{image_b64}",
                                "detail": "low",
                            },
                        },
                    ],
                }
            ],
            max_completion_tokens=500,
            response_format={"type": "json_object"},
        )

        result_text = response.choices[0].message.content
        result = json.loads(result_text)

    except Exception as e:
        print(f"Validation failed for {s3_key}: {e}")
        return {"s3_key": s3_key, "valid": False, "error": str(e)}

    result["s3_key"] = s3_key
    return result


def validate_s3_batch(category: str, source: str = "autodev", limit: int = 10) -> list[dict]:
    """Validate images from S3 bucket directly (no disk I/O)."""
    s3 = get_s3_client()
    prefix = f"processed/{category}/{source}/"

    # List objects in S3
    response = s3.list_objects_v2(Bucket=config.S3_BUCKET, Prefix=prefix, MaxKeys=limit)
    keys = [obj["Key"] for obj in response.get("Contents", [])]

    if not keys:
        print(f"No images found at s3://{config.S3_BUCKET}/{prefix}")
        return []

    # Build lookup from DB for original URLs
    db_images = storage.get_images(source, category)
    url_lookup = {img["id"]: img.get("original_url") for img in db_images}

    total = len(keys)
    print(f"Validating {total} images from S3...\n")
    results = []
    valid_count = 0

    for i, s3_key in enumerate(keys, 1):
        result = validate_s3_image(s3_key, category)
        results.append(result)

        is_valid = result.get("valid")
        if is_valid:
            valid_count += 1

        # Get image ID from s3_key (e.g., "processed/.../autodev_ABC123_0.jpg" -> "autodev_ABC123_0")
        image_id = Path(s3_key).stem
        original_url = url_lookup.get(image_id, "")

        # Print each result
        status = "✓" if is_valid else "✗"
        reason = result.get("reason") or result.get("rejection_reason") or result.get("error") or ""
        print(f"[{i}/{total}] {status} {image_id}")
        if original_url:
            print(f"         {original_url}")
        if reason:
            print(f"         {reason}")

    print(f"\n=== Validation Summary ===")
    print(f"Valid: {valid_count} | Invalid: {len(results) - valid_count}")

    return results


def validate_batch(
    source: str,
    category: str,
    limit: int = None,
    batch_size: int = 1,
    status: str = "processed",
    prompt_key: str = None,
    model: str = None,
) -> tuple:
    """Validate unvalidated images from DB.

    Args:
        source: Image source (e.g., 'pixabay', 'autodev', 'amazon_reviews')
        category: Category to filter images by
        limit: Max images to process
        batch_size: Number of images per LLM call (1 = serial, >1 = batched)
        status: Image status to filter by (default "processed")
        prompt_key: Validation prompt key (defaults to category if not specified)
        model: Model to use for validation (defaults to config.OPENAI_MODEL)
    """
    images = storage.get_images(source=source, category=category, status=status, with_s3_path=True)

    if not images:
        print(f"No images found for {category} from {source} with status={status}")
        return [], []

    # Filter to unvalidated only
    images = [img for img in images if not img.get("validated_at")]

    if not images:
        print(f"All images already validated for {category} from {source}")
        return [], []

    if limit:
        images = images[:limit]

    # Determine which prompt to use
    validation_prompt_key = prompt_key or category
    print(f"\n=== Validating {len(images)} images ({category}/{source}) ===")
    print(f"Prompt: {validation_prompt_key}")
    print(f"Model: {model or config.OPENAI_MODEL}")
    print(f"Batch size: {batch_size}\n")

    valid = []
    invalid = []

    # Process in batches
    for i in range(0, len(images), batch_size):
        batch = images[i:i + batch_size]

        if batch_size == 1:
            # Single image mode
            results = [validate_image(batch[0], validation_prompt_key, model=model)]
        else:
            # Batch mode
            results = validate_images_batch(batch, validation_prompt_key, batch_size, model=model)

        # Update DB for each result
        for img, result in zip(batch, results):
            new_status = "validated" if result.get("valid") else "rejected"
            storage.update_image_status(
                img["id"],
                new_status,
                valid=result.get("valid", False),
                validation_result=result,
                rejection_reason=result.get("reason"),
                validated_at=datetime.now().isoformat(),
            )

            if result.get("valid"):
                valid.append(img)
            else:
                invalid.append(img)

    print(f"\n=== Summary ===")
    print(f"Valid: {len(valid)} | Invalid: {len(invalid)}")

    return valid, invalid


def show_stats(source: str):
    """Show validation stats."""
    stats = storage.get_stats(source)
    print(f"\n=== {source} stats ===")
    for status, count in stats.items():
        print(f"  {status}: {count}")


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Validate images using LLM")
    parser.add_argument("source", help="Image source (e.g., 'amazon_reviews', 'pixabay') or 's3' for direct S3 mode")
    parser.add_argument("category", help="Category to filter by (for DB) or S3 prefix category")
    parser.add_argument("--limit", "-l", type=int, help="Max images to validate")
    parser.add_argument("--batch-size", "-b", type=int, default=1, help="Images per LLM call (default: 1)")
    parser.add_argument("--status", "-s", default="processed", help="Image status to filter (default: processed)")
    parser.add_argument("--prompt", "-p", help="Validation prompt key (defaults to category)")
    parser.add_argument("--model", "-m", help="Model to use (default: config.OPENAI_MODEL)")

    args = parser.parse_args()

    if args.source == "s3":
        # Direct S3 mode (legacy)
        results = validate_s3_batch(args.category, "autodev", args.limit or 10)
        print("\n=== Full Results ===")
        print(json.dumps(results, indent=2))
    else:
        validate_batch(
            source=args.source,
            category=args.category,
            limit=args.limit,
            batch_size=args.batch_size,
            status=args.status,
            prompt_key=args.prompt,
            model=args.model,
        )