#!/usr/bin/env python3
"""Classify raw UVC MJPEG frames without relying on an RCT browser."""

from __future__ import annotations

import argparse
import hashlib
import io
import json
import pathlib
import sys
import time
import warnings
import zlib
from collections import Counter
from typing import Any

APP4_MAGIC = b"LSVK"
BAND_ROWS = 64
BAND_BITS = 32
BAND_CELL = 4


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("inputs", nargs="*", help="MJPEG files or directories containing .jpg/.mjpg frames")
    parser.add_argument("--output-dir", default="", help="write report.json, report.txt, and frames.jsonl")
    parser.add_argument("--idle-mean-max", type=float, default=8.0)
    parser.add_argument("--idle-stddev-max", type=float, default=4.0)
    parser.add_argument("--self-test", action="store_true")
    return parser.parse_args()


def input_files(values: list[str]) -> list[pathlib.Path]:
    files: list[pathlib.Path] = []
    for value in values:
        path = pathlib.Path(value)
        if path.is_dir():
            files.extend(
                candidate
                for candidate in sorted(path.rglob("*"))
                if candidate.suffix.lower() in {".jpg", ".jpeg", ".mjpg", ".mjpeg"}
            )
        elif path.is_file():
            files.append(path)
    return files


def split_mjpeg(data: bytes) -> list[bytes]:
    starts: list[int] = []
    cursor = 0
    while True:
        start = data.find(b"\xff\xd8", cursor)
        if start < 0:
            break
        starts.append(start)
        cursor = start + 2
    if not starts:
        return [data]
    frames: list[bytes] = []
    for index, start in enumerate(starts):
        next_start = starts[index + 1] if index + 1 < len(starts) else len(data)
        end = data.find(b"\xff\xd9", start + 2, next_start)
        frames.append(data[start : end + 2] if end >= 0 else data[start:next_start])
    return frames


def inspect_jpeg(frame: bytes) -> dict[str, Any]:
    if len(frame) < 4 or not frame.startswith(b"\xff\xd8"):
        return {"valid": False, "reason": "missing_soi", "app4": None}
    pos = 2
    saw_scan = False
    app4: dict[str, Any] | None = None
    app4_range: tuple[int, int] | None = None
    while pos < len(frame):
        if frame[pos] != 0xFF:
            return {"valid": False, "reason": "marker_alignment", "app4": app4}
        marker_start = pos
        while pos < len(frame) and frame[pos] == 0xFF:
            pos += 1
        if pos >= len(frame):
            return {"valid": False, "reason": "truncated_marker", "app4": app4}
        marker = frame[pos]
        pos += 1
        if marker == 0xD9:
            valid = saw_scan and pos == len(frame)
            result = {"valid": valid, "reason": "ok" if valid else "trailing_or_no_scan", "app4": app4}
            if app4 is not None and app4_range is not None:
                start, end = app4_range
                original = frame[:start] + frame[end:]
                app4["actual_crc32"] = zlib.crc32(original) & 0xFFFFFFFF
                app4["crc_ok"] = app4["actual_crc32"] == app4["expected_crc32"]
            return result
        if marker in {0xD8, 0x00}:
            return {"valid": False, "reason": "nested_soi_or_stuffing", "app4": app4}
        if marker == 0x01 or 0xD0 <= marker <= 0xD7:
            continue
        if pos + 2 > len(frame):
            return {"valid": False, "reason": "truncated_length", "app4": app4}
        segment_len = int.from_bytes(frame[pos : pos + 2], "big")
        segment_end = pos + segment_len
        if segment_len < 2 or segment_end > len(frame):
            return {"valid": False, "reason": "truncated_segment", "app4": app4}
        payload = frame[pos + 2 : segment_end]
        if marker == 0xE4 and payload.startswith(APP4_MAGIC) and len(payload) >= 17:
            app4 = {
                "version": payload[4],
                "sequence": int.from_bytes(payload[5:13], "big"),
                "expected_crc32": int.from_bytes(payload[13:17], "big"),
            }
            app4_range = (marker_start, segment_end)
        pos = segment_end
        if marker != 0xDA:
            continue
        saw_scan = True
        while pos < len(frame):
            if frame[pos] != 0xFF:
                pos += 1
                continue
            entropy_marker = pos
            while pos < len(frame) and frame[pos] == 0xFF:
                pos += 1
            if pos >= len(frame):
                return {"valid": False, "reason": "truncated_entropy", "app4": app4}
            if frame[pos] == 0x00 or 0xD0 <= frame[pos] <= 0xD7:
                pos += 1
                continue
            pos = entropy_marker
            break
    return {"valid": False, "reason": "missing_eoi", "app4": app4}


