#!/usr/bin/env python3
"""TARGAFILE screen grab for DyzineFrame (shared P3 serial port).

Press S in DyzineFrame on the Pi to save under ~/p3-grabs/ then rclone-sync
to Dropbox Screenshots/P3 Grabs when configured.

Layout (local + Dropbox):
  ~/p3-grabs/pos5/T340_pos5_tidD249_20260805T205507Z.tga
  ~/p3-grabs/pos5/T340_pos5_tidD249_20260805T205507Z_raw.bin
  ~/p3-grabs/LATEST.txt   # relative path: pos5/....tga

Pos from FIXTURE_POS / fixture-bridge.env / hostname topboxN.
Local TID from LOCAL_TID / P3_GRAB_TID / FIXTURE_TID / fixture-bridge.env when set
(unknown if not configured).
"""
from __future__ import annotations

import os
import re
import socket
import struct
import subprocess
import threading
import time
import zlib
from datetime import datetime, timezone
from pathlib import Path
from typing import Callable

TGA_SIG = bytes([0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2, 0, 1, 0x10, 0])
PAYLOAD = 512 * 256 * 2
TOTAL = 18 + PAYLOAD
DEFAULT_OUT = Path.home() / "p3-grabs"
DEFAULT_PREFIX = "T340"
DEFAULT_WIDTH = 4

_ENV_POS_KEYS = ("FIXTURE_POS", "P3_GRAB_POS", "POS")
_ENV_TID_KEYS = ("LOCAL_TID", "P3_GRAB_TID", "FIXTURE_TID", "ETHERGOLD_LOCAL_TID")
_ENV_FILE_CANDIDATES = (
    Path("/etc/fixture/fixture-pos.env"),
    Path("/etc/fixture/fixture-bridge.env"),
    Path("/etc/ethergold/fixture-bridge.env"),
    Path.home() / "EtherGold" / "fixture-bridge.env",
    Path.home() / "EtherGold" / "fixture-pos.env",
)


def _read_env_file(path: Path) -> dict[str, str]:
    out: dict[str, str] = {}
    if not path.is_file():
        return out
    try:
        text = path.read_text(encoding="utf-8", errors="replace")
    except OSError:
        return out
    for line in text.splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        if line.startswith("export "):
            line = line[7:].strip()
        if "=" not in line:
            continue
        key, val = line.split("=", 1)
        key = key.strip()
        val = val.strip().strip("'\"")
        if key:
            out[key] = val
    return out


def _merged_env_file() -> dict[str, str]:
    merged: dict[str, str] = {}
    for path in _ENV_FILE_CANDIDATES:
        merged.update(_read_env_file(path))
    return merged


def resolve_fixture_pos() -> str:
    """Seat id as digits, e.g. '5'. Falls back to 'unk'."""
    for key in _ENV_POS_KEYS:
        raw = (os.environ.get(key) or "").strip()
        if raw.isdigit() and 1 <= int(raw) <= 99:
            return str(int(raw))
    file_env = _merged_env_file()
    for key in _ENV_POS_KEYS:
        raw = (file_env.get(key) or "").strip()
        if raw.isdigit() and 1 <= int(raw) <= 99:
            return str(int(raw))
    host = socket.gethostname().split(".")[0].lower()
    m = re.search(r"topbox(\d+)", host)
    if m:
        return str(int(m.group(1)))
    m = re.search(r"(?:^|[-_])pos(\d+)$", host)
    if m:
        return str(int(m.group(1)))
    return "unk"


def resolve_local_tid() -> str:
    """Local terminal TID hex upper, e.g. 'D249'. Falls back to 'unk'."""
    for key in _ENV_TID_KEYS:
        raw = (os.environ.get(key) or "").strip().upper()
        if re.fullmatch(r"[0-9A-F]{1,4}", raw):
            return raw.zfill(4) if len(raw) <= 4 else raw
    file_env = _merged_env_file()
    for key in _ENV_TID_KEYS:
        raw = (file_env.get(key) or "").strip().upper()
        if re.fullmatch(r"[0-9A-F]{1,4}", raw):
            return raw.zfill(4) if len(raw) <= 4 else raw
    return "unk"


