image_processor.py

"""
Image processing utilities for the Visual Memory Synthesizer.

Handles EXIF extraction, filename date parsing, image resizing,
and storage management.
"""

import base64
import io
import re
import shutil
from datetime import datetime
from pathlib import Path
from typing import Optional, Dict, Tuple

try:
    from PIL import Image
    from PIL.ExifTags import TAGS, GPSTAGS
    PIL_AVAILABLE = True
except ImportError:
    PIL_AVAILABLE = False

try:
    import pillow_heif
    pillow_heif.register_heif_opener()
    HEIF_AVAILABLE = True
except ImportError:
    HEIF_AVAILABLE = False


# Storage directory
DATA_DIR = Path(__file__).parent / "data"
IMAGES_DIR = DATA_DIR / "images"

# Supported image types
SUPPORTED_EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp", ".heic", ".heif"}


def ensure_images_dir() -> Path:
    """Ensure the images directory exists."""
    IMAGES_DIR.mkdir(parents=True, exist_ok=True)
    return IMAGES_DIR


def is_supported_image(path: Path) -> bool:
    """Check if the file is a supported image type."""
    return path.suffix.lower() in SUPPORTED_EXTENSIONS


def extract_exif(image_path: Path) -> Dict:
    """
    Extract EXIF metadata from an image.

    Returns:
        Dict with keys:
        - timestamp: datetime or None
        - latitude: float or None
        - longitude: float or None
        - width: int
        - height: int
        - camera: str or None
    """
    if not PIL_AVAILABLE:
        return {"error": "Pillow not installed"}

    result = {
        "timestamp": None,
        "latitude": None,
        "longitude": None,
        "width": None,
        "height": None,
        "camera": None
    }

    try:
        with Image.open(image_path) as img:
            result["width"] = img.width
            result["height"] = img.height

            # Get EXIF data
            exif_data = img._getexif()
            if not exif_data:
                return result

            # Parse EXIF tags
            exif = {}
            for tag_id, value in exif_data.items():
                tag = TAGS.get(tag_id, tag_id)
                exif[tag] = value

            # Extract timestamp
            for date_tag in ["DateTimeOriginal", "DateTimeDigitized", "DateTime"]:
                if date_tag in exif:
                    try:
                        dt_str = exif[date_tag]
                        result["timestamp"] = datetime.strptime(
                            dt_str, "%Y:%m:%d %H:%M:%S"
                        )
                        break
                    except (ValueError, TypeError):
                        continue

            # Extract camera info
            if "Make" in exif and "Model" in exif:
                result["camera"] = f"{exif['Make']} {exif['Model']}".strip()
            elif "Model" in exif:
                result["camera"] = exif["Model"]

            # Extract GPS coordinates
            if "GPSInfo" in exif:
                gps_info = exif["GPSInfo"]
                gps = {}
                for key, value in gps_info.items():
                    gps[GPSTAGS.get(key, key)] = value

                if "GPSLatitude" in gps and "GPSLatitudeRef" in gps:
                    lat = _convert_gps_coords(gps["GPSLatitude"])
                    if gps["GPSLatitudeRef"] == "S":
                        lat = -lat
                    result["latitude"] = lat

                if "GPSLongitude" in gps and "GPSLongitudeRef" in gps:
                    lon = _convert_gps_coords(gps["GPSLongitude"])
                    if gps["GPSLongitudeRef"] == "W":
                        lon = -lon
                    result["longitude"] = lon

    except Exception as e:
        result["error"] = str(e)

    return result


def _convert_gps_coords(coords) -> float:
    """Convert GPS coordinates from EXIF format to decimal degrees."""
    try:
        degrees = float(coords[0])
        minutes = float(coords[1])
        seconds = float(coords[2])
        return degrees + (minutes / 60.0) + (seconds / 3600.0)
    except (IndexError, TypeError, ValueError):
        return 0.0


def parse_filename_date(filename: str) -> Tuple[Optional[datetime], str]:
    """
    Attempt to parse a date from the filename.

    Common patterns:
    - IMG_YYYYMMDD_HHMMSS
    - YYYYMMDD_HHMMSS
    - YYYY-MM-DD_HH-MM-SS
    - Photo on YYYY-MM-DD
    - Screenshot YYYY-MM-DD

    Returns:
        Tuple of (datetime or None, confidence: "exact"|"month"|"fuzzy")
    """
    name = Path(filename).stem

    patterns = [
        # IMG_20241215_143022
        (r"IMG[_-]?(\d{4})(\d{2})(\d{2})[_-]?(\d{2})(\d{2})(\d{2})", "exact"),
        # 20241215_143022
        (r"^(\d{4})(\d{2})(\d{2})[_-](\d{2})(\d{2})(\d{2})", "exact"),
        # 2024-12-15_14-30-22
        (r"(\d{4})-(\d{2})-(\d{2})[_-](\d{2})-(\d{2})-(\d{2})", "exact"),
        # 2024-12-15 or 20241215
        (r"(\d{4})-?(\d{2})-?(\d{2})", "exact"),
        # Photo 12-15-24 (MM-DD-YY format)
        (r"(\d{1,2})-(\d{1,2})-(\d{2,4})", "month"),
    ]

    for pattern, confidence in patterns:
        match = re.search(pattern, name)
        if match:
            groups = match.groups()
            try:
                if len(groups) >= 6:
                    # Full datetime
                    year = int(groups[0])
                    month = int(groups[1])
                    day = int(groups[2])
                    hour = int(groups[3])
                    minute = int(groups[4])
                    second = int(groups[5])
                    return datetime(year, month, day, hour, minute, second), confidence
                elif len(groups) == 3:
                    year = int(groups[0])
                    month = int(groups[1])
                    day = int(groups[2])

                    # Handle 2-digit years
                    if year < 100:
                        year += 2000

                    # Validate reasonable range
                    if 1990 <= year <= 2100 and 1 <= month <= 12 and 1 <= day <= 31:
                        return datetime(year, month, day), confidence
            except (ValueError, IndexError):
                continue

    return None, "fuzzy"


