engine.py

"""
SanitizeAI Core Engine
Centralized, modular engine for image metadata stripping, invisible watermark destruction,
visual watermark removal, deep diffusion regeneration, and text/video sanitization.
"""

from __future__ import annotations

import io
import os
import re
import sys
import time
import shutil
import tempfile
import subprocess
from enum import Enum
from pathlib import Path
from typing import Any, BinaryIO, Literal, Union, Optional, Tuple
from PIL import Image

# Ensure local engine directory is accessible
ENGINE_DIR = Path(__file__).parent.resolve()
if str(ENGINE_DIR) not in sys.path:
    sys.path.insert(0, str(ENGINE_DIR))

# Core submodule imports
from cleaner import remove_ai_metadata, _extract_non_ai_metadata
from constants import AI_METADATA_KEYS, AI_KEYWORDS
from pixel_cleaner import (
    sanitize_frame_pixels,
    destroy_synthid_watermark,
    remove_visual_logo_watermark,
)
from pipeline import (
    sanitize_video_pipeline,
    sanitize_video_mp4,
    validate_mp4_container,
    calculate_logo_bbox,
    probe_video_metadata,
    inpaint_frame_logo,
)

# Optional Watermark Remover (PyTorch / Diffusers)
try:
    from watermark_remover import WatermarkRemover, is_watermark_removal_available
    HAS_DIFFUSION_REMOVER = is_watermark_removal_available()
except (ImportError, Exception):
    HAS_DIFFUSION_REMOVER = False
    WatermarkRemover = None


class CleanMode(str, Enum):
    SYNTHID_DESTRUCTOR = "synthid_destructor"
    VISUAL_PURGE = "visual_purge"
    METADATA_ONLY = "metadata_only"
    DEEP_DIFFUSION = "deep_diffusion"


def normalize_clean_mode(mode: str) -> CleanMode:
    """Normalize user or UI input strings into a canonical CleanMode."""
    m = mode.lower().strip()
    if "deep" in m or "diffusion" in m:
        return CleanMode.DEEP_DIFFUSION
    elif "meta" in m and "only" in m:
        return CleanMode.METADATA_ONLY
    elif "logo" in m or "visual" in m or "purge" in m:
        return CleanMode.VISUAL_PURGE
    elif "synthid" in m or "destructor" in m or "pixel" in m:
        return CleanMode.SYNTHID_DESTRUCTOR
    
    # Direct match attempt
    for cm in CleanMode:
        if cm.value == m:
            return cm
    return CleanMode.SYNTHID_DESTRUCTOR


def get_ffmpeg_path() -> str | None:
    """Find FFmpeg binary in local directory, PATH, or imageio-ffmpeg."""
    local_ffmpeg = ENGINE_DIR / "ffmpeg.exe"
    if local_ffmpeg.exists():
        return str(local_ffmpeg.resolve())
    
    ffmpeg_sys = shutil.which("ffmpeg")
    if ffmpeg_sys:
        return ffmpeg_sys

    try:
        import imageio_ffmpeg
        return imageio_ffmpeg.get_ffmpeg_exe()
    except Exception:
        pass
    return None


FFMPEG_PATH = get_ffmpeg_path()


def get_video_fps(video_path: Path | str, ffmpeg_path: str | None = None) -> float:
    """Probe exact original FPS from video using FFmpeg."""
    ff_path = ffmpeg_path or FFMPEG_PATH
    if not ff_path:
        return 30.0
    try:
        cmd = [ff_path, "-i", str(video_path)]
        res = subprocess.run(cmd, stderr=subprocess.PIPE, stdout=subprocess.PIPE, text=True)
        match = re.search(r'(\d+(?:\.\d+)?)\s+fps', res.stderr)
        if match:
            return float(match.group(1))
        match = re.search(r'(\d+(?:\.\d+)?)\s+tbr', res.stderr)
        if match:
            return float(match.group(1))
    except Exception:
        pass
    return 30.0


