nanobot/channels/telegram_media.py

@@ -2,11 +2,13 @@
 
 from __future__ import annotations
 
+import io
 import mimetypes
 from enum import Enum
 from pathlib import Path
 
 from loguru import logger
+from PIL import Image
 
 try:
     import magic
@@ -14,6 +16,13 @@ try:
 except ImportError:
     HAS_MAGIC = False
 
+try:
+    from pillow_heif import register_heif_opener
+    register_heif_opener()
+    HAS_HEIF = True
+except ImportError:
+    HAS_HEIF = False
+
 
 class MediaKind(Enum):
     """Media type classification."""
@@ -91,3 +100,108 @@ def is_heic_format(path: str) -> bool:
     """
     ext = Path(path).suffix.lower()
     return ext in (".heic", ".heif")
+
+
+async def optimize_image(path: str, max_bytes: int = 6_000_000) -> bytes:
+    """
+    Optimize image to fit under size limit.
+
+    Strategy:
+    1. Convert HEIC to JPEG if needed
+    2. PNG with alpha → preserve with compression levels [6,7,8,9]
+    3. JPEG/PNG without alpha → resize + quality grid
+
+    Sizes: [2048, 1536, 1280, 1024, 800] px (max dimension)
+    Qualities: [80, 70, 60, 50, 40] (JPEG only)
+
+    Args:
+        path: Path to image file
+        max_bytes: Maximum size in bytes (default 6MB for Telegram)
+
+    Returns:
+        Optimized image bytes
+
+    Raises:
+        ValueError: If image cannot be optimized under limit
+    """
+    # Load image
+    img = Image.open(path)
+
+    # Convert HEIC to JPEG
+    if is_heic_format(path):
+        if not HAS_HEIF:
+            raise ValueError("pillow-heif not available for HEIC conversion")
+        # Convert to RGB (HEIC → JPEG)
+        if img.mode != "RGB":
+            img = img.convert("RGB")
+        return await _optimize_jpeg(img, max_bytes)
+
+    # PNG with alpha channel - preserve it
+    if img.mode == "RGBA" or img.mode == "LA":
+        return await _optimize_png(img, max_bytes)
+
+    # Everything else → convert to JPEG and optimize
+    if img.mode != "RGB":
+        img = img.convert("RGB")
+    return await _optimize_jpeg(img, max_bytes)
+
+
+async def _optimize_jpeg(img: Image.Image, max_bytes: int) -> bytes:
+    """Optimize JPEG with size/quality grid."""
+    sizes = [2048, 1536, 1280, 1024, 800]
+    qualities = [80, 70, 60, 50, 40]
+
+    original_width, original_height = img.size
+
+    for size in sizes:
+        # Skip if image is already smaller
+        if max(original_width, original_height) <= size:
+            resized = img
+        else:
+            # Resize maintaining aspect ratio
+            resized = img.copy()
+            resized.thumbnail((size, size), Image.Resampling.LANCZOS)
+
+        for quality in qualities:
+            buf = io.BytesIO()
+            resized.save(buf, format="JPEG", quality=quality, optimize=True)
+            data = buf.getvalue()
+
+            if len(data) <= max_bytes:
+                return data
+
+    # If we get here, even smallest size/quality is too large
+    raise ValueError(f"Cannot optimize image under {max_bytes} bytes")
+
+
+async def _optimize_png(img: Image.Image, max_bytes: int) -> bytes:
+    """Optimize PNG while preserving alpha channel."""
+    compress_levels = [6, 7, 8, 9]
+    sizes = [2048, 1536, 1280, 1024, 800]
+
+    original_width, original_height = img.size
+
+    for size in sizes:
+        # Skip if image is already smaller
+        if max(original_width, original_height) <= size:
+            resized = img
+        else:
+            resized = img.copy()
+            resized.thumbnail((size, size), Image.Resampling.LANCZOS)
+
+        for compress_level in compress_levels:
+            buf = io.BytesIO()
+            resized.save(buf, format="PNG", compress_level=compress_level, optimize=True)
+            data = buf.getvalue()
+
+            if len(data) <= max_bytes:
+                return data
+
+    # Fallback: try converting to JPEG if still too large
+    if img.mode == "RGBA":
+        # Create white background
+        background = Image.new("RGB", img.size, (255, 255, 255))
+        background.paste(img, mask=img.split()[3])  # Use alpha as mask
+        return await _optimize_jpeg(background, max_bytes)
+
+    raise ValueError(f"Cannot optimize PNG under {max_bytes} bytes")

tests/test_telegram_media.py

@@ -1,5 +1,10 @@
 """Tests for Telegram media handling."""
 
+import io
+from pathlib import Path
+
+import pytest
+from PIL import Image
 
 
 def test_detect_mime_from_jpeg():
@@ -109,3 +114,64 @@ def test_is_heic_format():
     assert is_heic_format("photo.HEIC") is True
     assert is_heic_format("photo.heif") is True
     assert is_heic_format("photo.jpg") is False
+
+
+@pytest.mark.asyncio
+async def test_optimize_image_jpeg_quality():
+    """Test JPEG optimization reduces size with quality ladder."""
+    from nanobot.channels.telegram_media import optimize_image
+
+    # Create a large test image (3000x3000 RGB)
+    img = Image.new("RGB", (3000, 3000), color="red")
+    buf = io.BytesIO()
+    img.save(buf, format="JPEG", quality=95)
+    original_bytes = buf.getvalue()
+    original_size = len(original_bytes)
+
+    # Write to temp file
+    import tempfile
+    with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f:
+        f.write(original_bytes)
+        temp_path = f.name
+
+    try:
+        # Optimize to 1MB max
+        optimized = await optimize_image(temp_path, max_bytes=1_000_000)
+
+        # Should be smaller than original
+        assert len(optimized) < original_size
+        # Should be under limit
+        assert len(optimized) <= 1_000_000
+        # Should still be valid JPEG
+        assert optimized.startswith(b'\xff\xd8\xff')
+    finally:
+        Path(temp_path).unlink()
+
+
+@pytest.mark.asyncio
+async def test_optimize_image_png_preserve_alpha():
+    """Test PNG with alpha channel is preserved."""
+    from nanobot.channels.telegram_media import optimize_image
+
+    # Create PNG with alpha channel
+    img = Image.new("RGBA", (1000, 1000), color=(255, 0, 0, 128))
+    buf = io.BytesIO()
+    img.save(buf, format="PNG")
+    original_bytes = buf.getvalue()
+
+    import tempfile
+    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f:
+        f.write(original_bytes)
+        temp_path = f.name
+
+    try:
+        optimized = await optimize_image(temp_path, max_bytes=5_000_000)
+
+        # Should still be PNG (PNG magic bytes)
+        assert optimized.startswith(b'\x89PNG')
+
+        # Load and verify alpha channel preserved
+        img_opt = Image.open(io.BytesIO(optimized))
+        assert img_opt.mode == "RGBA"
+    finally:
+        Path(temp_path).unlink()