from __future__ import annotations import asyncio import base64 import copy import datetime import mimetypes import os import re import xxhash from pathlib import Path from typing import TYPE_CHECKING, Any if TYPE_CHECKING: from ...core import app from . import mgr as storage_mgr DEFAULT_RETENTION_DAYS = 30 DEFAULT_MAX_SIZE_MB = 0 MEDIA_DIR = 'media_cache' SAFE_MEDIA_FILENAME = re.compile(r'^[a-f0-9]{32,64}(\.[a-zA-Z0-9]{1,10})?$') class MediaCache: """Content-addressable storage cache for images and media attachments. Deduplicates media files using xxHash3-128, offloads payloads from SQLite to StorageProvider, and implements LRU and age-based retention cleanup. """ def __init__(self, ap: app.Application, storage_mgr: storage_mgr.StorageMgr): self.ap = ap self.storage_mgr = storage_mgr @staticmethod def hash_bytes(data: bytes) -> str: """Compute content-addressable hash for binary data.""" return xxhash.xxh3_128_hexdigest(data) @staticmethod def parse_data_url(data_url: str) -> tuple[bytes, str] | None: """Parse a data URL or raw base64 string into bytes and mime type.""" if not data_url or not isinstance(data_url, str): return None try: if data_url.startswith('data:'): split_index = data_url.find(';base64,') if split_index != -1: mime_type = data_url[5:split_index] b64_data = data_url[split_index + 8 :] return base64.b64decode(b64_data), mime_type # Try raw base64 if sufficiently long if len(data_url) > 20 and not data_url.startswith(('http://', 'https://', '/')): return base64.b64decode(data_url), 'application/octet-stream' except Exception: return None return None @staticmethod def guess_extension(mime_type: str | None, default: str = '.jpg') -> str: """Guess appropriate file extension from MIME type.""" if not mime_type: return default mime_lower = mime_type.lower() if 'png' in mime_lower: return '.png' if 'webp' in mime_lower: return '.webp' if 'gif' in mime_lower: return '.gif' if 'jpeg' in mime_lower or 'jpg' in mime_lower: return '.jpg' ext = mimetypes.guess_extension(mime_type) if ext == '.jpe': return '.jpg' return ext or default async def save_media(self, data: bytes, mime_type: str | None = None) -> tuple[str, str, int]: """Save media bytes into content-addressable storage cache. Returns: Tuple of (hash_str, storage_key, byte_size) """ hash_str = self.hash_bytes(data) ext = self.guess_extension(mime_type) storage_key = f'{MEDIA_DIR}/{hash_str}{ext}' provider = self.storage_mgr.storage_provider if not await provider.exists(storage_key): await provider.save(storage_key, data) else: await self.touch(storage_key) return hash_str, storage_key, len(data) async def get_media(self, filename_or_key: str) -> tuple[bytes, str] | None: """Retrieve media bytes and mime type by key or filename.""" filename = os.path.basename(filename_or_key) if not SAFE_MEDIA_FILENAME.match(filename): return None storage_key = f'{MEDIA_DIR}/{filename}' provider = self.storage_mgr.storage_provider if not await provider.exists(storage_key): return None data = await self.storage_mgr._load_object_bounded(storage_key) mime_type = mimetypes.guess_type(filename)[0] or 'application/octet-stream' await self.touch(storage_key) return data, mime_type async def touch(self, storage_key: str) -> None: """Update access/modified time of a media file for LRU tracking.""" provider = getattr(self.storage_mgr, 'storage_provider', None) if provider is not None and provider.__class__.__name__ == 'LocalStorageProvider': full_path = os.path.join('data', 'storage', storage_key) if os.path.exists(full_path): now = datetime.datetime.now().timestamp() try: await asyncio.to_thread(os.utime, full_path, (now, now)) except Exception: # Failures to update mtime are intentionally ignored because LRU touch is opportunistic. pass async def cleanup( self, retention_days: int = DEFAULT_RETENTION_DAYS, max_size_mb: int = DEFAULT_MAX_SIZE_MB, ) -> dict[str, int]: """Perform age-based and LRU size-based cleanup on media cache. Args: retention_days: Retain media accessed within this many days (default 30). max_size_mb: Maximum total size in MB (0 means unlimited). Returns: Dictionary of cleanup metrics. """ provider = getattr(self.storage_mgr, 'storage_provider', None) if provider is None or provider.__class__.__name__ != 'LocalStorageProvider': return {'expired_deleted': 0, 'size_deleted': 0, 'bytes_freed': 0} target_dir = Path('data/storage') / MEDIA_DIR if not target_dir.exists() or not target_dir.is_dir(): return {'expired_deleted': 0, 'size_deleted': 0, 'bytes_freed': 0} now = datetime.datetime.now().timestamp() cutoff = (now - retention_days * 86400) if retention_days > 0 else 0 expired_deleted = 0 size_deleted = 0 bytes_freed = 0 remaining: list[tuple[Path, int, float]] = [] for entry in target_dir.iterdir(): if not entry.is_file(): continue try: stat = entry.stat() except OSError: continue if cutoff > 0 and stat.st_mtime < cutoff: try: entry.unlink(missing_ok=True) expired_deleted += 1 bytes_freed += stat.st_size except OSError: # Best-effort cleanup; file may already be gone or temporarily inaccessible. pass else: remaining.append((entry, stat.st_size, stat.st_mtime)) if max_size_mb > 0: max_bytes = max_size_mb * 1024 * 1024 total_bytes = sum(item[1] for item in remaining) if total_bytes > max_bytes: remaining.sort(key=lambda item: item[2]) for path, size, _ in remaining: if total_bytes <= max_bytes: break try: path.unlink(missing_ok=True) size_deleted += 1 bytes_freed += size total_bytes -= size except OSError: # Best-effort cleanup; file may already be gone or temporarily inaccessible. continue return { 'expired_deleted': expired_deleted, 'size_deleted': size_deleted, 'bytes_freed': bytes_freed, } async def externalize_chain_dump(self, chain_dump: Any) -> Any: """Recursively extract raw base64 media into cache and replace with references.""" if isinstance(chain_dump, list): return [await self.externalize_chain_dump(item) for item in chain_dump] if isinstance(chain_dump, dict): node = copy.copy(chain_dump) node_type = node.get('type') if node_type == 'Image': b64 = node.get('base64') if b64 and isinstance(b64, str): try: parsed = self.parse_data_url(b64) if parsed is not None: raw_bytes, mime_type = parsed hash_str, storage_key, size = await self.save_media(raw_bytes, mime_type) filename = os.path.basename(storage_key) current_url = node.get('url') or '' if current_url and not current_url.startswith('data:'): node['original_url'] = current_url node['url'] = f'/api/v1/files/media/{filename}' node['base64'] = None node['hash'] = hash_str node['storage_key'] = storage_key node['size'] = size node['mime_type'] = mime_type except Exception as e: if hasattr(self.ap, 'logger') and self.ap.logger: self.ap.logger.warning(f'Failed to externalize image to media cache: {e}') node['base64'] = None for k, v in list(node.items()): if isinstance(v, (list, dict)): node[k] = await self.externalize_chain_dump(v) return node return chain_dump