"""Hive-partitioned Parquet writer layout."""

from __future__ import annotations

import time
from pathlib import Path

import duckdb

from virtual_drone.flight import Pose
from virtual_drone.sensors import imu_row
from virtual_drone.writer import PartitionWriter, make_writers


def test_partition_path_follows_hive_layout(tmp_path: Path) -> None:
    ts_ns = int(time.time() * 1e9)
    writer = PartitionWriter(tmp_path, "telemetry", "flt-01", "dr-01", "imu")
    writer.append(imu_row(
        Pose(ts_ns=ts_ns, x=1.0, y=2.0, z=3.0, roll=0.0, pitch=0.0, yaw=0.0, vx=0.0, vy=0.0, vz=0.0),
        __import__("random").Random(0),
        ts_ns,
    ))
    writer.seal()

    files = list(tmp_path.rglob("*.parquet"))
    assert len(files) >= 1
    path_str = str(files[0])
    assert "dataset=telemetry" in path_str
    assert "flight=flt-01" in path_str
    assert "drone=dr-01" in path_str
    assert "sensor=imu" in path_str
    assert "year=" in path_str and "hour=" in path_str


def test_seal_compacts_current_blocks(tmp_path: Path) -> None:
    writer = PartitionWriter(tmp_path, "state", "flt-02", "dr-02", None)
    base_ts = 1_700_000_000_000_000_000
    for i in range(3):
        writer.append({"ts_ns": base_ts + i * 5_000_000_000, "pos_x": float(i)})
    writer.seal()

    sealed = list(tmp_path.rglob("data.parquet"))
    assert len(sealed) == 1
    assert not list(tmp_path.rglob("current/min_*.parquet"))
    rows = duckdb.sql(f"SELECT count(*) FROM read_parquet('{sealed[0]}')").fetchone()[0]
    assert rows == 3


def test_make_writers_registers_expected_datasets(tmp_path: Path) -> None:
    writers = make_writers(tmp_path, "flt-03", "dr-03", ["imu", "battery"])
    assert set(writers) == {"imu", "battery", "detections", "state"}


def test_duckdb_reads_written_partition(tmp_path: Path) -> None:
    writers = make_writers(tmp_path, "flt-04", "dr-04", ["imu"])
    ts_ns = int(time.time() * 1e9)
    pose = Pose(ts_ns=ts_ns, x=0.0, y=0.0, z=0.0, roll=0.0, pitch=0.0, yaw=0.0, vx=0.0, vy=0.0, vz=0.0)
    writers["imu"].append(imu_row(pose, __import__("random").Random(1), ts_ns))
    for w in writers.values():
        w.seal()

    count = duckdb.sql(f"""
        SELECT count(*) FROM read_parquet(
            '{tmp_path}/dataset=telemetry/**/*.parquet',
            hive_partitioning=true, union_by_name=true)
    """).fetchone()[0]
    assert count >= 1