"""
ThinkStorm Image Intake & Sanitization Service
Handles safe image validation, EXIF/GPS metadata stripping, token-optimized vision resizing,
durable artifact persistence, and canonical metadata generation.
"""

import io
import os
import hashlib
from pathlib import Path
from typing import Optional, Dict, Any, Tuple
from PIL import Image, ImageOps

from ..config import config, BASE_DIR

ARTIFACTS_DIR = BASE_DIR / "data" / "artifacts"
ARTIFACTS_DIR.mkdir(parents=True, exist_ok=True)

# Map PIL format strings to canonical MIME types and file extensions
FORMAT_TO_MIME = {
    "JPEG": "image/jpeg",
    "PNG": "image/png",
    "WEBP": "image/webp"
}

MIME_TO_EXT = {
    "image/jpeg": ".jpg",
    "image/png": ".png",
    "image/webp": ".webp"
}

class ImageValidationError(Exception):
    """Raised when an uploaded or received file fails image validation."""
    pass

class ImageTooLargeError(ImageValidationError):
    """Raised when an image exceeds configured size limit."""
    pass

class ImageFormatError(ImageValidationError):
    """Raised when image format is unsupported or invalid."""
    pass

def sanitize_filename(filename: Optional[str], default_ext: str = ".jpg") -> str:
    """Strips path traversal elements and unsafe characters from original filename."""
    if not filename:
        return f"submission_image{default_ext}"
    base = os.path.basename(filename).strip()
    # Remove null bytes and non-printable characters
    clean = "".join(c for c in base if c.isalnum() or c in "._- ")
    if not clean or clean.startswith("."):
        return f"submission_image{default_ext}"
    return clean[:128]

def validate_and_sanitize_image(
    raw_bytes: bytes,
    original_filename: str = "",
    source: str = "web",
    idea_id: Optional[str] = None
) -> Tuple[bytes, Dict[str, Any]]:
    """
    Validates, strips metadata, and generates canonical sanitized image bytes and metadata.
    
    Returns:
        Tuple[bytes, Dict[str, Any]]: (sanitized_bytes, metadata_dict)
    
    Raises:
        ImageTooLargeError: If payload exceeds max size.
        ImageFormatError: If data is corrupted or unsupported format.
    """
    if not raw_bytes:
        raise ImageFormatError("Image payload cannot be empty.")

    max_bytes = config.max_image_upload_bytes
    if len(raw_bytes) > max_bytes:
        raise ImageTooLargeError(
            f"Image size ({len(raw_bytes)} bytes) exceeds maximum limit of {max_bytes} bytes ({max_bytes // (1024 * 1024)}MB)."
        )

    # 1. Open and verify image structure with Pillow
    try:
        in_stream = io.BytesIO(raw_bytes)
        img = Image.open(in_stream)
        img_format = img.format
    except Exception as e:
        raise ImageFormatError(f"Could not decode image header: {str(e)}")

    if not img_format or img_format.upper() not in FORMAT_TO_MIME:
        allowed_str = ", ".join(config.allowed_image_mime_types)
        raise ImageFormatError(
            f"Unsupported image format: '{img_format}'. Supported formats are: {allowed_str} (JPEG, PNG, WebP)."
        )

    canonical_mime = FORMAT_TO_MIME[img_format.upper()]

    # 2. Decode pixel data to ensure file is not truncated or corrupted
    try:
        # Respect EXIF orientation before stripping EXIF metadata
        img = ImageOps.exif_transpose(img)
        img.load()
    except Exception as e:
        raise ImageFormatError(f"Malformed or corrupt image pixel stream: {str(e)}")

    # 3. Create sanitized copy (stripping EXIF, GPS, camera metadata, comments)
    out_format = img_format.upper()
    out_stream = io.BytesIO()

    try:
        if out_format == "JPEG":
            # Convert palette/RGBA modes to RGB for JPEG
            if img.mode in ("RGBA", "LA", "P"):
                rgb_img = Image.new("RGB", img.size, (255, 255, 255))
                if img.mode == "P":
                    img = img.convert("RGBA")
                rgb_img.paste(img, mask=img.split()[-1] if img.mode == "RGBA" else None)
                img = rgb_img
            elif img.mode != "RGB":
                img = img.convert("RGB")
            img.save(out_stream, format="JPEG", quality=92, optimize=True)

        elif out_format == "PNG":
            # Preserve RGB / RGBA transparency
            if img.mode not in ("RGB", "RGBA", "L", "LA"):
                img = img.convert("RGBA" if "transparency" in img.info else "RGB")
            img.save(out_stream, format="PNG", optimize=True)

        elif out_format == "WEBP":
            if img.mode not in ("RGB", "RGBA"):
                img = img.convert("RGBA" if "transparency" in img.info else "RGB")
            img.save(out_stream, format="WEBP", quality=90, method=6)
        else:
            raise ImageFormatError(f"Unsupported save format: {out_format}")

    except Exception as e:
        raise ImageFormatError(f"Failed to encode sanitized image: {str(e)}")

    sanitized_bytes = out_stream.getvalue()
    width, height = img.size
    sha256_hash = hashlib.sha256(sanitized_bytes).hexdigest()
    ext = MIME_TO_EXT.get(canonical_mime, ".jpg")
    safe_filename = sanitize_filename(original_filename, default_ext=ext)

    artifact_id = f"IMG-{idea_id}" if idea_id else f"IMG-{sha256_hash[:12].upper()}"

    metadata: Dict[str, Any] = {
        "present": True,
        "artifact_id": artifact_id,
        "original_filename": safe_filename,
        "mime_type": canonical_mime,
        "size_bytes": len(sanitized_bytes),
        "width": width,
        "height": height,
        "sha256": sha256_hash,
        "source": source,
        "vision_analysis_run_id": None
    }

    return sanitized_bytes, metadata

