from __future__ import annotations

from datetime import date, datetime, timezone

import polars as pl
import pytest

from uforia.storage import DatasetName, DatasetStore
from uforia.storage.datasets import DATASET_SPECS
from uforia.storage.db import DbSettings, TimescaleDatasetStore


class FakeDb:
    def __init__(
        self,
        frame: pl.DataFrame | None = None,
        *,
        fail_read: bool = False,
        fail_write: bool = False,
    ) -> None:
        self.frame = frame
        self.fail_read = fail_read
        self.fail_write = fail_write
        self.writes: list[tuple[DatasetName, int]] = []

    def write(self, dataset: DatasetName, frame: pl.DataFrame) -> int:
        if self.fail_write:
            raise RuntimeError("db write down")
        self.writes.append((dataset, frame.height))
        return frame.height

    def read(self, dataset: DatasetName, **_: object) -> pl.DataFrame:
        if self.fail_read:
            raise RuntimeError("db down")
        return self.frame if self.frame is not None else pl.DataFrame()

    def read_latest(self, dataset: DatasetName, **_: object) -> pl.DataFrame:
        return self.read(dataset)

    def read_latest_slice(self, dataset: DatasetName, **_: object) -> pl.DataFrame:
        return self.read(dataset)

    def latest_logical_ts(self, dataset: DatasetName, filters: dict[str, object] | None = None):
        del dataset, filters
        return None

    def health(self) -> dict[str, object]:
        return {"enabled": True, "status": "ok"}

    def verify_with_frame(
        self,
        dataset: DatasetName,
        frame: pl.DataFrame | None,
        *,
        start: datetime | None = None,
        end: datetime | None = None,
        filters: dict[str, object] | None = None,
    ) -> dict[str, object]:
        expected = frame if frame is not None else pl.DataFrame()
        return {
            "dataset": dataset.value,
            "source_all_runs_rows": expected.height,
            "db_all_runs_rows": 0 if self.frame is None else self.frame.height,
            "start": start,
            "end": end,
            "filters": filters or {},
        }


def _sample_dvol(close: float = 51.0) -> pl.DataFrame:
    ts = datetime(2026, 3, 1, 12, tzinfo=timezone.utc)
    return pl.DataFrame(
        [
            {
                "ts": ts,
                "date": ts.date(),
                "venue": "deribit",
                "underlying": "BTC",
                "source": "dvol",
                "dvol_open": 50.0,
                "dvol_high": 52.0,
                "dvol_low": 49.0,
                "dvol_close": close,
                "quality_flag": "ok",
            }
        ]
    )


def test_store_db_write_calls_backend(tmp_path) -> None:
    fake_db = FakeDb()
    store = DatasetStore(
        tmp_path,
        db_settings=DbSettings(
            enabled=True,
            dsn="postgresql://unused",
            schema="uforia",
            read_mode="db",
            write_mode="db",
            strict_datasets=frozenset(),
        ),
        db_backend=fake_db,  # type: ignore[arg-type]
    )

    written = store.write(DatasetName.DVOL_SNAPSHOT, _sample_dvol())

    assert len(written) == 1
    assert fake_db.writes == [(DatasetName.DVOL_SNAPSHOT, 1)]


def test_store_db_write_mode_raises_on_db_failure(tmp_path) -> None:
    store = DatasetStore(
        tmp_path,
        db_settings=DbSettings(
            enabled=True,
            dsn="postgresql://unused",
            schema="uforia",
            read_mode="db",
            write_mode="db",
            strict_datasets=frozenset(),
        ),
        db_backend=FakeDb(fail_write=True),  # type: ignore[arg-type]
    )

    with pytest.raises(RuntimeError, match="db write down"):
        store.write(DatasetName.DVOL_SNAPSHOT, _sample_dvol())


def test_store_db_read_mode_prefers_db_backend(tmp_path) -> None:
    expected = _sample_dvol(close=77.0).with_columns(
        pl.lit("run-one").alias("run_id"),
        pl.lit(datetime(2026, 3, 1, 12, 1, tzinfo=timezone.utc)).alias("run_ts"),
        pl.col("ts").alias("logical_ts"),
        pl.lit("immutable").alias("write_mode"),
    )
    store = DatasetStore(
        tmp_path,
        db_settings=DbSettings(
            enabled=True,
            dsn="postgresql://unused",
            schema="uforia",
            read_mode="db",
            write_mode="db",
            strict_datasets=frozenset(),
        ),
        db_backend=FakeDb(expected),  # type: ignore[arg-type]
    )

    frame = store.read(DatasetName.DVOL_SNAPSHOT)

    assert frame.height == 1
    assert frame.item(0, "dvol_close") == 77.0


