#!/usr/bin/env python3
"""One-use, local-only customer-surface evidence capture."""
from __future__ import annotations

import argparse
import hashlib
import json
import os
import stat
import struct
import sys
import tempfile
import zlib
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Callable

from PIL import Image, UnidentifiedImageError

ROOT = Path(__file__).resolve().parent
CONTRACT = ROOT / "capture-contract.json"
VAULT = ROOT / "vault"
RECEIPT = ROOT / "capture-receipt.redacted.json"
MAX_BYTES = 20 * 1024 * 1024
MAX_DIMENSION = 16_384
MAX_PIXELS = 100_000_000
ALLOWED_MODES = {"1", "L", "LA", "P", "RGB", "RGBA"}
ATTESTATION = "I_ATTEST_EXACTLY_ONE_VISIBLE_MESSAGE_MATCHES_THE_DURABLE_DELIVERED_TEXT"
PNG_SIG = b"\x89PNG\r\n\x1a\n"

class CaptureError(RuntimeError):
    pass

def canonical(value: Any) -> bytes:
    return (json.dumps(value, ensure_ascii=True, sort_keys=True, separators=(",", ":")) + "\n").encode()

def digest(data: bytes) -> str:
    return hashlib.sha256(data).hexdigest()

def _write_all(fd: int, data: bytes, write: Callable[[int, bytes], int] = os.write) -> None:
    view = memoryview(data)
    while view:
        written = write(fd, bytes(view))
        if written <= 0:
            raise CaptureError("incomplete write")
        view = view[written:]

def _fsync_directory(path: Path) -> None:
    fd = os.open(path, os.O_RDONLY | os.O_DIRECTORY | os.O_CLOEXEC)
    try:
        os.fsync(fd)
    finally:
        os.close(fd)

def atomic_write(path: Path, data: bytes, mode: int = 0o600) -> None:
    path.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
    os.chmod(path.parent, 0o700)
    fd, tmp_name = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent)
    tmp = Path(tmp_name)
    try:
        os.fchmod(fd, mode)
        _write_all(fd, data)
        os.fsync(fd)
        os.close(fd)
        fd = -1
        if path.exists() or path.is_symlink():
            raise CaptureError(f"one-use output already exists: {path.name}")
        # link(2) is an atomic no-replace publication: it cannot overwrite a
        # concurrently created one-use destination, unlike rename(2).
        os.link(tmp, path, follow_symlinks=False)
        tmp.unlink()
        _fsync_directory(path.parent)
    finally:
        if fd >= 0:
            os.close(fd)
        try:
            tmp.unlink()
        except FileNotFoundError:
            pass

def load_contract() -> tuple[dict[str, Any], bytes]:
    flags = os.O_RDONLY | os.O_CLOEXEC | os.O_NOFOLLOW
    fd = os.open(CONTRACT, flags)
    try:
        st = os.fstat(fd)
        if not stat.S_ISREG(st.st_mode) or stat.S_IMODE(st.st_mode) != 0o600:
            raise CaptureError("contract must be a mode-0600 regular file")
        raw = b""
        while True:
            chunk = os.read(fd, 65536)
            if not chunk:
                break
            raw += chunk
    finally:
        os.close(fd)
    value = json.loads(raw)
    required = {
        "candidate", "core_candidate", "hermes_wheel_sha256", "profile_wheel_sha256",
        "config_sha256", "customer_chat_id", "expected_provider_message_id",
        "expected_provider_timestamp", "delivery_row_key", "delivery_idempotency_key",
        "lifecycle_run_id", "lifecycle_session_id", "durable_text_sha256",
    }
    if set(value.get("binding", {})) != required or value.get("schema") != "task26-customer-surface-capture-contract-v1":
        raise CaptureError("capture contract schema/binding mismatch")
    for key, item in value["binding"].items():
        if not isinstance(item, str) or not item:
            raise CaptureError(f"invalid contract binding: {key}")
    return value, raw

