atm/scripts/repro_ss_resume.py

r"""Reproduce /ss and /resume Telegram screenshot pipelines for offline inspection.

Brings the `m2d` window to front (Win32 trick — same as live loop), captures via mss
using the active config's chart_window_region, then runs the EXACT same annotators
that /ss and /resume use:
  - _save_inspect_frame  → /ss path (top-3 dots per detected strip + caption)
  - _save_annotated_frame → /resume path (cyan rect on dot_roi/sub_roi)

Outputs are saved under logs/repro/ alongside a JSON summary of detections.

Usage:
  .\.venv\Scripts\python.exe scripts\repro_ss_resume.py
  .\.venv\Scripts\python.exe scripts\repro_ss_resume.py --no-focus
"""
from __future__ import annotations

import argparse
import json
import sys
import time
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
SRC = ROOT / "src"
if str(SRC) not in sys.path:
    sys.path.insert(0, str(SRC))

import cv2  # noqa: E402
import numpy as np  # noqa: E402

from atm.config import Config  # noqa: E402
from atm.layout import detect_strips  # noqa: E402
from atm.main import (  # noqa: E402
    _focus_window_by_title,
    _format_inspect_caption,
    _save_annotated_frame,
    _save_inspect_frame,
)
from atm.vision import crop_roi  # noqa: E402


def _capture_via_region(cfg) -> np.ndarray | None:
    import mss  # type: ignore[import-untyped]

    reg = cfg.chart_window_region
    if reg is None:
        # Fallback: grab full primary monitor
        with mss.mss() as sct:
            mon = sct.monitors[1]
            img = sct.grab(mon)
            return cv2.cvtColor(np.array(img), cv2.COLOR_BGRA2BGR)
    with mss.mss() as sct:
        mon = {"top": reg.y, "left": reg.x, "width": reg.w, "height": reg.h}
        img = sct.grab(mon)
        return cv2.cvtColor(np.array(img), cv2.COLOR_BGRA2BGR)


def main() -> int:
    p = argparse.ArgumentParser()
    p.add_argument("--no-focus", action="store_true", help="Skip Win32 focus call.")
    p.add_argument("--delay", type=float, default=0.5, help="Seconds to sleep after focus.")
    p.add_argument(
        "--out",
        type=Path,
        default=ROOT / "logs" / "repro",
        help="Output directory for annotated frames + JSON.",
    )
    args = p.parse_args()

    cfg = Config.load_current(ROOT / "configs")
    args.out.mkdir(parents=True, exist_ok=True)

    # 1) Focus
    if not args.no_focus and cfg.window_title:
        title = _focus_window_by_title(cfg.window_title)
        print(f"[focus] needle={cfg.window_title!r} -> {title!r}")
    if args.delay > 0:
        time.sleep(args.delay)

    # 2) Capture
    frame = _capture_via_region(cfg)
    if frame is None:
        print("[capture] FAILED — no frame")
        return 1
    print(f"[capture] frame shape={frame.shape}")

    now = time.time()
    ts_str = time.strftime("%Y%m%d_%H%M%S")

    # Save raw too so we can hand-inspect what mss actually grabbed
    raw_path = args.out / f"{ts_str}_raw.png"
    cv2.imwrite(str(raw_path), frame)
    print(f"[raw]    {raw_path}")

    # 3) Detect strips (same logic as live multi-chart split)
    strip_h = cfg.dot_roi.h
    min_strip_px = max(150, strip_h * 8)
    min_gap_px = max(20, int(strip_h * 0.8))
    palette = {
        name: (spec.rgb, spec.tolerance)
        for name, spec in cfg.colors.items()
        if name != "background"
    }
    full_dot_crop = crop_roi(frame, cfg.dot_roi)
    raw_strips = detect_strips(full_dot_crop, palette, min_gap_px, min_strip_px)
    # Translate back to absolute frame coords
    from atm.config import ROI  # noqa: E402

    strips = [
        ROI(x=cfg.dot_roi.x + r.x, y=cfg.dot_roi.y + r.y, w=r.w, h=r.h)
        for r in raw_strips
    ]
    print(f"[strips] dot_roi={cfg.dot_roi}  detected={len(strips)} strips")
    for i, s in enumerate(strips):
        print(f"         strip[{i}] x={s.x} y={s.y} w={s.w} h={s.h}")

    # Also dump a copy of the full dot_roi crop for visual sanity-check
    crop_path = args.out / f"{ts_str}_dot_roi_crop.png"
    cv2.imwrite(str(crop_path), full_dot_crop)
    print(f"[crop]   {crop_path}")

    # 4) /ss inspect-annotate (top-3 per strip, FSM-pick markers)
    inspect_path, detections = _save_inspect_frame(
        frame, cfg, args.out, now, audit=None,
        strips=strips if strips else None,
    )
    caption = _format_inspect_caption(detections, cfg)
    print(f"[ss]     {inspect_path}")
    print(f"[ss]     caption:\n{caption}")

    # 5) /resume annotate (cyan rect on dot_roi or first strip)
    roi_for_resume = strips[0] if strips else cfg.dot_roi
    resume_path = _save_annotated_frame(
        frame, cfg, args.out, "resume_repro", now, audit=None, roi=roi_for_resume,
    )
    print(f"[resume] {resume_path}")

    # 6) JSON summary
    summary = {
        "ts": ts_str,
        "frame_shape": list(frame.shape),
        "dot_roi": {"x": cfg.dot_roi.x, "y": cfg.dot_roi.y, "w": cfg.dot_roi.w, "h": cfg.dot_roi.h},
        "strips": [
            {"x": s.x, "y": s.y, "w": s.w, "h": s.h} for s in strips
        ],
        "detections": [
            {
                "strip_idx": d["strip_idx"],
                "idx": d["idx"],
                "name": d["name"],
                "rgb": list(d["rgb"]),
                "distance": round(float(d["distance"]), 3),
                "confidence": round(float(d["confidence"]), 3),
                "pos_abs": list(d["pos_abs"]),
            }
            for d in detections
        ],
        "files": {
            "raw": str(raw_path),
            "dot_roi_crop": str(crop_path),
            "inspect": str(inspect_path) if inspect_path else None,
            "resume": str(resume_path) if resume_path else None,
        },
        "ss_caption": caption,
    }
    json_path = args.out / f"{ts_str}_summary.json"
    json_path.write_text(json.dumps(summary, indent=2), encoding="utf-8")
    print(f"[json]   {json_path}")
    return 0


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