"""
Metadata handling pipeline: extract, classify, strip.

Implements the extract-then-strip model for image metadata:
1. Extract all EXIF/metadata from raw bytes
2. Classify fields as evidentiary (GPS, timestamp) or dangerous (device serial)
3. Preserve evidentiary fields in a sidecar record
4. Strip all metadata from the stored/display copy

This resolves the tension between steganography (strip everything to protect
sources) and attestation (preserve everything to prove provenance).
"""

from __future__ import annotations

import io
from dataclasses import dataclass, field


# Fields classified by risk level:
# - EVIDENTIARY: valuable for proving provenance (location, time)
# - DANGEROUS: could identify the source's device or person
# - NEUTRAL: neither helpful nor harmful

EVIDENTIARY_FIELDS = {
    "captured_at",      # When the photo was taken
    "location",         # GPS coordinates (proves location)
    "width",            # Image dimensions
    "height",
}

DANGEROUS_FIELDS = {
    "device",           # Camera make/model/software/serial — fingerprints the source
}

# Granular device sub-fields
DANGEROUS_DEVICE_FIELDS = {
    "serial_hash",      # Device serial number hash
    "software",         # Firmware version (narrows device population)
}

SAFE_DEVICE_FIELDS = {
    "make",             # Broad manufacturer (e.g., "Apple") — low risk
    "model",            # Device model — moderate risk, useful for sensor analysis
}


@dataclass
class MetadataExtraction:
    """Result of extracting and classifying image metadata."""

    # Raw EXIF data as extracted by verisoo
    raw_exif: dict = field(default_factory=dict)

    # Classified subsets
    evidentiary: dict = field(default_factory=dict)  # Safe to include in attestation
    dangerous: dict = field(default_factory=dict)     # Should be redacted or stored separately
    all_fields: dict = field(default_factory=dict)    # Everything (for secure sidecar)

    # Hashes
    original_sha256: str = ""   # Hash of original bytes (with EXIF)
    stripped_sha256: str = ""   # Hash of stripped bytes (without EXIF)


def extract_and_classify(image_data: bytes) -> MetadataExtraction:
    """Extract EXIF metadata from image bytes and classify by risk level.

    Returns a MetadataExtraction with evidentiary (safe for attestation)
    and dangerous (should be redacted) fields separated.
    """
    import hashlib

    result = MetadataExtraction()
    result.original_sha256 = hashlib.sha256(image_data).hexdigest()

    # Extract raw EXIF via verisoo's extractor
    try:
        from soosef.verisoo.attestation import extract_exif_metadata

        raw = extract_exif_metadata(image_data)
        result.raw_exif = raw
        result.all_fields = dict(raw)
    except ImportError:
        return result

    # Classify fields
    for key in EVIDENTIARY_FIELDS:
        if key in raw:
            result.evidentiary[key] = raw[key]

    # Handle device fields granularly
    device = raw.get("device")
    if device:
        # Keep safe device fields in evidentiary
        safe_device = {}
        dangerous_device = {}
        device_dict = device if isinstance(device, dict) else {}
        if hasattr(device, "__dataclass_fields__"):
            from dataclasses import asdict
            device_dict = asdict(device)

        for k, v in device_dict.items():
            if v is None:
                continue
            if k in DANGEROUS_DEVICE_FIELDS:
                dangerous_device[k] = v
            elif k in SAFE_DEVICE_FIELDS:
                safe_device[k] = v
            else:
                safe_device[k] = v  # Default to safe for unknown fields

        if safe_device:
            result.evidentiary["device"] = safe_device
        if dangerous_device:
            result.dangerous["device"] = dangerous_device

    return result


def strip_metadata(image_data: bytes) -> bytes:
    """Strip all metadata from image bytes. Returns clean image bytes."""
    import hashlib

    try:
        from PIL import Image

        img = Image.open(io.BytesIO(image_data))
        clean = io.BytesIO()
        # Re-save without copying info/exif — strips all metadata
        img.save(clean, format=img.format or "PNG")
        return clean.getvalue()
    except Exception:
        # Not an image or Pillow can't handle it — return as-is
        return image_data


def extract_strip_pipeline(image_data: bytes) -> tuple[MetadataExtraction, bytes]:
    """Full extract-then-strip pipeline.

    Returns:
        (extraction, stripped_bytes) — extraction contains classified metadata,
        stripped_bytes is the clean image for storage/display.
    """
    import hashlib

    extraction = extract_and_classify(image_data)
    stripped = strip_metadata(image_data)
    extraction.stripped_sha256 = hashlib.sha256(stripped).hexdigest()
    return extraction, stripped
