#!/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. Frames are registered while you scan;
after the stop press the page is assembled from all of them at once (see
myscan/pagegraph.py) and saved as <outdir>/scan-YYYYmmdd-HHMMSS.png. The mouse goes
back to mouse mode between scans (and on Ctrl-C).

With --view, a browser viewer (myscan/webview.py) shows the camera and the page as it
is being stitched, and has Start/Stop buttons that work like the scan button."""
import argparse
import csv
import os
import queue
import sys
import threading
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.live import LiveSession  # noqa: E402
from myscan.pagegraph import PageGraph  # noqa: E402
from myscan.xu import XU  # noqa: E402


def wait_press(x, viewer=None, poll=0.05):
    """Block until the scan button goes down then up again, or Start is clicked."""
    prev = 1
    while True:
        if viewer and viewer.take_start():
            return
        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, undistort=None, gains=None, viewer=None):
    """Capture until the button is pressed again. Frames are registered (and saved) on a
    worker thread so capture never waits; returns (PageGraph or None, LiveSession or None)."""
    g, live, rows, work, saves = None, None, [], queue.Queue(), queue.Queue()

    def worker():
        while (item := work.get()) is not None:
            n, im, v = item
            f = g.add(im, v)
            if live:
                live.on_frame(f)

    def saver():
        while (item := saves.get()) is not None:
            cv2.imwrite(os.path.join(frames_dir, f"f{item[0]:05d}.png"), item[1])
    th = None
    sv = threading.Thread(target=saver, daemon=True) if frames_dir else None
    if sv:
        sv.start()
    if viewer:
        viewer.take_stop()          # ignore a Stop clicked before the scan started
    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 g is None:
                g = PageGraph(im.shape, flat=flat, undistort=undistort)
                if viewer:
                    live = LiveSession(g, viewer)
                th = threading.Thread(target=worker, daemon=True)
                th.start()
            work.put((n, im, rec.v))
            if sv:
                saves.put((n, im))
            if live:
                live.raw = im
            if frames_dir:
                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 viewer and viewer.take_stop():
                break
            if t >= a.max_time:
                print("max scan time reached", flush=True)
                break
            if viewer:
                viewer.set_status(t=t, frames=n + 1, backlog=work.qsize())
            if t - last >= 1:
                print(f"\r  {t:5.1f}s  {n + 1} frames  ({g.n} registered)  E6={s.e6}", end="", flush=True)
                last = t
    print()
    if viewer:
        viewer.set_status(state="assembling", message="finishing registration")
    if th:
        work.put(None)
        th.join()
    if sv:
        saves.put(None)
        sv.join()
    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 g, live


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-undistort", action="store_true", help="skip the EEPROM lens correction")
    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")
    ap.add_argument("--view", type=int, nargs="?", const=8000, metavar="PORT",
                    help="serve a live browser view on PORT (default 8000)")
    a = ap.parse_args()
    os.makedirs(a.outdir, exist_ok=True)
    viewer = None
    if a.view is not None:
        from myscan.webview import Viewer
        viewer = Viewer(a.view)
        print("live view: " + "  ".join(viewer.urls()), flush=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)
    undistort = None if a.no_undistort else calib.Undistort(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)
                if viewer:
                    viewer.set_status(state="idle", message="press the scan button or Start")
                wait_press(x, viewer)
            first = False
            print("scanning - press the scan button again to stop", flush=True)
            if viewer:
                viewer.set_status(state="scanning", message="press the scan button or Stop to finish")
            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)
            g, live = scan_page(x, a, frames_dir, flat, undistort, gains, viewer)
            if live:
                live.stop_solving()
                live.close()
            if g is None or g.n < 2:
                print("no frames captured")
                if viewer:
                    viewer.set_status(state="idle", message="no frames captured")
            else:
                print("assembling the page...", flush=True)
                t0 = time.time()
                if viewer:
                    viewer.set_status(message="solving")
                page = g.image(white_balance=not a.no_wb)
                path = os.path.join(a.outdir, name + ".png")
                cv2.imwrite(path, page)
                print(f"saved {path}  ({time.time() - t0:.1f} s; {g.stats})", flush=True)
                if viewer:
                    s = min(1.0, 1600 / max(page.shape[:2]))
                    viewer.publish("page", cv2.resize(page, None, fx=s, fy=s, interpolation=cv2.INTER_AREA))
                    viewer.set_result(path)
                    viewer.set_status(state="done", message=f"saved {path}", backlog=0)
            if a.once:
                break
    except KeyboardInterrupt:
        print("\nbye")


if __name__ == "__main__":
    main()