def save_image_artifact(idea_id: str, sanitized_bytes: bytes, metadata: Dict[str, Any]) -> str:
    """
    Saves the canonical sanitized image file to the idea's artifact directory.
    Returns relative or absolute path to saved artifact.
    """
    idea_dir = ARTIFACTS_DIR / idea_id
    idea_dir.mkdir(parents=True, exist_ok=True)

    artifact_id = metadata.get("artifact_id") or f"IMG-{idea_id}"
    mime = metadata.get("mime_type", "image/jpeg")
    ext = MIME_TO_EXT.get(mime, ".jpg")

    file_path = idea_dir / f"{artifact_id}{ext}"
    file_path.write_bytes(sanitized_bytes)

    # Also keep a predictable reference file for easy retrieval
    canonical_link = idea_dir / f"reference_image{ext}"
    if canonical_link != file_path:
        canonical_link.write_bytes(sanitized_bytes)

    return str(file_path)

def get_image_artifact_path(idea_id: str, metadata: Optional[Dict[str, Any]] = None) -> Optional[Path]:
    """Retrieves file path to an idea's saved reference image."""
    idea_dir = ARTIFACTS_DIR / idea_id
    if not idea_dir.exists():
        return None

    if metadata and metadata.get("artifact_id"):
        artifact_id = metadata["artifact_id"]
        mime = metadata.get("mime_type", "image/jpeg")
        ext = MIME_TO_EXT.get(mime, ".jpg")
        path = idea_dir / f"{artifact_id}{ext}"
        if path.exists():
            return path

    # Fallback to search any image in idea artifact directory
    for candidate_ext in [".jpg", ".jpeg", ".png", ".webp"]:
        ref = idea_dir / f"reference_image{candidate_ext}"
        if ref.exists():
            return ref
        img_match = list(idea_dir.glob(f"IMG-*{candidate_ext}"))
        if img_match:
            return img_match[0]

    return None

def create_token_optimized_vision_payload(
    sanitized_bytes: bytes,
    mime_type: str = "image/jpeg"
) -> Tuple[bytes, str]:
    """
    Optimizes and downscales image dimensions to bound vision LLM token consumption.
    Bounds maximum dimension to `config.vision_max_dimension` (default 1536px)
    and compresses to efficient JPEG to minimize token expenditure across vision grid tiles.
    
    Returns:
        Tuple[bytes, str]: (optimized_bytes, optimized_mime_type)
    """
    try:
        img = Image.open(io.BytesIO(sanitized_bytes))
        width, height = img.size
        max_dim = config.vision_max_dimension

        # Downscale if larger than max_dim while preserving aspect ratio
        if width > max_dim or height > max_dim:
            img.thumbnail((max_dim, max_dim), Image.Resampling.LANCZOS)

        # Convert to RGB if needed for JPEG compression
        if img.mode in ("RGBA", "LA", "P"):
            rgb_img = Image.new("RGB", img.size, (255, 255, 255))
            if img.mode == "P":
                img = img.convert("RGBA")
            rgb_img.paste(img, mask=img.split()[-1] if img.mode == "RGBA" else None)
            img = rgb_img
        elif img.mode != "RGB":
            img = img.convert("RGB")

        out_buf = io.BytesIO()
        img.save(
            out_buf,
            format="JPEG",
            quality=config.vision_jpeg_quality,
            optimize=True
        )
        return out_buf.getvalue(), "image/jpeg"

    except Exception:
        # Fallback to original sanitized bytes if downscaling fails
        return sanitized_bytes, mime_type