#!/usr/bin/env python3
"""Publish bounded kernel-side UVC/UDC error counters for Lesavka diagnostics."""

from __future__ import annotations

import argparse
import json
import os
import pathlib
import re
import signal
import subprocess
import time
from collections import Counter, deque
from typing import Any

PATTERNS = {
    "dwc2_errors": re.compile(r"\bdwc2\b.*\b(error|failed|timeout|incomplete|shutdown|disconnect)", re.I),
    "uvc_errors": re.compile(
        r"\buvc(video)?\b.*\b(error|failed|timeout|corrupt|invalid|overflow|incomplete|non-zero status)",
        re.I,
    ),
    "udc_errors": re.compile(r"\b(udc|gadget)\b.*\b(error|failed|timeout|shutdown|disconnect|overflow)", re.I),
    "urb_errors": re.compile(r"\burb\b.*\b(error|failed|timeout|resubmit|non-zero|overflow)", re.I),
    "usb_resets": re.compile(r"\busb\b.*\b(reset|disconnect|re-enumerat)", re.I),
}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--stats-path", default=os.environ.get("LESAVKA_UVC_KERNEL_STATS_PATH", "/run/lesavka-uvc-kernel-stats.json"))
    parser.add_argument("--recent-lines", type=int, default=500)
    parser.add_argument("--flush-seconds", type=float, default=5.0)
    parser.add_argument("--self-test", action="store_true")
    return parser.parse_args()


def write_atomic(path: pathlib.Path, payload: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(f".tmp.{os.getpid()}")
    temporary.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
    temporary.replace(path)


def journal_process(lines: int, follow: bool) -> subprocess.Popen[str]:
    command = ["journalctl", "-k", "--no-pager", "-o", "json", "-n", str(max(0, lines))]
    if follow:
        command.insert(2, "-f")
    return subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)


def classify_message(message: str) -> list[str]:
    return [name for name, pattern in PATTERNS.items() if pattern.search(message)]


def self_test() -> int:
    cases = {
        "dwc2 error: incomplete request": {"dwc2_errors"},
        "uvcvideo error: Non-zero status (-75) in video completion handler": {"uvc_errors"},
        "usb 1-1: reset high-speed USB device": {"usb_resets"},
        "ordinary unrelated kernel line": set(),
    }
    for message, expected in cases.items():
        actual = set(classify_message(message))
        if actual != expected:
            raise AssertionError(f"{message!r}: expected {expected}, got {actual}")
    print("self-test: pass")
    return 0


def main() -> int:
    args = parse_args()
    if args.self_test:
        return self_test()
    path = pathlib.Path(args.stats_path)
    counters: Counter[str] = Counter()
    recent: deque[dict[str, Any]] = deque(maxlen=40)
    baseline_counters: dict[str, int] = {}
    baseline_recent: list[dict[str, Any]] = []
    started = int(time.time() * 1000)
    running = True

    def stop(_signum: int, _frame: Any) -> None:
        nonlocal running
        running = False

    signal.signal(signal.SIGTERM, stop)
    signal.signal(signal.SIGINT, stop)

    def consume(raw: str) -> None:
        try:
            entry = json.loads(raw)
        except json.JSONDecodeError:
            return
        message = str(entry.get("MESSAGE") or "")
        matched = classify_message(message)
        if not matched:
            return
        counters["matched_lines"] += 1
        counters.update(matched)
        recent.append(
            {
                "unix_us": int(entry.get("__REALTIME_TIMESTAMP") or 0),
                "categories": matched,
                "message": message[:800],
            }
        )

    def snapshot() -> dict[str, Any]:
        return {
            "schema": "lesavka.uvc-kernel-watch.v1",
            "pid": os.getpid(),
            "started_unix_ms": started,
            "updated_unix_ms": int(time.time() * 1000),
            "counters": {
                name: counters[name]
                for name in ["matched_lines", "watcher_restarts", *PATTERNS]
            },
            "recent": list(recent),
            "startup_baseline_counters": baseline_counters,
            "startup_baseline_recent": baseline_recent,
        }

    initial = journal_process(args.recent_lines, False)
    stdout, _stderr = initial.communicate(timeout=15)
    for line in stdout.splitlines():
        consume(line)
    baseline_counters = {
        name: counters[name] for name in ["matched_lines", *PATTERNS]
    }
    baseline_recent = list(recent)
    counters.clear()
    recent.clear()
    write_atomic(path, snapshot())

    follow = journal_process(0, True)
    assert follow.stdout is not None
    try:
        while running:
            line = follow.stdout.readline()
            if line:
                consume(line)
                write_atomic(path, snapshot())
            elif follow.poll() is not None:
                counters["watcher_restarts"] += 1
                time.sleep(max(0.1, args.flush_seconds))
                follow = journal_process(0, True)
                assert follow.stdout is not None
            else:
                time.sleep(0.1)
    finally:
        follow.terminate()
        try:
            follow.wait(timeout=3)
        except subprocess.TimeoutExpired:
            follow.kill()
        write_atomic(path, snapshot())
    return 0


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