def _strict_png(data: bytes) -> tuple[int, int]:
    if not data.startswith(PNG_SIG):
        raise CaptureError("invalid PNG signature")
    pos, chunks, idat, width, height = 8, 0, bytearray(), 0, 0
    seen_ihdr = seen_idat = False
    while pos + 12 <= len(data):
        length = struct.unpack(">I", data[pos:pos + 4])[0]
        kind = data[pos + 4:pos + 8]
        end = pos + 12 + length
        if end > len(data):
            raise CaptureError("truncated PNG chunk")
        payload = data[pos + 8:pos + 8 + length]
        expected_crc = struct.unpack(">I", data[pos + 8 + length:end])[0]
        if zlib.crc32(kind + payload) & 0xFFFFFFFF != expected_crc:
            raise CaptureError("PNG CRC mismatch")
        chunks += 1
        if chunks == 1:
            if kind != b"IHDR" or length != 13:
                raise CaptureError("PNG IHDR missing")
            width, height, depth, color, compression, filtering, interlace = struct.unpack(">IIBBBBB", payload)
            if depth not in {1, 2, 4, 8, 16} or color not in {0, 2, 3, 4, 6} or compression or filtering or interlace not in {0, 1}:
                raise CaptureError("unsupported PNG header")
            seen_ihdr = True
        elif kind == b"IHDR":
            raise CaptureError("duplicate PNG IHDR")
        if kind == b"IDAT":
            seen_idat = True
            idat.extend(payload)
        elif seen_idat and kind not in {b"IDAT", b"IEND"}:
            # Ancillary chunks after IDAT are valid, but a later IDAT is not.
            seen_idat = False
        if kind == b"IEND":
            if length or end != len(data):
                raise CaptureError("PNG has trailing/polyglot bytes")
            if not seen_ihdr or not idat:
                raise CaptureError("PNG image data missing")
            try:
                decompressor = zlib.decompressobj()
                decompressor.decompress(bytes(idat), MAX_BYTES * 64)
                if not decompressor.eof or decompressor.unused_data:
                    raise CaptureError("invalid PNG compressed stream")
            except zlib.error as exc:
                raise CaptureError("invalid PNG compressed stream") from exc
            return width, height
        pos = end
    raise CaptureError("PNG IEND missing")

def _strict_jpeg(data: bytes) -> tuple[int, int]:
    if len(data) < 4 or data[:2] != b"\xff\xd8":
        raise CaptureError("invalid JPEG signature")
    pos, dimensions, in_scan = 2, None, False
    sof = {0xC0, 0xC1, 0xC2, 0xC3, 0xC5, 0xC6, 0xC7, 0xC9, 0xCA, 0xCB, 0xCD, 0xCE, 0xCF}
    while pos < len(data):
        if in_scan:
            marker = data.find(b"\xff", pos)
            if marker < 0:
                break
            pos = marker
            while pos < len(data) and data[pos] == 0xFF:
                pos += 1
            if pos >= len(data):
                break
            code = data[pos]
            pos += 1
            if code == 0x00 or 0xD0 <= code <= 0xD7:
                continue
            in_scan = False
        else:
            if data[pos] != 0xFF:
                raise CaptureError("invalid JPEG marker stream")
            while pos < len(data) and data[pos] == 0xFF:
                pos += 1
            if pos >= len(data):
                break
            code = data[pos]
            pos += 1
        if code == 0xD9:
            if pos != len(data) or dimensions is None:
                raise CaptureError("JPEG has trailing/polyglot bytes or no dimensions")
            return dimensions
        if code in {0xD8, 0x01} or 0xD0 <= code <= 0xD7:
            continue
        if pos + 2 > len(data):
            raise CaptureError("truncated JPEG segment")
        length = struct.unpack(">H", data[pos:pos + 2])[0]
        if length < 2 or pos + length > len(data):
            raise CaptureError("invalid JPEG segment length")
        payload = data[pos + 2:pos + length]
        if code in sof:
            if len(payload) < 6:
                raise CaptureError("truncated JPEG SOF")
            height, width = struct.unpack(">HH", payload[1:5])
            dimensions = (width, height)
        pos += length
        if code == 0xDA:
            in_scan = True
    raise CaptureError("JPEG EOI missing")

def validate_image(data: bytes) -> dict[str, Any]:
    if not data or len(data) > MAX_BYTES:
        raise CaptureError(f"image size must be 1..{MAX_BYTES} bytes")
    if data.startswith(PNG_SIG):
        fmt, dimensions = "PNG", _strict_png(data)
    elif data.startswith(b"\xff\xd8"):
        fmt, dimensions = "JPEG", _strict_jpeg(data)
    else:
        raise CaptureError("only real PNG or JPEG input is accepted")
    Image.MAX_IMAGE_PIXELS = MAX_PIXELS
    try:
        from io import BytesIO
        with Image.open(BytesIO(data)) as image:
            image.load()
            decoded = (image.width, image.height)
            mode = image.mode
            decoded_format = image.format
    except (UnidentifiedImageError, OSError, ValueError) as exc:
        raise CaptureError("image decoder rejected input") from exc
    width, height = dimensions
    if decoded_format != fmt or decoded != dimensions:
        raise CaptureError("image format/dimension mismatch")
    if width < 1 or height < 1 or width > MAX_DIMENSION or height > MAX_DIMENSION or width * height > MAX_PIXELS:
        raise CaptureError("image dimensions exceed policy")
    if mode not in ALLOWED_MODES:
        raise CaptureError(f"image mode is not allowed: {mode}")
    return {"format": fmt, "width": width, "height": height, "mode": mode, "bytes": len(data), "sha256": digest(data)}