def analyze_image_metadata(img_input: Union[Path, str, bytes, Image.Image]) -> dict[str, Any]:
    """
    Analyze image metadata to detect AI metadata fields, prompt tags, and C2PA signatures.
    """
    detected_keys: list[str] = []
    has_c2pa = False

    def _inspect_img(img: Image.Image):
        nonlocal has_c2pa
        for key, val in img.info.items():
            key_lower = str(key).lower()
            val_str = str(val).lower()
            is_ai = (
                any(k.lower() in key_lower for k in AI_METADATA_KEYS)
                or any(kw in key_lower for kw in AI_KEYWORDS)
                or any(kw in val_str for kw in ["stable diffusion", "comfyui", "midjourney", "dall-e", "c2pa"])
            )
            if is_ai:
                detected_keys.append(str(key))
            if "c2pa" in key_lower or "jumbf" in key_lower or "c2pa" in val_str:
                has_c2pa = True

    try:
        if isinstance(img_input, (str, Path)):
            with Image.open(img_input) as img:
                _inspect_img(img)
        elif isinstance(img_input, bytes):
            with Image.open(io.BytesIO(img_input)) as img:
                _inspect_img(img)
        elif isinstance(img_input, Image.Image):
            _inspect_img(img_input)
    except Exception:
        pass

    return {
        "detected_keys": detected_keys,
        "has_c2pa": has_c2pa,
        "total_fields": len(detected_keys),
    }


def clean_text_content(text: str, remove_ai_phrases: bool = False) -> tuple[str, dict[str, Any]]:
    """
    Clean zero-width characters, hidden markers, Cyrillic homoglyphs, and AI conversational markers.
    """
    stats = {
        "zero_width": 0,
        "control_chars": 0,
        "homoglyphs": 0,
        "ai_phrases": 0,
        "initial_len": len(text),
    }

    # 1. Zero-width spaces & invisible characters
    zero_width_chars = [
        '\u200B', '\u200C', '\u200D', '\uFEFF', '\u00AD',
        '\u2060', '\u200E', '\u200F', '\u202A', '\u202B',
        '\u202C', '\u202D', '\u202E', '\u2061', '\u2062',
        '\u2063', '\u2064',
    ]

    cleaned = text
    for zw in zero_width_chars:
        count = cleaned.count(zw)
        if count > 0:
            stats["zero_width"] += count
            cleaned = cleaned.replace(zw, '')

    # 2. Other invisible control characters
    control_pattern = re.compile(r'[\x00-\x08\x0B\x0C\x0E-\x1F\x7F-\x9F]')
    ctrl_matches = len(control_pattern.findall(cleaned))
    stats["control_chars"] += ctrl_matches
    cleaned = control_pattern.sub('', cleaned)

    # 3. Common Homoglyphs (Cyrillic to Latin lookalikes)
    homoglyph_map = {
        'а': 'a', 'е': 'e', 'о': 'o', 'р': 'p', 'с': 'c', 'х': 'x', 'у': 'y',
        'А': 'A', 'В': 'B', 'E': 'E', 'К': 'K', 'М': 'M', 'Н': 'N', 'О': 'O',
        'Р': 'P', 'С': 'C', 'Т': 'T', 'Х': 'X',
    }
    has_cyrillic = any(char in homoglyph_map for char in cleaned)
    if has_cyrillic:
        for cyr, lat in homoglyph_map.items():
            cnt = cleaned.count(cyr)
            if cnt > 0:
                stats["homoglyphs"] += cnt
                cleaned = cleaned.replace(cyr, lat)

    # 4. Optional AI Conversational Marker Removal
    if remove_ai_phrases:
        ai_patterns = [
            r"(?i)^As an AI language model,\s*",
            r"(?i)^Certainly!\s+Here is\s*",
            r"(?i)^Sure,\s+I can help with that\.?\s*",
            r"(?i)^Here's a breakdown of\s*",
            r"(?i)If you have any further questions, feel free to ask!\s*$",
            r"(?i)I hope this helps!\s*$",
        ]
        for pat in ai_patterns:
            matches = len(re.findall(pat, cleaned))
            if matches > 0:
                stats["ai_phrases"] += matches
                cleaned = re.sub(pat, '', cleaned).strip()

    stats["final_len"] = len(cleaned)
    stats["total_removed"] = (
        stats["zero_width"] + stats["control_chars"] + stats["homoglyphs"] + stats["ai_phrases"]
    )
    return cleaned, stats