def decode_image(frame: bytes, allow_truncated: bool = False) -> Any | None:
    try:
        from PIL import Image, ImageFile

        previous = ImageFile.LOAD_TRUNCATED_IMAGES
        ImageFile.LOAD_TRUNCATED_IMAGES = allow_truncated
        try:
            image = Image.open(io.BytesIO(frame)).convert("L")
            image.load()
            return image
        finally:
            ImageFile.LOAD_TRUNCATED_IMAGES = previous
    except Exception:
        return None


def image_integrity_metrics(frame: bytes, allow_truncated: bool = False) -> dict[str, Any]:
    with warnings.catch_warnings(record=True) as caught:
        warnings.simplefilter("always")
        image = decode_image(frame, allow_truncated=allow_truncated)
    decode_warnings = [str(item.message) for item in caught]
    if image is None:
        return {
            "width": None,
            "height": None,
            "luma_mean": None,
            "luma_stddev": None,
            "bands": [],
            "first_bad_row": None,
            "decode_warnings": decode_warnings,
        }
    try:
        width, height = image.size
        histogram = image.resize((64, 36)).histogram()
        count = sum(histogram)
        mean = sum(value * occurrences for value, occurrences in enumerate(histogram)) / max(1, count)
        variance = (
            sum(((value - mean) ** 2) * occurrences for value, occurrences in enumerate(histogram))
            / max(1, count)
        )
        marker_width = BAND_BITS * BAND_CELL
        bands: list[dict[str, Any]] = []
        first_bad_row: int | None = None
        if width >= marker_width + 8 and height >= BAND_ROWS:
            x0 = width - marker_width - 4
            pixels = image.load()
            expected_bands = (height + BAND_ROWS - 1) // BAND_ROWS
            for band in range(expected_bands):
                y0 = band * BAND_ROWS + 4
                if y0 + BAND_CELL * 2 > height:
                    if first_bad_row is None:
                        first_bad_row = band * BAND_ROWS
                    break
                marker = 0
                confidence_sum = 0.0
                for bit in range(BAND_BITS):
                    values = [
                        pixels[x, y]
                        for y in range(y0, y0 + BAND_CELL * 2)
                        for x in range(x0 + bit * BAND_CELL, x0 + (bit + 1) * BAND_CELL)
                    ]
                    average = sum(values) / len(values)
                    if average >= 128.0:
                        marker |= 1 << bit
                    confidence_sum += abs(average - 128.0) / 127.0
                decoded_band = (marker >> 24) & 0xFF
                valid_band = decoded_band == (band & 0xFF)
                if not valid_band and first_bad_row is None:
                    first_bad_row = band * BAND_ROWS
                bands.append(
                    {
                        "band": band,
                        "row": band * BAND_ROWS,
                        "marker": marker,
                        "sequence": marker & 0x00FFFFFF,
                        "decoded_band": decoded_band,
                        "valid_band": valid_band,
                        "confidence": round(confidence_sum / BAND_BITS, 4),
                    }
                )
        return {
            "width": width,
            "height": height,
            "luma_mean": mean,
            "luma_stddev": variance**0.5,
            "bands": bands,
            "first_bad_row": first_bad_row,
            "decode_warnings": decode_warnings,
        }
    except Exception:
        return {
            "width": None,
            "height": None,
            "luma_mean": None,
            "luma_stddev": None,
            "bands": [],
            "first_bad_row": None,
            "decode_warnings": decode_warnings,
        }


