"""Detect and strip C2PA / AI-related metadata from PNG and JPEG (stdlib)."""

from __future__ import annotations

import re
import struct
import subprocess
import zlib
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

from common import which

PNG_SIG = b"\x89PNG\r\n\x1a\n"
JPEG_SOI = b"\xff\xd8"

# Chunk types / content patterns associated with C2PA / AI provenance
C2PA_MARKERS = (
    b"c2pa",
    b"C2PA",
    b"jumb",
    b"JUMB",
    b"c2ma",
    b"contentcredentials",
    b"contentauth",
    b"cai:",
    b"http://ns.adobe.com/xmp/InstanceID/",  # not solely C2PA
)

AI_META_HINTS = (
    b"c2pa",
    b"C2PA",
    b"contentcredentials",
    b"ContentCredentials",
    b"digitalSourceType",
    b"trainedAlgorithmicMedia",
    b"compositeWithTrainedAlgorithmicMedia",
    b"algorithmicMedia",
    b"AIGC",
    b"aigc",
    b"AI generated",
    b"Generated by",
    b"Claude",
    b"Anthropic",
    b"OpenAI",
    b"SynthID",
    b"dcterms:provenance",
)


@dataclass
class ImageInspectReport:
    path: str
    format: str  # png | jpeg | unknown
    has_c2pa: bool
    has_ai_metadata: bool
    findings: list[str] = field(default_factory=list)
    tools: dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> dict:
        return {
            "path": self.path,
            "format": self.format,
            "has_c2pa": self.has_c2pa,
            "has_ai_metadata": self.has_ai_metadata,
            "findings": self.findings,
            "tools": self.tools,
        }


def detect_format(data: bytes) -> str:
    if data.startswith(PNG_SIG):
        return "png"
    if data.startswith(JPEG_SOI):
        return "jpeg"
    return "unknown"


def _contains_any(blob: bytes, needles: tuple[bytes, ...]) -> list[str]:
    found = []
    lower = blob.lower()
    for n in needles:
        if n.lower() in lower:
            try:
                found.append(n.decode("ascii", errors="replace"))
            except Exception:
                found.append(repr(n))
    return found


def inspect_png(data: bytes) -> tuple[bool, bool, list[str]]:
    findings: list[str] = []
    has_c2pa = False
    has_ai = False
    if not data.startswith(PNG_SIG):
        return False, False, ["not a PNG"]
    pos = 8
    while pos + 8 <= len(data):
        length = struct.unpack(">I", data[pos : pos + 4])[0]
        ctype = data[pos + 4 : pos + 8]
        chunk_start = pos + 8
        chunk_end = chunk_start + length
        if chunk_end + 4 > len(data):
            findings.append(f"truncated chunk {ctype!r}")
            break
        payload = data[chunk_start:chunk_end]
        name = ctype.decode("latin-1", errors="replace")
        # Private/ancillary chunks sometimes used for JUMBF/C2PA
        if ctype in (b"caBX", b"juMB", b"jumb") or ctype.startswith(b"c2"):
            has_c2pa = True
            findings.append(f"PNG chunk {name} (possible C2PA container)")
        if ctype in (b"tEXt", b"zTXt", b"iTXt", b"eXIf"):
            hits = _contains_any(payload, AI_META_HINTS + C2PA_MARKERS)
            if hits:
                has_ai = True
                if any(h.lower() in ("c2pa", "contentcredentials", "jumb") for h in hits):
                    has_c2pa = True
                findings.append(f"PNG {name}: {', '.join(hits[:8])}")
        if ctype == b"IEND":
            break
        pos = chunk_end + 4  # skip CRC
    # Whole-file scan fallback
    whole = _contains_any(data, C2PA_MARKERS)
    if whole and not has_c2pa:
        has_c2pa = True
        findings.append(f"byte-scan C2PA markers: {', '.join(whole[:6])}")
    return has_c2pa, has_ai or has_c2pa, findings


def inspect_jpeg(data: bytes) -> tuple[bool, bool, list[str]]:
    findings: list[str] = []
    has_c2pa = False
    has_ai = False
    if not data.startswith(JPEG_SOI):
        return False, False, ["not a JPEG"]
    i = 2
    n = len(data)
    while i + 4 <= n:
        if data[i] != 0xFF:
            i += 1
            continue
        # skip fill
        while i < n and data[i] == 0xFF:
            i += 1
        if i >= n:
            break
        marker = data[i]
        i += 1
        if marker in (0xD8, 0xD9):  # SOI/EOI
            continue
        if marker == 0xDA:  # SOS — image data follows
            break
        if marker >= 0xD0 and marker <= 0xD7:  # RSTn
            continue
        if i + 2 > n:
            break
        seglen = struct.unpack(">H", data[i : i + 2])[0]
        if seglen < 2 or i + seglen > n:
            findings.append(f"bad segment length at marker 0x{marker:02X}")
            break
        payload = data[i + 2 : i + seglen]
        i += seglen

        # APP11 (0xEB) often holds JUMBF/C2PA
        if marker == 0xEB:
            has_c2pa = True
            findings.append("JPEG APP11 segment (JUMBF/C2PA common)")
        if marker in (0xE1, 0xE2, 0xED, 0xEE, 0xEB):  # APP1,2,13,14,11
            hits = _contains_any(payload, AI_META_HINTS + C2PA_MARKERS)
            if hits:
                has_ai = True
                if any(
                    h.lower() in ("c2pa", "contentcredentials", "jumb", "contentauth")
                    for h in hits
                ):
                    has_c2pa = True
                findings.append(f"JPEG APP{marker - 0xE0}: {', '.join(hits[:8])}")

    whole = _contains_any(data, C2PA_MARKERS)
    if whole and not has_c2pa:
        has_c2pa = True
        findings.append(f"byte-scan C2PA markers: {', '.join(whole[:6])}")
    return has_c2pa, has_ai or has_c2pa, findings


