from __future__ import annotations

import argparse
import csv
import sys
from pathlib import Path

import matplotlib.pyplot as plt


DEFAULT_CHANNELS = (
    "turn",
    "pulse",
    "gp0_falling",
    "gp1_falling",
)


def read_events(path: Path) -> tuple[list[str], list[list[float]]]:
    with path.open(newline="") as file:
        reader = csv.DictReader(file)
        fieldnames = set(reader.fieldnames or ())
        if "time_us" not in fieldnames:
            raise ValueError("CSV is missing column: time_us")

        channels = [channel for channel in DEFAULT_CHANNELS if channel in fieldnames]
        if not channels:
            raise ValueError("CSV has no plottable event columns")

        events = [[] for _ in channels]

        for row in reader:
            time_s = int(row["time_us"]) / 1_000_000
            for index, channel in enumerate(channels):
                if int(row[channel]):
                    events[index].append(time_s)

    return channels, events


def iter_csv_paths(path: Path) -> list[Path]:
    if path.is_dir():
        return sorted(path.rglob("*.csv"))
    return [path]


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Plot ESP32 event CSV files.")
    parser.add_argument("csv_path", type=Path, help="CSV file or directory of CSV files.")
    parser.add_argument("-o", "--output", type=Path, help="Optional image output path.")
    return parser.parse_args()


def plot_events(csv_path: Path, output: Path | None) -> bool:
    try:
        channels, events = read_events(csv_path)
    except (OSError, ValueError) as exc:
        print(f"Could not read {csv_path}: {exc}", file=sys.stderr)
        return False

    fig, ax = plt.subplots(figsize=(12, 5))
    ax.eventplot(events, orientation="horizontal", lineoffsets=range(len(channels)))
    ax.set_yticks(range(len(channels)), channels)
    ax.set_xlabel("time (s)")
    ax.set_title(str(csv_path))
    ax.grid(axis="x", alpha=0.25)
    fig.tight_layout()

    if output is not None:
        output.parent.mkdir(parents=True, exist_ok=True)
        fig.savefig(output, dpi=160)
        print(f"Wrote {output}")
        plt.close(fig)
    else:
        plt.show()
        plt.close(fig)

    return True


def output_path_for(csv_path: Path, output: Path | None, multiple: bool) -> Path | None:
    if output is None:
        return None
    if not multiple:
        return output
    return output / f"{csv_path.stem}.png"


def main() -> int:
    args = parse_args()

    if not args.csv_path.exists():
        print(f"Path does not exist: {args.csv_path}", file=sys.stderr)
        return 1

    csv_paths = iter_csv_paths(args.csv_path)
    if not csv_paths:
        print(f"No CSV files found under {args.csv_path}", file=sys.stderr)
        return 1

    multiple = len(csv_paths) > 1
    if multiple and args.output is not None and args.output.suffix:
        print("When plotting a directory, --output must be a directory.", file=sys.stderr)
        return 1

    failures = 0
    for csv_path in csv_paths:
        output = output_path_for(csv_path, args.output, multiple)
        if multiple and output is None:
            print(f"Plotting {csv_path}")
        if not plot_events(csv_path, output):
            failures += 1

    return 1 if failures else 0


if __name__ == "__main__":
    sys.exit(main())