def classify_frame(frame: bytes, previous_sequence: int | None, args: argparse.Namespace) -> dict[str, Any]:
    inspection = inspect_jpeg(frame)
    app4 = inspection.get("app4")
    sequence = app4.get("sequence") if isinstance(app4, dict) else None
    metrics = image_integrity_metrics(frame, allow_truncated=not inspection["valid"])
    mean = metrics["luma_mean"]
    stddev = metrics["luma_stddev"]
    bands = metrics["bands"]
    valid_band_sequences = [band["sequence"] for band in bands if band["valid_band"]]
    distinct_band_sequences = sorted(set(valid_band_sequences))
    expected_sequence = sequence & 0x00FFFFFF if sequence is not None else None
    mismatched_band = next(
        (
            band
            for band in bands
            if not band["valid_band"]
            or (expected_sequence is not None and band["sequence"] != expected_sequence)
        ),
        None,
    )
    band_splice = isinstance(app4, dict) and (
        len(distinct_band_sequences) > 1 or mismatched_band is not None
    )
    idle_black = mean is not None and stddev is not None and mean <= args.idle_mean_max and stddev <= args.idle_stddev_max
    if not inspection["valid"]:
        classification = "truncated"
    elif (isinstance(app4, dict) and not app4.get("crc_ok", False)) or band_splice:
        classification = "spliced"
    elif sequence is not None and previous_sequence is not None and sequence <= previous_sequence:
        classification = "stale_repeat"
    elif idle_black:
        classification = "idle_black"
    else:
        classification = "intact"
    first_bad_row = metrics["first_bad_row"] if classification == "truncated" else None
    seam_row = mismatched_band["row"] if classification == "spliced" and mismatched_band else None
    if classification == "truncated":
        verdict = f"truncated(first_bad_row={first_bad_row})"
    elif classification == "spliced":
        verdict = f"spliced(seqs={distinct_band_sequences}, seam_row={seam_row})"
    else:
        verdict = classification
    return {
        "classification": classification,
        "verdict": verdict,
        "bytes": len(frame),
        "sha256": hashlib.sha256(frame).hexdigest(),
        "jpeg_valid": inspection["valid"],
        "jpeg_reason": inspection["reason"],
        "sequence": sequence,
        "app4_crc_ok": app4.get("crc_ok") if isinstance(app4, dict) else None,
        "width": metrics["width"],
        "height": metrics["height"],
        "band_markers": bands,
        "band_sequences": distinct_band_sequences,
        "first_bad_row": first_bad_row,
        "seam_row": seam_row,
        "decode_warnings": metrics["decode_warnings"],
        "luma_mean": round(mean, 3) if mean is not None else None,
        "luma_stddev": round(stddev, 3) if stddev is not None else None,
    }


def add_app4(frame: bytes, sequence: int) -> bytes:
    payload = APP4_MAGIC + bytes([1]) + sequence.to_bytes(8, "big") + (zlib.crc32(frame) & 0xFFFFFFFF).to_bytes(4, "big")
    return frame[:2] + b"\xff\xe4" + (len(payload) + 2).to_bytes(2, "big") + payload + frame[2:]


