#!/usr/bin/env python3
"""Plot dryer CSV logs from capture_csv.py, the TUI, or firmware `log on`.

Example:
  python3 scripts/plot_logs.py logs/dryer_20260707_195354.csv
  python3 scripts/plot_logs.py logs/          # newest .csv in directory
  python3 scripts/plot_logs.py remote         # newest .csv on alex@10.81.16.44
  python3 scripts/plot_logs.py logs/foo.csv -o plot.png
"""

from __future__ import annotations

import argparse
import csv
import subprocess
import sys
import tempfile
from pathlib import Path

REMOTE_SSH = "alex@10.81.16.44"
REMOTE_LOG_DIRS = (
    "~/arduino-filament-dryer/logs",
    "~/voron-filament-dryer/logs",
    "~/logs",
)


def import_matplotlib():
    try:
        import matplotlib.pyplot as plt
    except ImportError:
        print("Install matplotlib: pip install matplotlib", file=sys.stderr)
        raise SystemExit(1) from None
    return plt


def fetch_remote_csv(remote_dirs: tuple[str, ...] = REMOTE_LOG_DIRS) -> Path:
    dir_list = " ".join(remote_dirs)
    find_cmd = (
        f"for d in {dir_list}; do "
        'if [ -d "$d" ]; then ls -t "$d"/*.csv 2>/dev/null; fi; done | head -1'
    )
    result = subprocess.run(
        ["ssh", REMOTE_SSH, find_cmd],
        capture_output=True,
        text=True,
        check=False,
    )
    remote_path = result.stdout.strip()
    if result.returncode != 0 or not remote_path:
        err = result.stderr.strip()
        print(f"No remote CSV found on {REMOTE_SSH}", file=sys.stderr)
        if err:
            print(err, file=sys.stderr)
        raise SystemExit(1)

    local_path = Path(tempfile.gettempdir()) / f"dryer_remote_{Path(remote_path).name}"
    scp = subprocess.run(
        ["scp", f"{REMOTE_SSH}:{remote_path}", str(local_path)],
        capture_output=True,
        text=True,
        check=False,
    )
    if scp.returncode != 0:
        print(f"scp failed for {REMOTE_SSH}:{remote_path}", file=sys.stderr)
        if scp.stderr.strip():
            print(scp.stderr.strip(), file=sys.stderr)
        raise SystemExit(1)

    print(f"Fetched {REMOTE_SSH}:{remote_path}", file=sys.stderr)
    return local_path


def resolve_csv(path: Path) -> Path:
    if path.is_dir():
        matches = sorted(path.glob("*.csv"), key=lambda p: p.stat().st_mtime, reverse=True)
        if not matches:
            print(f"No CSV files in {path}", file=sys.stderr)
            raise SystemExit(1)
        return matches[0]
    if not path.is_file():
        print(f"Not found: {path}", file=sys.stderr)
        raise SystemExit(1)
    return path


def load_rows(path: Path) -> tuple[list[str], list[dict[str, str]]]:
    with path.open(newline="", encoding="utf-8") as fh:
        reader = csv.DictReader(fh)
        if reader.fieldnames is None:
            print(f"Empty CSV: {path}", file=sys.stderr)
            raise SystemExit(1)
        rows = list(reader)
    return list(reader.fieldnames), rows


def column_float(rows: list[dict[str, str]], name: str) -> list[float | None]:
    out: list[float | None] = []
    for row in rows:
        raw = row.get(name, "").strip()
        if not raw:
            out.append(None)
            continue
        try:
            out.append(float(raw))
        except ValueError:
            out.append(None)
    return out


def time_axis(rows: list[dict[str, str]], fieldnames: list[str]) -> tuple[list[float], str]:
    if "ms" in fieldnames:
        ms = column_float(rows, "ms")
        if any(v is not None for v in ms):
            t0 = next(v for v in ms if v is not None)
            return [((v or t0) - t0) / 60000.0 for v in ms], "minutes since start"

    if "wall_time" in fieldnames:
        from datetime import datetime

        times: list[float] = []
        parsed: list[datetime] = []
        for row in rows:
            raw = row.get("wall_time", "").strip()
            if not raw:
                continue
            try:
                parsed.append(datetime.fromisoformat(raw))
            except ValueError:
                continue
        if parsed:
            t0 = parsed[0]
            for dt in parsed:
                times.append((dt - t0).total_seconds() / 60.0)
            return times, "minutes since start"

    return [float(i) for i in range(len(rows))], "sample"