def add_sanitizeai_watermark(image_pil: Image.Image) -> Image.Image:
    """
    Apply a discreet, professional 'SanitizeAI' watermark pill to the bottom-right of an image.
    Used for the $5 Starter tier.
    """
    from PIL import ImageDraw, ImageFont

    img = image_pil.convert("RGBA")
    overlay = Image.new("RGBA", img.size, (255, 255, 255, 0))
    draw = ImageDraw.Draw(overlay)

    w, h = img.size
    text = "SanitizeAI • Basic"

    # Adaptive font size based on image dimensions
    font_size = max(12, int(min(w, h) * 0.028))
    try:
        font = ImageFont.truetype("arial.ttf", font_size)
    except Exception:
        font = ImageFont.load_default()

    bbox = draw.textbbox((0, 0), text, font=font)
    text_w = bbox[2] - bbox[0]
    text_h = bbox[3] - bbox[1]

    pad_x = int(font_size * 0.6)
    pad_y = int(font_size * 0.4)
    pill_w = text_w + pad_x * 2
    pill_h = text_h + pad_y * 2

    margin = int(min(w, h) * 0.025)
    x1 = w - pill_w - margin
    y1 = h - pill_h - margin
    x2 = x1 + pill_w
    y2 = y1 + pill_h

    # Draw sleek dark translucent pill with indigo border
    draw.rounded_rectangle([x1, y1, x2, y2], radius=int(font_size * 0.35), fill=(10, 14, 26, 205), outline=(99, 102, 241, 160), width=1)
    draw.text((x1 + pad_x, y1 + pad_y - 1), text, font=font, fill=(241, 245, 249, 235))

    watermarked = Image.alpha_composite(img, overlay)
    if image_pil.mode in ("RGB", "L"):
        return watermarked.convert("RGB")
    return watermarked


def clean_image(
    image_input: Union[bytes, Image.Image, Path, str, BinaryIO],
    mode: Union[str, CleanMode] = CleanMode.SYNTHID_DESTRUCTOR,
    intensity: float = 0.5,
    logo_location: str = "bottom_right",
    keep_standard_metadata: bool = True,
    apply_watermark: bool = False,
    output_format: str | None = None,
    device: str = "auto",
) -> tuple[bytes, str]:
    """
    Unified Image Cleaning & Sanitization Function.
    
    Args:
        image_input: Raw image bytes, PIL Image, file path, or file-like object.
        mode: Sanitization mode ('synthid_destructor', 'visual_purge', 'metadata_only', 'deep_diffusion').
        intensity: Denoising / watermark destruction strength (0.01 - 1.0, default 0.5).
        logo_location: Corner location for visual logo purge ('bottom_right', 'bottom_left', 'top_right', 'top_left', 'all_corners').
        keep_standard_metadata: Whether to preserve non-AI EXIF/DPI metadata.
        apply_watermark: If True ($5 Starter plan), preserves original AI logos and overlays SanitizeAI watermark.
        output_format: Target format ('PNG', 'JPEG', 'WEBP'). Defaults to input format.
        device: Target inference device for diffusion mode ('auto', 'cpu', 'cuda').
        
    Returns:
        tuple of (cleaned_image_bytes, mime_type)
    """
    clean_mode = normalize_clean_mode(str(mode))
    intensity = max(0.01, min(1.0, float(intensity)))

    # Load PIL image and determine format
    orig_format = "PNG"
    if isinstance(image_input, (str, Path)):
        img_path = Path(image_input)
        raw_bytes = img_path.read_bytes()
    elif isinstance(image_input, bytes):
        raw_bytes = image_input
    elif hasattr(image_input, "read"):
        raw_bytes = image_input.read()
    elif isinstance(image_input, Image.Image):
        buf = io.BytesIO()
        image_input.save(buf, format=image_input.format or "PNG")
        raw_bytes = buf.getvalue()
    else:
        raise ValueError(f"Unsupported image input type: {type(image_input)}")

    with Image.open(io.BytesIO(raw_bytes)) as pil_img:
        orig_format = pil_img.format or "PNG"
        target_format = (output_format or orig_format).upper()
        if target_format in ("JPG", "JPEG"):
            target_format = "JPEG"
            mime_type = "image/jpeg"
        elif target_format == "WEBP":
            mime_type = "image/webp"
        else:
            target_format = "PNG"
            mime_type = "image/png"

    with tempfile.TemporaryDirectory() as tmp_dir:
        tmp_in = Path(tmp_dir) / f"input_raw.{target_format.lower()}"
        tmp_out = Path(tmp_dir) / f"cleaned_out.{target_format.lower()}"
        tmp_in.write_bytes(raw_bytes)

        if clean_mode == CleanMode.METADATA_ONLY:
            remove_ai_metadata(tmp_in, tmp_out, keep_standard=keep_standard_metadata)
            if apply_watermark:
                with Image.open(tmp_out) as tmp_pil:
                    wm_pil = add_sanitizeai_watermark(tmp_pil)
                    if target_format == "JPEG" and wm_pil.mode in ("RGBA", "P"):
                        wm_pil = wm_pil.convert("RGB")
                    wm_pil.save(tmp_out, format=target_format)
            cleaned_bytes = tmp_out.read_bytes()

        elif clean_mode == CleanMode.DEEP_DIFFUSION:
            if not HAS_DIFFUSION_REMOVER or WatermarkRemover is None:
                with Image.open(tmp_in) as img:
                    cleaned_pil = sanitize_frame_pixels(
                        img,
                        remove_synthid=True,
                        remove_logo=False if apply_watermark else False,
                        synthid_strength=intensity,
                    )
                    if apply_watermark:
                        cleaned_pil = add_sanitizeai_watermark(cleaned_pil)
                    if target_format == "JPEG" and cleaned_pil.mode in ("RGBA", "P"):
                        cleaned_pil = cleaned_pil.convert("RGB")
                    cleaned_pil.save(tmp_out, format=target_format)
                remove_ai_metadata(tmp_out, tmp_out, keep_standard=keep_standard_metadata)
            else:
                remover = WatermarkRemover(device=device)
                remover.remove_watermark(tmp_in, tmp_out, strength=min(0.2, intensity * 0.2))
                if apply_watermark:
                    with Image.open(tmp_out) as tmp_pil:
                        wm_pil = add_sanitizeai_watermark(tmp_pil)
                        if target_format == "JPEG" and wm_pil.mode in ("RGBA", "P"):
                            wm_pil = wm_pil.convert("RGB")
                        wm_pil.save(tmp_out, format=target_format)
                remove_ai_metadata(tmp_out, tmp_out, keep_standard=keep_standard_metadata)
            cleaned_bytes = tmp_out.read_bytes()

        else:
            # SynthID Destructor or Visual Logo Purge
            # If $5 plan (apply_watermark=True), do NOT purge original AI logo
            do_logo = False if apply_watermark else (clean_mode == CleanMode.VISUAL_PURGE)
            do_synthid = True
            loc = logo_location if do_logo else "bottom_right"

            with Image.open(tmp_in) as orig_pil:
                cleaned_pil = sanitize_frame_pixels(
                    orig_pil,
                    remove_synthid=do_synthid,
                    remove_logo=do_logo,
                    logo_location=loc,
                    synthid_strength=intensity,
                )
                if apply_watermark:
                    cleaned_pil = add_sanitizeai_watermark(cleaned_pil)
                if target_format == "JPEG" and cleaned_pil.mode in ("RGBA", "P"):
                    cleaned_pil = cleaned_pil.convert("RGB")
                cleaned_pil.save(tmp_out, format=target_format)

            remove_ai_metadata(tmp_out, tmp_out, keep_standard=keep_standard_metadata)
            cleaned_bytes = tmp_out.read_bytes()

    return cleaned_bytes, mime_type