def run_optional_tools(path: Path) -> dict[str, Any]:
    tools: dict[str, Any] = {}
    c2patool = which("c2patool")
    if c2patool:
        try:
            r = subprocess.run(
                [c2patool, str(path)],
                capture_output=True,
                text=True,
                timeout=30,
            )
            out = (r.stdout or "") + (r.stderr or "")
            tools["c2patool"] = {
                "available": True,
                "returncode": r.returncode,
                "snippet": out[:2000],
                "has_manifest": "claim" in out.lower()
                or "c2pa" in out.lower()
                or "manifest" in out.lower()
                and "no claim" not in out.lower()
                and "no jumbf" not in out.lower(),
            }
        except Exception as e:
            tools["c2patool"] = {"available": True, "error": str(e)}
    else:
        tools["c2patool"] = {"available": False}

    exiftool = which("exiftool")
    if exiftool:
        try:
            r = subprocess.run(
                [exiftool, "-G1", "-a", "-s", str(path)],
                capture_output=True,
                text=True,
                timeout=30,
            )
            out = r.stdout or ""
            interesting = [
                line
                for line in out.splitlines()
                if re.search(
                    r"c2pa|content.?credential|AIGC|digitalSource|XMP|EXIF|IPTC|jumb",
                    line,
                    re.I,
                )
            ]
            tools["exiftool"] = {
                "available": True,
                "interesting_lines": interesting[:50],
            }
        except Exception as e:
            tools["exiftool"] = {"available": True, "error": str(e)}
    else:
        tools["exiftool"] = {"available": False}
    return tools


def inspect_image(path: Path) -> ImageInspectReport:
    data = path.read_bytes()
    fmt = detect_format(data)
    if fmt == "png":
        has_c2pa, has_ai, findings = inspect_png(data)
    elif fmt == "jpeg":
        has_c2pa, has_ai, findings = inspect_jpeg(data)
    else:
        has_c2pa, has_ai, findings = False, False, ["unsupported format (MVP: PNG/JPEG)"]

    tools = run_optional_tools(path)
    # Elevate flags from tools
    ct = tools.get("c2patool") or {}
    if ct.get("has_manifest"):
        has_c2pa = True
        findings.append("c2patool reports a C2PA-related manifest")

    return ImageInspectReport(
        path=str(path),
        format=fmt,
        has_c2pa=has_c2pa,
        has_ai_metadata=has_ai,
        findings=findings,
        tools=tools,
    )


def _png_chunk(ctype: bytes, payload: bytes) -> bytes:
    crc = zlib.crc32(ctype)
    crc = zlib.crc32(payload, crc) & 0xFFFFFFFF
    return struct.pack(">I", len(payload)) + ctype + payload + struct.pack(">I", crc)


def strip_png(data: bytes, *, strip_all_text: bool = True) -> tuple[bytes, list[str]]:
    if not data.startswith(PNG_SIG):
        raise ValueError("not PNG")
    actions: list[str] = []
    out = bytearray(PNG_SIG)
    pos = 8
    while pos + 8 <= len(data):
        length = struct.unpack(">I", data[pos : pos + 4])[0]
        ctype = data[pos + 4 : pos + 8]
        chunk_start = pos + 8
        chunk_end = chunk_start + length
        if chunk_end + 4 > len(data):
            break
        payload = data[chunk_start:chunk_end]
        crc_bytes = data[chunk_end : chunk_end + 4]
        pos = chunk_end + 4
        name = ctype.decode("latin-1", errors="replace")

        drop = False
        if ctype in (b"eXIf", b"caBX") or ctype.startswith(b"c2"):
            drop = True
            actions.append(f"drop chunk {name}")
        elif ctype in (b"tEXt", b"zTXt", b"iTXt"):
            if strip_all_text or _contains_any(payload, AI_META_HINTS + C2PA_MARKERS):
                drop = True
                actions.append(f"drop chunk {name}")
        elif _contains_any(ctype + payload, C2PA_MARKERS) and ctype not in (
            b"IHDR",
            b"IDAT",
            b"IEND",
            b"PLTE",
            b"tRNS",
            b"gAMA",
            b"pHYs",
            b"sRGB",
            b"cHRM",
            b"iCCP",
        ):
            drop = True
            actions.append(f"drop chunk {name} (C2PA marker in payload)")

        if not drop:
            out.extend(struct.pack(">I", length) + ctype + payload + crc_bytes)
        if ctype == b"IEND":
            break
    if not actions:
        actions.append("no PNG metadata chunks removed (already clean or none matched)")
    return bytes(out), actions