def plot_csv(path: Path, output: Path | None, title: str | None = None) -> None:
    plt = import_matplotlib()

    fieldnames, rows = load_rows(path)
    if not rows:
        print(f"No data rows in {path}", file=sys.stderr)
        raise SystemExit(1)

    x, x_label = time_axis(rows, fieldnames)
    if len(x) != len(rows):
        x = [float(i) for i in range(len(rows))]
        x_label = "sample"

    fig, axes = plt.subplots(3, 1, figsize=(11, 8), sharex=True, constrained_layout=True)
    fig.suptitle(title or path.name)

    temp_ax = axes[0]
    for col, label, style in (
        ("avg_c", "avg", "-"),
        ("min_c", "min", "--"),
        ("max_c", "max", "--"),
        ("target_c", "target", ":"),
    ):
        if col not in fieldnames:
            continue
        y = column_float(rows, col)
        temp_ax.plot(x, y, style, label=label, linewidth=1.5 if col == "avg_c" else 1.0)
    for ch in (2, 3, 4, 5):
        col = f"ch{ch}_t"
        if col in fieldnames:
            y = column_float(rows, col)
            temp_ax.plot(x, y, "-", alpha=0.35, linewidth=0.8, label=f"ch{ch}")
    temp_ax.set_ylabel("°C")
    temp_ax.legend(loc="upper left", ncol=4, fontsize=8)
    temp_ax.grid(True, alpha=0.3)

    duty_ax = axes[1]
    if "heater_pct" in fieldnames:
        duty_ax.plot(x, column_float(rows, "heater_pct"), "C1-", label="heater %")
    if "heatlim_pct" in fieldnames:
        duty_ax.plot(x, column_float(rows, "heatlim_pct"), "C1--", alpha=0.6, label="heatlim %")
    if "fan_pct" in fieldnames:
        duty_ax.plot(x, column_float(rows, "fan_pct"), "C0-", label="fan %")
    duty_ax.set_ylabel("%")
    duty_ax.legend(loc="upper left", fontsize=8)
    duty_ax.grid(True, alpha=0.3)

    spread_ax = axes[2]
    if "spread_c" in fieldnames:
        spread_ax.plot(x, column_float(rows, "spread_c"), "C2-", label="spread")
    spread_ax.set_ylabel("°C")
    spread_ax.set_xlabel(x_label)
    spread_ax.legend(loc="upper left", fontsize=8)
    spread_ax.grid(True, alpha=0.3)

    if output is not None:
        fig.savefig(output, dpi=150)
        print(f"Wrote {output}")
    else:
        plt.show()


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Plot dryer CSV temperature logs")
    parser.add_argument(
        "csv",
        help='CSV file, directory (newest .csv), or "remote" for newest on ' + REMOTE_SSH,
    )
    parser.add_argument("-o", "--output", type=Path, help="Save PNG instead of opening a window")
    parser.add_argument(
        "--remote-dir",
        action="append",
        metavar="DIR",
        help=f"Remote log directory on {REMOTE_SSH} (repeatable; used with remote)",
    )
    args = parser.parse_args(argv)

    plot_title: str | None = None
    if args.csv == "remote":
        remote_dirs = tuple(args.remote_dir) if args.remote_dir else REMOTE_LOG_DIRS
        path = fetch_remote_csv(remote_dirs)
        plot_title = f"{REMOTE_SSH}:{path.name}"
    else:
        path = resolve_csv(Path(args.csv))

    if args.output is None:
        print(f"Plotting {plot_title or path}", file=sys.stderr)
    plot_csv(path, args.output, title=plot_title)
    return 0


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