def grab_utc_stamp() -> str:
    """Compact UTC for filenames: 20260805T205507Z."""
    return datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")


def make_grab_basename(
    *,
    prefix: str = DEFAULT_PREFIX,
    pos: str | None = None,
    tid: str | None = None,
    stamp: str | None = None,
) -> str:
    """Build base name without extension: T340_pos5_tidD249_20260805T205507Z."""
    pos_s = (pos if pos is not None else resolve_fixture_pos()).strip() or "unk"
    tid_s = (tid if tid is not None else resolve_local_tid()).strip().upper() or "unk"
    stamp_s = (stamp if stamp is not None else grab_utc_stamp()).strip() or grab_utc_stamp()
    # Sanitize path-hostile chars.
    pos_s = re.sub(r"[^0-9A-Za-z]", "", pos_s) or "unk"
    tid_s = re.sub(r"[^0-9A-Fa-f]", "", tid_s).upper() or "unk"
    stamp_s = re.sub(r"[^0-9A-Za-z]", "", stamp_s) or grab_utc_stamp()
    return f"{prefix}_pos{pos_s}_tid{tid_s}_{stamp_s}"


def grab_subdir_for_pos(pos: str | None = None) -> str:
    pos_s = (pos if pos is not None else resolve_fixture_pos()).strip() or "unk"
    pos_s = re.sub(r"[^0-9A-Za-z]", "", pos_s) or "unk"
    return f"pos{pos_s}"


def next_basename(out_dir: Path, *, prefix: str = DEFAULT_PREFIX, width: int = DEFAULT_WIDTH) -> str:
    """Legacy numeric sequence (kept for callers that still pass prefix only)."""
    pat = re.compile(rf"^{re.escape(prefix)}_(\d+)\.tga$", re.I)
    highest = 0
    for path in out_dir.glob(f"{prefix}_*.tga"):
        match = pat.match(path.name)
        if match:
            highest = max(highest, int(match.group(1)))
    counter = out_dir / f"{prefix}.counter"
    if counter.is_file():
        try:
            highest = max(highest, int(counter.read_text(encoding="ascii").strip()))
        except ValueError:
            pass
    return f"{prefix}_{highest + 1:0{width}d}"


def fetch_targa_on_serial(ser, *, wait_s: float = 150.0) -> bytes:
    ser.reset_input_buffer()
    ser.write(b"TARGAFILE\r")
    ser.flush()
    deadline = time.time() + wait_s
    chunks: list[bytes] = []
    last = 0.0
    while time.time() < deadline:
        waiting = getattr(ser, "in_waiting", 0) or 0
        if waiting:
            chunks.append(ser.read(waiting))
            last = time.time()
        elif chunks and last and (time.time() - last) >= 10.0:
            break
        time.sleep(0.01)
    return b"".join(chunks)


def build_tga(raw: bytes) -> bytes:
    idx = raw.find(TGA_SIG)
    if idx >= 0 and len(raw) >= idx + TOTAL:
        return raw[idx : idx + TOTAL]
    if idx >= 0:
        pixels = raw[idx + 18 :]
        pixels = pixels.ljust(PAYLOAD, b"\x00")[:PAYLOAD]
        return TGA_SIG + pixels
    if len(raw) >= PAYLOAD:
        return TGA_SIG + raw[-PAYLOAD:]
    return TGA_SIG + raw.ljust(PAYLOAD, b"\x00")[:PAYLOAD]


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


def rgb_to_png_bytes(width: int, height: int, rgb: bytes) -> bytes:
    """Uncompressed RGB888 → PNG (stdlib; no Pillow)."""
    if width < 1 or height < 1:
        raise ValueError("PNG size must be positive")
    row = width * 3
    if len(rgb) < row * height:
        raise ValueError("RGB buffer too short")
    raw = bytearray()
    for y in range(height):
        raw.append(0)
        raw.extend(rgb[y * row : (y + 1) * row])
    return (
        b"\x89PNG\r\n\x1a\n"
        + _png_chunk(b"IHDR", struct.pack(">IIBBBBB", width, height, 8, 2, 0, 0, 0))
        + _png_chunk(b"IDAT", zlib.compress(bytes(raw), 9))
        + _png_chunk(b"IEND", b"")
    )


