#!/usr/bin/env python3
"""Scan with the scan button, like the original app: press it to start, glide the
mouse over the page, press it again to stop. The page is stitched live and saved
as <outdir>/scan-YYYYmmdd-HHMMSS.png; the mouse goes back to mouse mode between
scans (and on Ctrl-C)."""
import argparse
import csv
import os
import sys
import time

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

from myscan.session import ScanSession, quiet_libjpeg  # noqa: E402

quiet_libjpeg()
import cv2  # noqa: E402

from myscan import calib  # noqa: E402
from myscan.stitch import Stitcher  # noqa: E402
from myscan.xu import XU  # noqa: E402


def wait_press(x, poll=0.05):
    """Block until the scan button goes down then up again."""
    prev = 1
    while True:
        b = x.scan_button()
        if b and not prev:
            break
        prev = b
        time.sleep(poll)
    while x.scan_button():
        time.sleep(poll)


def scan_page(x, a, frames_dir=None, flat=None, gains=None):
    st, rows = None, []
    with ScanSession(a.video, a.target, a.exposure, xu=x, gains=gains) as s:
        down, last = 0, 0.0
        for t, n, im, rec in s.frames(color=not a.gray):
            if st is None:
                st = Stitcher(im.shape, refine=not a.no_refine, flat=flat)
            st.add(im, rec.v)
            if frames_dir:
                cv2.imwrite(os.path.join(frames_dir, f"f{n:05d}.png"), im)
                rows.append([round(t, 4), n, rec.counter, *rec.v, rec.b14, rec.b15, rec.b16, rec.flags, s.e6,
                             round(float(im.mean()), 1), rec.raw.hex()])
            # Stop on a press: two "down" reads in a row (debounce), ignoring the first
            # 0.7 s so a bouncing start press can't end the scan.
            if n % 2 == 0:
                down = down + 1 if x.scan_button() else 0
                if down >= 2 and t > 0.7:
                    break
            if t >= a.max_time:
                print("max scan time reached", flush=True)
                break
            if t - last >= 1:
                print(f"\r  {t:5.1f}s  {st.n} frames  canvas {st.size[0]}x{st.size[1]}  E6={s.e6}",
                      end="", flush=True)
                last = t
    print()
    if frames_dir:   # same layout as tools/record.py, so tools/stitch.py can replay it
        with open(os.path.join(frames_dir, "records.csv"), "w", newline="") as fh:
            w = csv.writer(fh)
            w.writerow(["t", "frame", "counter", "v0", "v1", "v2", "v3", "v4", "v5",
                        "b14", "b15", "b16", "flags", "e6", "brightness", "raw"])
            w.writerows(rows)
    return st


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("-o", "--outdir", default="scans")
    ap.add_argument("--once", action="store_true", help="exit after one page")
    ap.add_argument("--now", action="store_true", help="start the first scan immediately")
    ap.add_argument("--video", default="/dev/video4")
    ap.add_argument("--target", type=float, default=200, help="auto-exposure target for paper (95th pct of brightest channel)")
    ap.add_argument("--exposure", type=int, help="fixed E6 value instead of auto-exposure")
    ap.add_argument("--gray", action="store_true")
    ap.add_argument("--no-wb", action="store_true", help="skip white-patch balancing")
    ap.add_argument("--no-flat", action="store_true", help="skip the EEPROM flat-field correction")
    ap.add_argument("--no-refine", action="store_true", help="lasers only, no image registration")
    ap.add_argument("--save-frames", action="store_true", help="also keep raw frames + records.csv per scan")
    ap.add_argument("--max-time", type=float, default=180, help="stop a scan after this many seconds")
    a = ap.parse_args()
    os.makedirs(a.outdir, exist_ok=True)
    x = XU(a.video)
    print("reading calibration from the mouse's EEPROM (cached after the first time)", flush=True)
    cal = calib.load(x)
    flat = None if a.no_flat else calib.FlatField(cal)
    gains = cal["gains_rgb"]
    first = True
    try:
        while True:
            if not (first and a.now):
                print("press the scan button to start (Ctrl-C to quit)", flush=True)
                wait_press(x)
            first = False
            print("scanning - press the scan button again to stop", flush=True)
            name = time.strftime("scan-%Y%m%d-%H%M%S")
            frames_dir = os.path.join(a.outdir, name) if a.save_frames else None
            if frames_dir:
                os.makedirs(frames_dir)
            st = scan_page(x, a, frames_dir, flat, gains)
            if st is None or st.n < 2:
                print("no frames captured")
            else:
                path = os.path.join(a.outdir, name + ".png")
                cv2.imwrite(path, st.image(white_balance=not a.no_wb))
                print(f"saved {path}  ({st.n} frames, {st.stats})", flush=True)
            if a.once:
                break
    except KeyboardInterrupt:
        print("\nbye")


if __name__ == "__main__":
    main()