"""Laptop receiver: read the M5Stack over USB serial or WiFi (UDP) and run the sampling SOP.

    python record.py --list-ports
    python record.py --serial /dev/tty.usbserial-XXXX --out data/2026-11-03_E1/raw.csv
    python record.py --serial COM3 --out data/2026-11-03_E1/raw.csv            (Windows)
    python record.py --udp 5005 --out data/2026-11-03_E1/raw.csv
    python record.py --serial COM3 --monitor                                   (just watch the numbers)

For every sample the program asks for the sample ID, then walks you through
    purge   (clean air, --purge seconds)  -> baseline
    sample  (sample in the box, --sample) -> steady value
    recover (clean air, --recover)        -> recovery time
and tags every reading with sample_id, phase and the second inside the phase.
Rows are appended to raw.csv, so one file can hold a whole day.
Close the UIFlow web IDE / serial monitor first: only one program can open the port.
"""
from __future__ import annotations

import argparse
import csv
import socket
import sys
import threading
import time
from datetime import datetime
from pathlib import Path

from enose.protocol import RAW_COLUMNS, Reading, parse_line


class Recorder:
    """Shared state between the reading thread (writes rows) and the SOP prompts (sets the phase)."""

    def __init__(self, out_path: Path | None) -> None:
        self.lock = threading.Lock()
        self.sample_id = ""
        self.phase = "idle"
        self.phase_start = time.monotonic()
        self.latest: Reading | None = None
        self.last_time = 0.0
        self.error = ""
        self.count = 0
        self.writer = None
        self.used_ids: set[str] = set()
        if out_path is not None:
            out_path.parent.mkdir(parents=True, exist_ok=True)
            new_file = not out_path.exists()
            if not new_file:
                with out_path.open(encoding="utf-8") as existing:
                    self.used_ids = {row["sample_id"] for row in csv.DictReader(existing) if row["sample_id"]}
            self.file = out_path.open("a", newline="", encoding="utf-8")
            self.writer = csv.writer(self.file)
            if new_file:
                self.writer.writerow(RAW_COLUMNS)

    def set_phase(self, sample_id: str, phase: str) -> None:
        with self.lock:
            self.sample_id, self.phase, self.phase_start = sample_id, phase, time.monotonic()

    def handle(self, reading: Reading) -> None:
        with self.lock:
            self.latest = reading
            self.last_time = time.monotonic()
            self.count += 1
            if self.writer is None or self.phase in ("idle", "monitor"):
                return
            second = int(time.monotonic() - self.phase_start)
            self.writer.writerow([datetime.now().isoformat(timespec="seconds"), self.sample_id, self.phase, second,
                                  *reading.as_list()])
            self.file.flush()


def serial_lines(port: str, baud: int):
    import serial  # pyserial

    with serial.Serial(port, baud, timeout=2) as connection:
        while True:
            yield connection.readline().decode("utf-8", errors="replace")


def udp_lines(port: int):
    receiver = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    receiver.bind(("0.0.0.0", port))
    print(f"等待 UDP 資料（port {port}）。M5Stack 的 LAPTOP_IP 要填這台筆電的 IP。")
    while True:
        data, _ = receiver.recvfrom(1024)
        yield data.decode("utf-8", errors="replace")


def reader(lines, recorder: Recorder) -> None:
    try:
        for line in lines:
            reading = parse_line(line)
            if reading is not None:
                recorder.handle(reading)
    except Exception as problem:  # unplugged cable, port closed ...
        recorder.error = f"{type(problem).__name__}: {problem}"


def countdown(recorder: Recorder, sample_id: str, phase: str, seconds: int, message: str) -> None:
    recorder.set_phase(sample_id, phase)
    print(f"\n[{phase}] {message}（{seconds} 秒）")
    end = time.monotonic() + seconds
    while (left := end - time.monotonic()) > 0:
        latest = recorder.latest
        status = (f"gas={latest.gas_ohm:>9.0f}Ω eth={latest.ethanol_raw:>6.0f} co2={latest.co2_ppm:>5.0f}ppm "
                  f"RH={latest.rh_pct:4.1f}%") if latest else "還沒收到資料…"
        if recorder.error:
            status = f"⚠ 連線中斷，沒有在記錄！（{recorder.error}）"
        elif latest and time.monotonic() - recorder.last_time > 3:
            status = f"⚠ 已 {time.monotonic() - recorder.last_time:.0f} 秒沒收到資料，檢查線材／WiFi"
        print(f"\r  剩 {left:5.0f} 秒  {status}", end="", flush=True)
        time.sleep(min(1.0, left))
    # Stop tagging while the box is being opened or the next prompt waits for Enter.
    recorder.set_phase(sample_id, "idle")
    print()


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    source = parser.add_mutually_exclusive_group()
    source.add_argument("--serial", help="serial port, e.g. /dev/tty.usbserial-XXXX or COM3")
    source.add_argument("--udp", type=int, help="UDP port to listen on, e.g. 5005")
    parser.add_argument("--list-ports", action="store_true")
    parser.add_argument("--baud", type=int, default=115200)
    parser.add_argument("--out", type=Path, help="raw.csv to append to")
    parser.add_argument("--monitor", action="store_true", help="only show live values, save nothing")
    parser.add_argument("--purge", type=int, default=120)
    parser.add_argument("--sample", type=int, default=240)
    parser.add_argument("--recover", type=int, default=300)
    arguments = parser.parse_args()

    if arguments.list_ports:
        from serial.tools import list_ports

        for port in list_ports.comports():
            print(port.device, "-", port.description)
        return
    if not (arguments.serial or arguments.udp):
        parser.error("請指定 --serial 或 --udp")
    if not arguments.monitor and arguments.out is None:
        parser.error("請指定 --out（或用 --monitor 只看數值）")

    recorder = Recorder(None if arguments.monitor else arguments.out)
    lines = serial_lines(arguments.serial, arguments.baud) if arguments.serial else udp_lines(arguments.udp)
    threading.Thread(target=reader, args=(lines, recorder), daemon=True).start()

    if arguments.monitor:
        countdown(recorder, "", "monitor", 10 ** 9, "即時數值，按 Ctrl+C 結束")
        return
    try:
        while True:
            sample_id = input("\n樣本編號（例如 B1-r03-d04；直接按 Enter 結束）：").strip()
            if not sample_id:
                break
            if sample_id in recorder.used_ids:
                print(f"{sample_id} 已經量過（在 CSV 裡），同一個編號會把兩次量測混在一起。"
                      "重量請用新編號，例如加 -r2。")
                continue
            recorder.used_ids.add(sample_id)  # even an interrupted run already wrote rows
            countdown(recorder, sample_id, "purge", arguments.purge, "採樣盒通乾淨空氣，量基準")
            input("→ 放入樣本、蓋上盒子後按 Enter")
            countdown(recorder, sample_id, "sample", arguments.sample, "量測樣本")
            input("→ 取出樣本、打開換氣後按 Enter")
            countdown(recorder, sample_id, "recover", arguments.recover, "換氣回復")
            recorder.used_ids.add(sample_id)
            print(f"完成 {sample_id}，累計收到 {recorder.count} 筆讀值")
    except KeyboardInterrupt:
        print("\n中斷；已寫入的資料都保留在 CSV。")
    sys.exit(0)


if __name__ == "__main__":
    main()