def tga16_to_png_bytes(tga: bytes) -> bytes:
    """T340+ TARGAFILE 16-bit TGA (BGR555) → PNG. Origin bit 5 = top-left."""
    if len(tga) < 18:
        raise ValueError("TGA shorter than header")
    width = tga[12] | (tga[13] << 8)
    height = tga[14] | (tga[15] << 8)
    bpp = tga[16]
    desc = tga[17]
    if bpp != 16:
        raise ValueError(f"need 16-bit TGA, got bpp={bpp}")
    if width < 1 or height < 1 or width > 4096 or height > 4096:
        raise ValueError(f"implausible TGA size {width}x{height}")
    n = width * height
    payload = tga[18 : 18 + n * 2]
    if len(payload) < n * 2:
        payload = payload.ljust(n * 2, b"\x00")
    origin_top = bool(desc & 0x20)
    rgb = bytearray(n * 3)
    for i in range(n):
        pix = payload[i * 2] | (payload[i * 2 + 1] << 8)
        blue = (pix & 0x1F) * 255 // 31
        green = ((pix >> 5) & 0x1F) * 255 // 31
        red = ((pix >> 10) & 0x1F) * 255 // 31
        src_y = i // width
        src_x = i % width
        dst_y = src_y if origin_top else (height - 1 - src_y)
        o = (dst_y * width + src_x) * 3
        rgb[o] = red
        rgb[o + 1] = green
        rgb[o + 2] = blue
    return rgb_to_png_bytes(width, height, bytes(rgb))


def save_targa_capture(
    raw: bytes,
    out_dir: Path = DEFAULT_OUT,
    *,
    basename: str = "",
    prefix: str = DEFAULT_PREFIX,
    pos: str | None = None,
    tid: str | None = None,
    stamp: str | None = None,
    segregate_by_pos: bool = True,
) -> tuple[Path, Path]:
    """Write TGA + raw under out_dir[/posN/]. Returns (tga_path, raw_path)."""
    out_dir = Path(out_dir)
    pos_s = pos if pos is not None else resolve_fixture_pos()
    tid_s = tid if tid is not None else resolve_local_tid()
    stamp_s = stamp if stamp is not None else grab_utc_stamp()

    if segregate_by_pos:
        dest_dir = out_dir / grab_subdir_for_pos(pos_s)
    else:
        dest_dir = out_dir
    dest_dir.mkdir(parents=True, exist_ok=True)

    base = basename.strip() or make_grab_basename(
        prefix=prefix, pos=pos_s, tid=tid_s, stamp=stamp_s
    )
    # Avoid clobber if same-second press.
    candidate = dest_dir / f"{base}.tga"
    if candidate.exists():
        n = 2
        while (dest_dir / f"{base}_{n}.tga").exists():
            n += 1
        base = f"{base}_{n}"

    raw_path = dest_dir / f"{base}_raw.bin"
    tga_path = dest_dir / f"{base}.tga"

    if len(raw) < 1000:
        raise RuntimeError(f"only {len(raw)} bytes from TARGAFILE — terminal enabled?")

    raw_path.write_bytes(raw)
    tga = build_tga(raw)
    tga_path.write_bytes(tga)
    png_path = dest_dir / f"{base}.png"
    try:
        png_path.write_bytes(tga16_to_png_bytes(tga))
    except Exception as exc:
        print(f"P3 grab PNG convert failed: {exc}", flush=True)

    try:
        rel = tga_path.relative_to(out_dir).as_posix()
    except ValueError:
        rel = tga_path.name
    (out_dir / "LATEST.txt").write_text(f"{rel}\n", encoding="ascii")

    # Optional counter still tracks how many grabs this tree has taken (debug).
    counter = out_dir / f"{prefix}.counter"
    try:
        prev = int(counter.read_text(encoding="ascii").strip()) if counter.is_file() else 0
    except ValueError:
        prev = 0
    counter.write_text(f"{prev + 1}\n", encoding="ascii")

    return tga_path, raw_path


