from __future__ import annotations

from datetime import datetime, timezone

import polars as pl

from uforia.storage import DatasetName, DatasetStore, ReadMode, RunStatus
from uforia.storage.datasets import DATASET_SPECS


def test_store_keeps_all_runs_and_latest_read_in_memory(tmp_path) -> None:
    store = DatasetStore(tmp_path)
    ts = datetime(2026, 3, 1, 12, tzinfo=timezone.utc)
    base = {
        "ts": ts,
        "date": ts.date(),
        "venue": "deribit",
        "underlying": "BTC",
        "source": "dvol",
        "dvol_open": 50.0,
        "dvol_high": 52.0,
        "dvol_low": 49.0,
        "quality_flag": "ok",
    }
    first = pl.DataFrame([{**base, "dvol_close": 51.0}])
    second = pl.DataFrame([{**base, "dvol_close": 53.0}])
    store.write(
        DatasetName.DVOL_SNAPSHOT,
        first,
        run_id="run-one",
        run_ts=datetime(2026, 3, 1, 12, 1, tzinfo=timezone.utc),
        status=RunStatus.SUCCESS,
    )
    store.write(
        DatasetName.DVOL_SNAPSHOT,
        second,
        run_id="run-two",
        run_ts=datetime(2026, 3, 1, 12, 2, tzinfo=timezone.utc),
        status=RunStatus.SUCCESS,
    )

    latest = store.read(DatasetName.DVOL_SNAPSHOT)
    all_runs = store.read(DatasetName.DVOL_SNAPSHOT, read_mode=ReadMode.ALL_RUNS)

    assert latest.height == 1
    assert latest.item(0, "dvol_close") == 53.0
    assert all_runs.height == 2
    assert not list(tmp_path.rglob("*.parquet"))


def test_store_returns_empty_frame_with_schema_for_missing_dataset(tmp_path) -> None:
    store = DatasetStore(tmp_path)

    frame = store.read(DatasetName.VVIX_SERIES)

    assert frame.is_empty()
    assert frame.columns == list(DATASET_SPECS[DatasetName.VVIX_SERIES].schema.keys())


def test_store_compact_deduplicates_latest_rows_in_memory(tmp_path) -> None:
    store = DatasetStore(tmp_path)
    ts = datetime(2026, 3, 1, 12, tzinfo=timezone.utc)
    base = {
        "ts": ts,
        "date": ts.date(),
        "venue": "deribit",
        "underlying": "BTC",
        "source": "dvol",
        "dvol_open": 50.0,
        "dvol_high": 52.0,
        "dvol_low": 49.0,
        "quality_flag": "ok",
    }
    store.write(
        DatasetName.DVOL_SNAPSHOT,
        pl.DataFrame([{**base, "dvol_close": 51.0}]),
        run_id="run-one",
        run_ts=datetime(2026, 3, 1, 12, 1, tzinfo=timezone.utc),
    )
    store.write(
        DatasetName.DVOL_SNAPSHOT,
        pl.DataFrame([{**base, "dvol_close": 53.0}]),
        run_id="run-two",
        run_ts=datetime(2026, 3, 1, 12, 2, tzinfo=timezone.utc),
    )

    result = store.compact(DatasetName.DVOL_SNAPSHOT)
    frame = store.read(DatasetName.DVOL_SNAPSHOT, read_mode=ReadMode.ALL_RUNS)

    assert result == {"partitions": 1, "rows": 1, "deleted_files": 0}
    assert frame.height == 1
    assert frame.item(0, "dvol_close") == 53.0