def read_source(path: Path) -> tuple[bytes, int]:
    if not path.is_absolute():
        raise CaptureError("image path must be an absolute Linux filesystem path")
    flags = os.O_RDONLY | os.O_CLOEXEC | os.O_NOFOLLOW
    try:
        fd = os.open(path, flags)
    except OSError as exc:
        raise CaptureError(f"cannot securely open image: {exc.strerror}") from exc
    try:
        before = os.fstat(fd)
        mode = stat.S_IMODE(before.st_mode)
        if not stat.S_ISREG(before.st_mode) or before.st_nlink != 1:
            raise CaptureError("image must be a single-link regular file")
        if mode & 0o111 or mode & 0o022:
            raise CaptureError("image source must not be executable or group/world-writable")
        if before.st_size < 1 or before.st_size > MAX_BYTES:
            raise CaptureError(f"image size must be 1..{MAX_BYTES} bytes")
        parts, remaining = [], MAX_BYTES + 1
        while remaining:
            chunk = os.read(fd, min(1024 * 1024, remaining))
            if not chunk:
                break
            parts.append(chunk)
            remaining -= len(chunk)
        data = b"".join(parts)
        after = os.fstat(fd)
        if (before.st_dev, before.st_ino, before.st_size, before.st_mtime_ns) != (after.st_dev, after.st_ino, after.st_size, after.st_mtime_ns):
            raise CaptureError("image changed while being read")
        if len(data) != before.st_size:
            raise CaptureError("incomplete image read")
        return data, mode
    finally:
        os.close(fd)

def capture(path: Path, attestation: str, *, now: datetime | None = None) -> dict[str, Any]:
    if attestation != ATTESTATION:
        raise CaptureError("exact operator attestation is required")
    if not path.is_absolute():
        raise CaptureError("image path must be an absolute Linux filesystem path")
    contract, contract_raw = load_contract()
    if RECEIPT.exists() or RECEIPT.is_symlink() or VAULT.exists() or VAULT.is_symlink():
        raise CaptureError("capture is one-use and has already been consumed")
    data, source_mode = read_source(path)
    image = validate_image(data)
    extension = ".png" if image["format"] == "PNG" else ".jpg"
    VAULT.mkdir(mode=0o700)
    os.chmod(VAULT, 0o700)
    target = VAULT / f"customer-surface{extension}"
    atomic_write(target, data)
    if stat.S_IMODE(target.stat().st_mode) != 0o600 or target.read_bytes() != data:
        target.unlink(missing_ok=True)
        raise CaptureError("vault copy verification failed")
    captured = (now or datetime.now(timezone.utc)).astimezone(timezone.utc).isoformat().replace("+00:00", "Z")
    receipt = {
        "schema": "task26-customer-surface-capture-receipt-v1",
        "status": "PASS_CAPTURED_ONCE",
        "binding": contract["binding"],
        "contract_sha256": digest(contract_raw),
        "capture": {
            "captured_at_utc": captured,
            "image_bytes": image["bytes"],
            "image_format": image["format"],
            "image_height": image["height"],
            "image_mode": image["mode"],
            "image_sha256": image["sha256"],
            "image_width": image["width"],
            "source_mode": f"{source_mode:04o}",
            "vault_mode": "0600",
            "vault_path": str(target.relative_to(ROOT)),
        },
        "operator_attestation": {
            "exactly_one_visible_message": True,
            "matches_durable_delivered_text": True,
            "token": ATTESTATION,
        },
        "privacy": {
            "chat_attachment_is_filesystem_evidence": False,
            "customer_answers_in_json": False,
            "image_is_protected_evidence": True,
            "ocr_performed": False,
            "raw_delivered_text_in_json": False,
        },
    }
    try:
        atomic_write(RECEIPT, canonical(receipt))
        _fsync_directory(ROOT)
    except Exception:
        target.unlink(missing_ok=True)
        try:
            VAULT.rmdir()
        except OSError:
            pass
        raise
    return receipt

def dry_run() -> dict[str, Any]:
    contract, raw = load_contract()
    return {
        "schema": "task26-customer-surface-capture-dry-run-v1",
        "status": "PASS_AWAITING_FINAL_RUN_IMAGE",
        "contract_sha256": digest(raw),
        "candidate": contract["binding"]["candidate"],
        "live_actions": 0,
        "network_actions": 0,
        "ocr_actions": 0,
        "vault_present": VAULT.exists(),
        "receipt_present": RECEIPT.exists(),
    }

def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("image", nargs="?", type=Path, help="absolute Linux path to one PNG/JPEG")
    parser.add_argument("--attest", help="exact attestation token")
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args(argv)
    try:
        if args.dry_run:
            if args.image is not None or args.attest is not None:
                raise CaptureError("dry-run accepts no image or attestation")
            result = dry_run()
        else:
            if args.image is None:
                raise CaptureError("one image path is required")
            result = capture(args.image, args.attest or "")
        sys.stdout.buffer.write(canonical(result))
        return 0
    except (CaptureError, OSError, json.JSONDecodeError) as exc:
        sys.stderr.write(f"capture rejected: {exc}\n")
        return 2

if __name__ == "__main__":
    raise SystemExit(main())