def find_sync_script() -> Path | None:
    home = Path.home()
    candidates = [
        home / "EtherGold" / "sync" / "p3_grab_sync.sh",
        home / "EtherGold" / "4-Board Test Fixture" / "sync" / "p3_grab_sync.sh",
        home / "EtherGold" / "p3-grab" / "p3_grab_sync.sh",
    ]
    for path in candidates:
        if path.is_file():
            return path
    return None


def run_p3_grab_sync() -> None:
    script = find_sync_script()
    if script is None:
        print("P3 grab: no p3_grab_sync.sh — files stay on Pi only", flush=True)
        return
    subprocess.run(["bash", str(script)], check=False)


_LISTENER_STARTED = False
_FLAG_LISTENER_STARTED = False


def _p3_grab_flag_worker(on_request: Callable[[], None], flag_path: Path) -> None:
    print(f"P3 grab: HTTP flag watcher {flag_path}", flush=True)
    while True:
        try:
            if flag_path.is_file():
                try:
                    flag_path.unlink()
                except OSError:
                    time.sleep(0.2)
                    continue
                try:
                    on_request()
                except Exception as exc:
                    print(f"P3 grab flag error: {exc}", flush=True)
        except Exception:
            pass
        time.sleep(0.4)


def start_p3_grab_flag_listener(on_request: Callable[[], None], flag_path: str | Path) -> bool:
    """EtherGold HTTP writes run/p3_grab.request; DF consumes it (same TARGAFILE as S)."""
    global _FLAG_LISTENER_STARTED
    if _FLAG_LISTENER_STARTED:
        return True
    path = Path(flag_path)
    threading.Thread(
        target=_p3_grab_flag_worker,
        args=(on_request, path),
        daemon=True,
        name="p3-grab-flag",
    ).start()
    _FLAG_LISTENER_STARTED = True
    return True


def _find_keyboard_device():
    from evdev import InputDevice, ecodes, list_devices

    for path in list_devices():
        candidate = InputDevice(path)
        keys = candidate.capabilities().get(ecodes.EV_KEY, [])
        if ecodes.KEY_S not in keys:
            continue
        name = candidate.name.lower()
        if "keyboard" in name or "kbd" in candidate.phys:
            return candidate
    for path in list_devices():
        candidate = InputDevice(path)
        keys = candidate.capabilities().get(ecodes.EV_KEY, [])
        if ecodes.KEY_S in keys:
            return candidate
    return None


def _evdev_grab_worker(on_s: Callable[[], None]) -> None:
    from evdev import ecodes

    dev = _find_keyboard_device()
    if dev is None:
        print("P3 grab: no keyboard device — S key only via pygame window focus", flush=True)
        return

    print(f"P3 grab: S key listener on {dev.path} ({dev.name})", flush=True)
    last = 0.0
    debounce_s = 2.0
    for event in dev.read_loop():
        if event.type != ecodes.EV_KEY or event.code != ecodes.KEY_S or event.value != 1:
            continue
        now = time.monotonic()
        if now - last < debounce_s:
            continue
        last = now
        try:
            on_s()
        except Exception as exc:
            print(f"P3 grab hotkey error: {exc}", flush=True)


def start_evdev_grab_listener(on_s: Callable[[], None]) -> bool:
    """Physical keyboard S -> callback; works regardless of pygame loop state."""
    global _LISTENER_STARTED
    if _LISTENER_STARTED:
        return True
    try:
        import evdev  # noqa: F401
    except ImportError:
        print("P3 grab: install python3-evdev for S key at all times", flush=True)
        return False

    threading.Thread(
        target=_evdev_grab_worker,
        args=(on_s,),
        daemon=True,
        name="p3-grab-evdev",
    ).start()
    _LISTENER_STARTED = True
    return True