def test_store_db_read_latest_uses_db_latest_backend(tmp_path) -> None:
    expected = _sample_dvol(close=77.0).with_columns(
        pl.lit("run-one").alias("run_id"),
        pl.lit(datetime(2026, 3, 1, 12, 1, tzinfo=timezone.utc)).alias("run_ts"),
        pl.col("ts").alias("logical_ts"),
        pl.lit("immutable").alias("write_mode"),
    )
    store = DatasetStore(
        tmp_path,
        db_settings=DbSettings(
            enabled=True,
            dsn="postgresql://unused",
            schema="uforia",
            read_mode="db",
            write_mode="db",
            strict_datasets=frozenset(),
        ),
        db_backend=FakeDb(expected),  # type: ignore[arg-type]
    )

    frame = store.read_latest(DatasetName.DVOL_SNAPSHOT)

    assert frame.height == 1
    assert frame.item(0, "dvol_close") == 77.0


def test_store_db_read_latest_slice_uses_db_backend(tmp_path) -> None:
    expected = _sample_dvol(close=77.0).with_columns(
        pl.lit("run-one").alias("run_id"),
        pl.lit(datetime(2026, 3, 1, 12, 1, tzinfo=timezone.utc)).alias("run_ts"),
        pl.col("ts").alias("logical_ts"),
        pl.lit("immutable").alias("write_mode"),
    )
    store = DatasetStore(
        tmp_path,
        db_settings=DbSettings(
            enabled=True,
            dsn="postgresql://unused",
            schema="uforia",
            read_mode="db",
            write_mode="db",
            strict_datasets=frozenset(),
        ),
        db_backend=FakeDb(expected),  # type: ignore[arg-type]
    )

    frame = store.read_latest_slice(DatasetName.DVOL_SNAPSHOT, columns=["ts", "dvol_close"])

    assert frame.height == 1
    assert frame.columns == ["ts", "dvol_close"]
    assert frame.item(0, "dvol_close") == 77.0


def test_store_db_verify_reports_db_only_summary(tmp_path) -> None:
    store = DatasetStore(
        tmp_path,
        db_settings=DbSettings(
            enabled=True,
            dsn="postgresql://unused",
            schema="uforia",
            read_mode="db",
            write_mode="db",
            strict_datasets=frozenset(),
        ),
        db_backend=FakeDb(_sample_dvol(close=70.0)),  # type: ignore[arg-type]
    )

    result = store.db_verify(DatasetName.DVOL_SNAPSHOT)

    assert result["dataset"] == DatasetName.DVOL_SNAPSHOT.value
    assert result["db_all_runs_rows"] == 1
    assert result["source_all_runs_rows"] == 0


def test_db_backfill_is_removed(tmp_path) -> None:
    store = DatasetStore(
        tmp_path,
        db_settings=DbSettings(
            enabled=True,
            dsn="postgresql://unused",
            schema="uforia",
            read_mode="db",
            write_mode="db",
            strict_datasets=frozenset(),
        ),
        db_backend=FakeDb(),  # type: ignore[arg-type]
    )

    with pytest.raises(RuntimeError, match="Parquet backfill has been removed"):
        store.db_backfill(DatasetName.DVOL_SNAPSHOT)


def test_timescale_read_backfills_missing_db_columns(monkeypatch) -> None:
    settings = DbSettings(
        enabled=True,
        dsn="postgresql://unused",
        schema="uforia",
        read_mode="db",
        write_mode="db",
        strict_datasets=frozenset(),
    )
    store = TimescaleDatasetStore(settings)

    class FakeCursor:
        def execute(self, sql, params=None) -> None:
            self.sql = sql
            self.params = params

        def fetchall(self):
            return [
                {
                    "ts": datetime(2026, 3, 1, 12, tzinfo=timezone.utc),
                    "date": date(2026, 3, 1),
                    "asset": "BTC",
                    "source": "derived",
                    "vol_carry_score": 1.25,
                    "run_id": "run-one",
                    "run_ts": datetime(2026, 3, 1, 12, 1, tzinfo=timezone.utc),
                    "logical_ts": datetime(2026, 3, 1, 12, tzinfo=timezone.utc),
                    "write_mode": "immutable",
                }
            ]

        def __enter__(self):
            return self

        def __exit__(self, exc_type, exc, tb) -> None:
            return None

    class FakeConn:
        def cursor(self):
            return FakeCursor()

        def __enter__(self):
            return self

        def __exit__(self, exc_type, exc, tb) -> None:
            return None

    monkeypatch.setattr(store, "_conn", lambda: FakeConn())
    monkeypatch.setattr(
        store,
        "_table_columns",
        lambda dataset: frozenset(
            {
                "ts",
                "date",
                "asset",
                "source",
                "vol_carry_score",
                "run_id",
                "run_ts",
                "logical_ts",
                "write_mode",
            }
        ),
    )

    frame = store.read(DatasetName.OPTIONS_SIGNAL_SERIES)

    assert frame.height == 1
    assert "surface_factor_score" in frame.columns


