#!/usr/bin/env python3
"""Run fan step-response test and capture sr,... serial log lines to CSV.

The firmware holds heater duty fixed, steps fan PWM, and logs temperature
every second. Use the output to see how chamber temp responds to fan changes.

Example:
  ./step_response.py
  ./step_response.py --temp 45 --heater 35 -o logs/stepresp.csv
"""

from __future__ import annotations

import argparse
import re
import sys
import time
from datetime import datetime, timezone
from pathlib import Path

from capture_csv import decode_line, open_serial, resolve_port

SR_RE = re.compile(
    r"^sr,(?P<ms>\d+),(?P<phase>\w+),(?P<step>\d+/\d+),"
    r"(?P<heater>\d+),(?P<fan>\d+),"
    r"(?P<avg>[\d.]+),(?P<min>[\d.]+),(?P<max>[\d.]+),(?P<spread>[\d.]+)"
    r"(?:,(?P<temps>.*))?$"
)
DONE_RE = re.compile(r"^stepresp: done")
FAIL_RE = re.compile(r"^stepresp: abort")

HEADER = (
    "wall_time,ms,phase,step,heater_pct,fan_pwm,fan_pct,avg_c,min_c,max_c,spread_c,"
    "ch2_t,ch3_t,ch4_t,ch5_t"
)


def fan_pct(pwm: int) -> int:
    return (pwm * 100) // 255


def send_command(ser, command: str, timeout: float = 3.0) -> list[str]:
    ser.write((command.strip() + "\n").encode("utf-8"))
    ser.flush()
    lines: list[str] = []
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        raw = ser.readline()
        if not raw:
            continue
        line = decode_line(raw)
        if not line:
            continue
        lines.append(line)
        if line.startswith("OK ") or line.startswith("ERR "):
            break
    return lines


def parse_sr_line(line: str) -> dict | None:
    match = SR_RE.match(line)
    if not match:
        return None
    data = match.groupdict()
    temps = data.pop("temps") or ""
    channels = (temps.split(",") + ["", "", "", ""])[:4]
    data["ch2_t"], data["ch3_t"], data["ch4_t"], data["ch5_t"] = channels
    data["fan_pwm"] = data.pop("fan")
    data["heater_pct"] = data.pop("heater")
    return data


def main() -> int:
    parser = argparse.ArgumentParser(description="Capture fan step-response data")
    parser.add_argument("-p", "--port", help="Serial port (default: auto-detect)")
    parser.add_argument("-b", "--baud", type=int, default=115200)
    parser.add_argument("--temp", type=float, default=45.0, help="Target temperature (C)")
    parser.add_argument("--heater", type=float, default=35.0, help="Fixed heater duty (%%)")
    parser.add_argument(
        "-o",
        "--output",
        type=Path,
        help="Output CSV (default: logs/stepresp_YYYYMMDD_HHMMSS.csv)",
    )
    args = parser.parse_args()

    port = resolve_port(args.port)
    out = args.output
    if out is None:
        out = Path("logs") / f"stepresp_{datetime.now():%Y%m%d_%H%M%S}.csv"

    print(f"Port: {port}", file=sys.stderr)
    print(f"Output: {out}", file=sys.stderr)
    print(f"Command: stepresp {args.temp:g} {args.heater:g}", file=sys.stderr)
    print("Ctrl+C to stop\n", file=sys.stderr)

    out.parent.mkdir(parents=True, exist_ok=True)
    row_count = 0

    with open_serial(port, args.baud) as ser, out.open("w", encoding="utf-8") as fh:
        fh.write(HEADER + "\n")
        ser.reset_input_buffer()

        lines = send_command(ser, f"stepresp {args.temp:g} {args.heater:g}")
        for line in lines:
            print(line, flush=True)
        if any(line.startswith("ERR ") for line in lines):
            return 1

        try:
            while True:
                raw = ser.readline()
                if not raw:
                    continue
                line = decode_line(raw)
                if not line:
                    continue

                if line.startswith("sr,"):
                    data = parse_sr_line(line)
                    if data is None:
                        print(f"WARN: bad sr line: {line}", file=sys.stderr)
                        continue
                    pwm = int(data["fan_pwm"])
                    wall = datetime.now(timezone.utc).isoformat(timespec="seconds")
                    row = [
                        wall,
                        data["ms"],
                        data["phase"],
                        data["step"],
                        data["heater_pct"],
                        str(pwm),
                        str(fan_pct(pwm)),
                        data["avg"],
                        data["min"],
                        data["max"],
                        data["spread"],
                        data["ch2_t"],
                        data["ch3_t"],
                        data["ch4_t"],
                        data["ch5_t"],
                    ]
                    fh.write(",".join(row) + "\n")
                    fh.flush()
                    row_count += 1
                    if row_count % 30 == 0:
                        print(
                            f"  {data['phase']} step {data['step']} "
                            f"fan={pwm} avg={data['avg']}C spread={data['spread']}C",
                            flush=True,
                        )
                    continue

                if DONE_RE.match(line) or FAIL_RE.match(line):
                    print(line, flush=True)
                    break

                if line.startswith("stepresp:"):
                    print(line, flush=True)
        except KeyboardInterrupt:
            print("\nStopping…", file=sys.stderr)
            send_command(ser, "stepresp stop")

    print(f"Wrote {row_count} rows to {out}", file=sys.stderr)
    return 0


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