def strip_jpeg(data: bytes, *, strip_all_app: bool = True) -> tuple[bytes, list[str]]:
    if not data.startswith(JPEG_SOI):
        raise ValueError("not JPEG")
    actions: list[str] = []
    out = bytearray(JPEG_SOI)
    i = 2
    n = len(data)
    while i < n:
        if data[i] != 0xFF:
            # unexpected; copy rest
            out.extend(data[i:])
            actions.append("copied remainder after non-marker byte")
            break
        while i < n and data[i] == 0xFF:
            i += 1
        if i >= n:
            break
        marker = data[i]
        i += 1

        if marker == 0xD9:  # EOI
            out.extend(b"\xff\xd9")
            break
        if marker == 0xD8:  # nested SOI?
            continue
        if 0xD0 <= marker <= 0xD7:
            out.extend(bytes([0xFF, marker]))
            continue
        if marker == 0xDA:  # SOS — copy through end
            # need length of SOS header then entropy-coded data to EOI
            if i + 2 > n:
                break
            seglen = struct.unpack(">H", data[i : i + 2])[0]
            # Find EOI from here carefully: after SOS segment header, scan for FF D9
            # not preceded by stuffed FF 00 issues — simple approach: copy from FF DA to end
            sos_start = i - 2  # points at 0xFF before marker... actually marker already consumed
            # Reconstruct: FF DA + rest of file
            out.extend(b"\xff\xda")
            out.extend(data[i:])
            actions.append("preserved entropy-coded scan (SOS→EOF)")
            break

        if i + 2 > n:
            break
        seglen = struct.unpack(">H", data[i : i + 2])[0]
        if seglen < 2 or i + seglen > n:
            out.extend(data[i - 2 :])  # best effort
            actions.append("truncated segment; copied remainder")
            break
        payload = data[i + 2 : i + seglen]
        next_i = i + seglen

        # APP0 is JFIF — keep for compatibility unless full strip of all APP
        keep = False
        drop = False
        if 0xE0 <= marker <= 0xEF:  # APPn
            if marker == 0xEB:  # APP11 JUMBF/C2PA
                drop = True
                actions.append("drop APP11 (C2PA/JUMBF)")
            elif strip_all_app and marker != 0xE0:
                # keep APP0 (JFIF) by default
                drop = True
                actions.append(f"drop APP{marker - 0xE0}")
            elif _contains_any(payload, AI_META_HINTS + C2PA_MARKERS):
                drop = True
                actions.append(f"drop APP{marker - 0xE0} (AI/C2PA markers)")
            else:
                keep = True
        elif marker in (0xFE,):  # COM
            drop = True
            actions.append("drop COM comment")
        else:
            keep = True

        if keep and not drop:
            out.extend(bytes([0xFF, marker]))
            out.extend(data[i : i + seglen])
        i = next_i

    if not actions:
        actions.append("no JPEG APP segments removed")
    return bytes(out), actions


def clean_image(
    path: Path,
    dest: Path,
    *,
    strip_all_metadata: bool = True,
) -> dict[str, Any]:
    data = path.read_bytes()
    fmt = detect_format(data)
    if fmt == "png":
        cleaned, actions = strip_png(data, strip_all_text=strip_all_metadata)
    elif fmt == "jpeg":
        cleaned, actions = strip_jpeg(data, strip_all_app=strip_all_metadata)
    else:
        raise ValueError(f"unsupported format: {fmt}")

    # Optional exiftool pass for residual tags
    dest.parent.mkdir(parents=True, exist_ok=True)
    dest.write_bytes(cleaned)
    exiftool = which("exiftool")
    if exiftool and strip_all_metadata:
        try:
            subprocess.run(
                [
                    exiftool,
                    "-all=",
                    "-overwrite_original",
                    str(dest),
                ],
                capture_output=True,
                text=True,
                timeout=60,
                check=False,
            )
            actions.append("exiftool -all= pass")
        except Exception as e:
            actions.append(f"exiftool failed: {e}")

    after = inspect_image(dest)
    return {
        "input": str(path),
        "output": str(dest),
        "format": fmt,
        "actions": actions,
        "bytes_in": len(data),
        "bytes_out": dest.stat().st_size,
        "still_has_c2pa": after.has_c2pa,
        "still_has_ai_metadata": after.has_ai_metadata,
        "post_findings": after.findings,
    }