def prepare_for_vision(
    image_path: Path,
    max_size: int = 2048,
    quality: int = 85
) -> str:
    """
    Prepare an image for vision model analysis.

    Resizes if necessary and returns base64 encoded string.

    Args:
        image_path: Path to the image file
        max_size: Maximum dimension (width or height) in pixels
        quality: JPEG quality for compression (1-100)

    Returns:
        Base64 encoded image string
    """
    if not PIL_AVAILABLE:
        # Fallback: just read and encode
        with open(image_path, "rb") as f:
            return base64.b64encode(f.read()).decode("utf-8")

    with Image.open(image_path) as img:
        # Convert to RGB if necessary (handles RGBA, P mode, etc.)
        if img.mode in ("RGBA", "P"):
            img = img.convert("RGB")
        elif img.mode != "RGB":
            img = img.convert("RGB")

        # Resize if larger than max_size
        if img.width > max_size or img.height > max_size:
            ratio = min(max_size / img.width, max_size / img.height)
            new_size = (int(img.width * ratio), int(img.height * ratio))
            img = img.resize(new_size, Image.LANCZOS)

        # Encode to JPEG in memory
        buffer = io.BytesIO()
        img.save(buffer, format="JPEG", quality=quality)
        buffer.seek(0)

        return base64.b64encode(buffer.read()).decode("utf-8")


def copy_to_storage(
    source_path: Path,
    preserve_original: bool = True
) -> Path:
    """
    Copy an image to the storage directory.

    Uses timestamp-based naming to avoid conflicts.

    Args:
        source_path: Path to the source image
        preserve_original: If False, moves instead of copying

    Returns:
        Path to the stored image
    """
    ensure_images_dir()

    # Generate unique filename
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    original_name = source_path.stem
    extension = source_path.suffix.lower()

    # Sanitize original name
    safe_name = re.sub(r"[^\w\-]", "_", original_name)[:50]

    new_filename = f"{timestamp}_{safe_name}{extension}"
    dest_path = IMAGES_DIR / new_filename

    # Handle name collision
    counter = 1
    while dest_path.exists():
        new_filename = f"{timestamp}_{safe_name}_{counter}{extension}"
        dest_path = IMAGES_DIR / new_filename
        counter += 1

    if preserve_original:
        shutil.copy2(source_path, dest_path)
    else:
        shutil.move(source_path, dest_path)

    return dest_path


def get_image_dimensions(image_path: Path) -> Tuple[int, int]:
    """Get image dimensions without loading full image into memory."""
    if not PIL_AVAILABLE:
        return (0, 0)

    try:
        with Image.open(image_path) as img:
            return img.size
    except Exception:
        return (0, 0)


def generate_thumbnail(
    image_path: Path,
    size: Tuple[int, int] = (200, 200)
) -> Optional[str]:
    """
    Generate a thumbnail as base64 for preview.

    Args:
        image_path: Path to the image
        size: Thumbnail dimensions (width, height)

    Returns:
        Base64 encoded thumbnail or None on failure
    """
    if not PIL_AVAILABLE:
        return None

    try:
        with Image.open(image_path) as img:
            if img.mode in ("RGBA", "P"):
                img = img.convert("RGB")

            img.thumbnail(size, Image.LANCZOS)

            buffer = io.BytesIO()
            img.save(buffer, format="JPEG", quality=70)
            buffer.seek(0)

            return base64.b64encode(buffer.read()).decode("utf-8")
    except Exception:
        return None


def get_processing_info() -> Dict:
    """
    Get information about available processing capabilities.

    Returns:
        Dict with capability flags and supported formats
    """
    return {
        "pil_available": PIL_AVAILABLE,
        "heif_available": HEIF_AVAILABLE,
        "supported_extensions": list(SUPPORTED_EXTENSIONS),
        "images_dir": str(IMAGES_DIR),
        "max_recommended_size": 2048
    }