"""Streaming CSV recorder for exact-timestamp-aligned QTFM sensor data.

A row holds one shared millisecond timestamp and every sensor's field. By
default no gradiometer difference columns are written — with many sensors
they multiply the file size while adding nothing that cannot be recomputed
with ``tools/gradiometer_parser.py``. Passing a ``reference`` sensor to
:meth:`CSVRecorder.start` opts into a self-contained file that also stores
every ``reference − sensor`` column.
"""

from __future__ import annotations

import csv
from datetime import datetime, timezone
from pathlib import Path
from typing import TextIO

from sensors import difference_label
from serial_connection import QTFMMeasurement


class CSVRecorder:
    """Write one CSV row per timestamp shared by every recorded sensor."""

    MAX_PENDING = 4096

    def __init__(self) -> None:
        self.active = False
        self.path: Path | None = None
        self.last_error = ""
        self.rows = 0
        self._file: TextIO | None = None
        self._writer: object | None = None
        self._sensors: tuple[str, ...] = ("A", "B")
        self._reference: str | None = None
        self._pending: dict[int, dict[str, float]] = {}
        self._gradiometer_enabled = False

    @property
    def sensors(self) -> tuple[str, ...]:
        """The sensor set captured when the recording started."""
        return self._sensors

    @property
    def reference(self) -> str | None:
        """The difference reference, or ``None`` for a fields-only file."""
        return self._reference

    def difference_names(self) -> tuple[str, ...]:
        """The recorded difference channels (empty for a fields-only file)."""
        if self._reference is None:
            return ()
        return tuple(
            difference_label(self._reference, key)
            for key in self._sensors
            if key != self._reference
        )

    @staticmethod
    def output_path(selected_path: str | Path) -> Path:
        selected = Path(selected_path)
        base = selected.with_suffix("") if selected.suffix else selected
        return Path(f"{base}_gradiometer.csv")

    @staticmethod
    def _field_header(sensor: str, serial_number: str | None) -> str:
        serial = (serial_number or "unknown").strip() or "unknown"
        return f"sensor_{sensor.lower()}_nt [{serial}]"

    @staticmethod
    def _difference_header(reference: str, sensor: str) -> str:
        return f"{reference.lower()}_minus_{sensor.lower()}_nt"

    def header_row(
        self,
        serial_numbers: dict[str, str | None],
    ) -> tuple[str, ...]:
        """Return the CSV header for the configured sensors."""
        difference_columns: tuple[str, ...] = ()
        if self._reference is not None:
            difference_columns = tuple(
                self._difference_header(self._reference, key)
                for key in self._sensors
                if key != self._reference
            )
        return (
            "pc_time_utc",
            "timestamp_ms",
            *(
                self._field_header(key, serial_numbers.get(key))
                for key in self._sensors
            ),
            *difference_columns,
        )

    def start(
        self,
        selected_path: str | Path,
        serial_numbers: dict[str, str | None] | None = None,
        sensors: tuple[str, ...] = ("A", "B"),
        reference: str | None = None,
    ) -> Path:
        """Begin recording; a non-``None`` reference adds difference columns."""
        if self.active:
            raise RuntimeError("A CSV recording is already active.")

        sensors = tuple(sensors)
        if reference is not None and reference not in sensors:
            reference = sensors[0] if sensors else None

        path = self.output_path(selected_path)
        path.parent.mkdir(parents=True, exist_ok=True)
        self.last_error = ""
        self.rows = 0
        self._sensors = sensors
        self._reference = reference
        self._pending = {}

        try:
            self._file = path.open("w", newline="", encoding="utf-8")
            self._writer = csv.writer(self._file)
            self._writer.writerow(self.header_row(serial_numbers or {}))
        except Exception:
            self._close_file()
            raise

        self.path = path
        self.active = True
        return path

    def set_gradiometer_enabled(self, enabled: bool) -> None:
        if enabled != self._gradiometer_enabled:
            self._pending = {}
        self._gradiometer_enabled = enabled

    def add_measurement(
        self,
        sensor: str,
        measurement: QTFMMeasurement,
    ) -> bool:
        if not self.active or sensor not in self._sensors:
            return False

        try:
            if (
                self._gradiometer_enabled
                and measurement.valid
                and measurement.timestamp_ms is not None
            ):
                self._add_candidate(
                    sensor,
                    measurement.timestamp_ms,
                    measurement.value_nt,
                )
            return True
        except Exception as exc:
            self.last_error = str(exc)
            self.stop()
            return False

    def _add_candidate(
        self,
        sensor: str,
        timestamp_ms: int,
        value_nt: float,
    ) -> None:
        group = self._pending.setdefault(timestamp_ms, {})
        group[sensor] = value_nt

        if len(group) < len(self._sensors):
            while len(self._pending) > self.MAX_PENDING:
                self._pending.pop(next(iter(self._pending)))
            return

        self._pending.pop(timestamp_ms, None)
        difference_values: tuple[str, ...] = ()
        if self._reference is not None:
            reference_value = group[self._reference]
            difference_values = tuple(
                f"{reference_value - group[key]:.9f}"
                for key in self._sensors
                if key != self._reference
            )
        now_utc = datetime.now(timezone.utc).isoformat(timespec="milliseconds")
        self._writer.writerow(
            (
                now_utc,
                timestamp_ms,
                *(f"{group[key]:.9f}" for key in self._sensors),
                *difference_values,
            )
        )
        self.rows += 1
        if self.rows % 250 == 0:
            self._file.flush()

    def stop(self) -> Path | None:
        self._close_file()
        self.active = False
        self._pending = {}
        return self.path

    def _close_file(self) -> None:
        if self._file is not None:
            try:
                self._file.flush()
                self._file.close()
            except OSError:
                pass
        self._file = None
        self._writer = None