def test_timescale_read_latest_forces_latest_mode(monkeypatch) -> None:
    settings = DbSettings(
        enabled=True,
        dsn="postgresql://unused",
        schema="uforia",
        read_mode="db",
        write_mode="db",
        strict_datasets=frozenset(),
    )
    store = TimescaleDatasetStore(settings)

    captured: dict[str, object] = {}

    def fake_read(
        dataset: DatasetName,
        *,
        start=None,
        end=None,
        filters=None,
        read_mode=None,
        columns=None,
    ) -> pl.DataFrame:
        captured["dataset"] = dataset
        captured["read_mode"] = read_mode
        del start, end, filters, columns
        return pl.DataFrame()

    monkeypatch.setattr(store, "read", fake_read)

    store.read_latest(DatasetName.DVOL_SNAPSHOT)

    assert captured["dataset"] == DatasetName.DVOL_SNAPSHOT
    assert captured["read_mode"] == "latest" or str(captured["read_mode"]).endswith("LATEST")


def test_timescale_read_latest_slice_returns_requested_columns(monkeypatch) -> None:
    settings = DbSettings(
        enabled=True,
        dsn="postgresql://unused",
        schema="uforia",
        read_mode="db",
        write_mode="db",
        strict_datasets=frozenset(),
    )
    store = TimescaleDatasetStore(settings)
    captured: dict[str, object] = {}

    class FakeCursor:
        def execute(self, sql, params=None) -> None:
            captured["sql"] = sql
            captured["params"] = params

        def fetchall(self):
            return [
                {
                    "ts": datetime(2026, 3, 1, 12, tzinfo=timezone.utc),
                    "dvol_close": 77.0,
                    "run_ts": datetime(2026, 3, 1, 12, 1, tzinfo=timezone.utc),
                    "logical_ts": datetime(2026, 3, 1, 12, tzinfo=timezone.utc),
                }
            ]

        def __enter__(self):
            return self

        def __exit__(self, exc_type, exc, tb) -> None:
            return None

    class FakeConn:
        def cursor(self):
            return FakeCursor()

        def __enter__(self):
            return self

        def __exit__(self, exc_type, exc, tb) -> None:
            return None

    monkeypatch.setattr(store, "_conn", lambda: FakeConn())
    monkeypatch.setattr(
        store,
        "_table_columns",
        lambda dataset: frozenset(
            {"ts", "date", "venue", "underlying", "source", "dvol_close", "run_ts", "logical_ts"}
        ),
    )

    frame = store.read_latest_slice(
        DatasetName.DVOL_SNAPSHOT,
        filters={"source": "dvol", "underlying": "BTC"},
        columns=["ts", "dvol_close"],
    )

    assert frame.columns == ["ts", "dvol_close"]
    assert frame.item(0, "dvol_close") == 77.0
    assert "WITH filtered_rows AS" in str(captured["sql"])
    assert "MAX(COALESCE(\"logical_ts\", \"ts\"))" in str(captured["sql"])


def test_timescale_init_adds_missing_columns_for_existing_tables() -> None:
    settings = DbSettings(
        enabled=True,
        dsn="postgresql://unused",
        schema="uforia",
        read_mode="db",
        write_mode="db",
        strict_datasets=frozenset(),
    )
    store = TimescaleDatasetStore(settings)
    executed: list[tuple[str, object]] = []

    class FakeCursor:
        def execute(self, sql, params=None) -> None:
            executed.append((sql, params))
            self.sql = sql

        def fetchall(self):
            if "information_schema.columns" not in self.sql:
                return []
            return [
                {"column_name": "ts"},
                {"column_name": "date"},
                {"column_name": "underlying"},
                {"column_name": "signal_state"},
                {"column_name": "run_id"},
                {"column_name": "run_ts"},
                {"column_name": "logical_ts"},
                {"column_name": "write_mode"},
            ]

    store._create_table(  # type: ignore[attr-defined]
        FakeCursor(),
        DATASET_SPECS[DatasetName.OPTIONS_SIGNAL_SERIES],
        timescaledb_available=False,
    )

    alter_statements = [sql for sql, _ in executed if "ALTER TABLE" in sql]
    assert any('"surface_factor_score"' in sql for sql in alter_statements)