def clean_video(
    input_video: Union[Path, str],
    output_video: Union[Path, str],
    remove_visible_logo: bool = True,
    remove_synthid: bool = True,
    remove_metadata: bool = True,
    logo_location: str = "bottom_right",
    custom_bbox: Optional[Tuple[int, int, int, int]] = None,
    engine_backend: str = "opencv",
    crf_quality: int = 18,
    progress_callback: Optional[Any] = None,
) -> Path:
    """
    Unified Video Sanitization Entry Point.
    Rebuilds clean MP4 containers while purging visible logos, invisible SynthID watermarks,
    and all metadata/C2PA manifests.
    """
    if not FFMPEG_PATH:
        raise RuntimeError("FFmpeg executable not found. Please install FFmpeg to process videos.")
    
    return sanitize_video_pipeline(
        input_video=Path(input_video),
        output_video=Path(output_video),
        ffmpeg_path=FFMPEG_PATH,
        remove_visible_logo=remove_visible_logo,
        remove_synthid=remove_synthid,
        remove_metadata=remove_metadata,
        logo_location=logo_location,
        custom_bbox=custom_bbox,
        engine_backend=engine_backend,
        crf_quality=crf_quality,
        progress_callback=progress_callback,
    )


def get_engine_status() -> dict[str, Any]:
    """Return status of all subsystem components."""
    return {
        "engine_version": "0.1.29",
        "ffmpeg_available": FFMPEG_PATH is not None,
        "ffmpeg_path": FFMPEG_PATH,
        "diffusion_remover_available": HAS_DIFFUSION_REMOVER,
        "supported_modes": [m.value for m in CleanMode],
    }