def self_test(args: argparse.Namespace) -> int:
    from PIL import Image

    def marked_image(top_sequence: int, bottom_sequence: int | None = None) -> Any:
        image = Image.new("RGB", (256, 192))
        pixels = image.load()
        x0 = image.width - BAND_BITS * BAND_CELL - 4
        for y in range(image.height):
            for x in range(image.width):
                pixels[x, y] = ((x * 3) & 0xFF, (y * 5) & 0xFF, ((x + y) * 2) & 0xFF)
        for band in range((image.height + BAND_ROWS - 1) // BAND_ROWS):
            sequence = top_sequence if bottom_sequence is None or band < 2 else bottom_sequence
            marker = (sequence & 0x00FFFFFF) | ((band & 0xFF) << 24)
            for bit in range(BAND_BITS):
                value = 245 if marker & (1 << bit) else 12
                for y in range(band * BAND_ROWS + 4, band * BAND_ROWS + 12):
                    for x in range(x0 + bit * BAND_CELL, x0 + (bit + 1) * BAND_CELL):
                        pixels[x, y] = (value, value, value)
        return image

    image = marked_image(7)
    encoded = io.BytesIO()
    image.save(encoded, format="JPEG", quality=90)
    intact = add_app4(encoded.getvalue(), 7)
    spliced = bytearray(intact)
    entropy_index = max(24, len(spliced) - 12)
    while entropy_index > 24 and spliced[entropy_index] == 0xFF:
        entropy_index -= 1
    spliced[entropy_index] ^= 0x01
    visibly_spliced = io.BytesIO()
    marked_image(7, 8).save(visibly_spliced, format="JPEG", quality=95)
    visibly_spliced_frame = add_app4(visibly_spliced.getvalue(), 7)
    black = io.BytesIO()
    Image.new("RGB", (256, 192), (0, 0, 0)).save(black, format="JPEG", quality=90)

    cases = [
        (intact, None, "intact"),
        (intact, 7, "stale_repeat"),
        (bytes(spliced), 6, "spliced"),
        (visibly_spliced_frame, 6, "spliced"),
        (intact[:-9], 6, "truncated"),
        (black.getvalue(), None, "idle_black"),
    ]
    for frame, previous, expected in cases:
        actual = classify_frame(frame, previous, args)["classification"]
        if actual != expected:
            raise AssertionError(f"expected {expected}, got {actual}")
    localized = classify_frame(visibly_spliced_frame, 6, args)
    if localized["seam_row"] != 128 or localized["band_sequences"] != [7, 8]:
        raise AssertionError(f"splice localization failed: {localized}")
    print("self-test: pass")
    return 0


def main() -> int:
    args = parse_args()
    if args.self_test:
        return self_test(args)
    files = input_files(args.inputs)
    if not files:
        raise SystemExit("no MJPEG inputs found")
    records: list[dict[str, Any]] = []
    previous_sequence: int | None = None
    for path in files:
        for index, frame in enumerate(split_mjpeg(path.read_bytes())):
            record = classify_frame(frame, previous_sequence, args)
            record.update({"source": str(path), "source_frame_index": index})
            records.append(record)
            if record["sequence"] is not None:
                previous_sequence = int(record["sequence"])
    counts = Counter(record["classification"] for record in records)
    report = {
        "schema": "lesavka.uvc-mjpeg-integrity.v1",
        "generated_unix_ms": int(time.time() * 1000),
        "frames": len(records),
        "classification_counts": dict(sorted(counts.items())),
        "intact": counts["intact"],
        "truncated": counts["truncated"],
        "spliced": counts["spliced"],
        "stale_repeat": counts["stale_repeat"],
        "idle_black": counts["idle_black"],
        "verdict": "pass" if len(records) > 0 and counts["intact"] == len(records) else "fail",
    }
    text = (
        f"verdict: {report['verdict']}\nframes: {report['frames']}\n"
        f"intact: {report['intact']}\ntruncated: {report['truncated']}\n"
        f"spliced: {report['spliced']}\nstale_repeat: {report['stale_repeat']}\n"
        f"idle_black: {report['idle_black']}\n"
    )
    if args.output_dir:
        output = pathlib.Path(args.output_dir)
        output.mkdir(parents=True, exist_ok=True)
        (output / "report.json").write_text(json.dumps(report, indent=2) + "\n")
        (output / "report.txt").write_text(text)
        (output / "frames.jsonl").write_text("".join(json.dumps(record) + "\n" for record in records))
    sys.stdout.write(text)
    return 0 if report["verdict"] == "pass" else 1


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