"""Hive-partitioned Parquet writer: the hot `current/` path plus sealing.

Layout (identical to the on-board and warehouse design):

    dataset=<name>/flight=<id>/drone=<id>[/sensor=<type>]/
        year=YYYY/month=MM/day=DD/hour=HH/
            current/min_MM.parquet   <- open window, small blocks
            data.parquet             <- sealed on window close / shutdown
"""

from __future__ import annotations

import time
from pathlib import Path
from typing import Any

import pyarrow as pa
import pyarrow.parquet as pq

Row = dict[str, Any]


def _hour_dir(root: Path, dataset: str, flight: str, drone: str, sensor: str | None, ts_ns: int) -> Path:
    t = time.gmtime(ts_ns / 1e9)
    parts = [f"dataset={dataset}", f"flight={flight}", f"drone={drone}"]
    if sensor is not None:
        parts.append(f"sensor={sensor}")
    parts += [
        f"year={t.tm_year:04d}", f"month={t.tm_mon:02d}",
        f"day={t.tm_mday:02d}", f"hour={t.tm_hour:02d}",
    ]
    return root.joinpath(*parts)


class PartitionWriter:
    """Buffers rows per minute and flushes small blocks into current/."""

    def __init__(self, root: Path, dataset: str, flight: str, drone: str, sensor: str | None) -> None:
        self._root = root
        self._dataset = dataset
        self._flight = flight
        self._drone = drone
        self._sensor = sensor
        self._buffer: list[Row] = []
        self._minute: int | None = None
        self._touched: set[Path] = set()

    def append(self, row: Row) -> None:
        minute = int(row["ts_ns"] // 60_000_000_000)
        if self._minute is not None and minute != self._minute:
            self._flush()
        self._minute = minute
        self._buffer.append(row)

    def _flush(self) -> None:
        if not self._buffer or self._minute is None:
            return
        first_ts = self._buffer[0]["ts_ns"]
        hour_dir = _hour_dir(self._root, self._dataset, self._flight, self._drone, self._sensor, first_ts)
        current = hour_dir / "current"
        current.mkdir(parents=True, exist_ok=True)
        table = pa.Table.from_pylist(self._buffer)
        minute_of_hour = self._minute % 60
        pq.write_table(table, current / f"min_{minute_of_hour:02d}.parquet", compression="zstd")
        self._touched.add(hour_dir)
        self._buffer = []

    def seal(self) -> None:
        """Compact every touched hour's current/ blocks into one data.parquet."""
        self._flush()
        for hour_dir in sorted(self._touched):
            current = hour_dir / "current"
            blocks = sorted(current.glob("min_*.parquet"))
            if not blocks:
                continue
            table = pa.concat_tables([pq.read_table(b) for b in blocks])
            pq.write_table(table, hour_dir / "data.parquet", compression="zstd", compression_level=9)
            for b in blocks:
                b.unlink()
            current.rmdir()


def make_writers(root: Path, flight: str, drone: str, sensor_names: list[str]) -> dict[str, PartitionWriter]:
    """One telemetry writer per sensor + detections + state."""
    writers = {
        name: PartitionWriter(root, "telemetry", flight, drone, name)
        for name in sensor_names
    }
    writers["detections"] = PartitionWriter(root, "detections", flight, drone, None)
    writers["state"] = PartitionWriter(root, "state", flight, drone, None)
    return writers