from __future__ import annotations

from datetime import datetime, timezone

import polars as pl

from uforia.cli import _merge_with_existing_vvix
from uforia.storage import DatasetName, DatasetStore


def test_merge_with_existing_vvix_preserves_existing_qvvix(tmp_path) -> None:
    store = DatasetStore(tmp_path)
    ts = datetime(2026, 3, 9, 14, tzinfo=timezone.utc)
    store.write(
        DatasetName.VVIX_SERIES,
        pl.DataFrame(
            [
                {
                    "ts": ts,
                    "date": ts.date(),
                    "underlying": "BTC",
                    "source": "dvol",
                    "vol_index_value": 57.0,
                    "rvvix_7d": None,
                    "rvvix_14d": None,
                    "rvvix_30d": None,
                    "qvvix_30d": 88.0,
                    "model_version": "qvvix-var-v1",
                    "quality_flag": "ok",
                }
            ]
        ),
    )
    incoming = pl.DataFrame(
        [
            {
                "ts": ts,
                "date": ts.date(),
                "underlying": "BTC",
                "source": "dvol",
                "vol_index_value": 57.0,
                "rvvix_7d": 101.0,
                "rvvix_14d": 99.0,
                "rvvix_30d": 97.0,
                "qvvix_30d": None,
                "model_version": "rvvix-v1",
                "quality_flag": "ok",
            }
        ]
    )

    merged = _merge_with_existing_vvix(store, incoming, source="dvol", underlying="BTC")

    assert merged.item(0, "rvvix_14d") == 99.0
    assert merged.item(0, "qvvix_30d") == 88.0


def test_merge_with_existing_vvix_preserves_existing_rvvix(tmp_path) -> None:
    store = DatasetStore(tmp_path)
    ts = datetime(2026, 3, 9, 14, tzinfo=timezone.utc)
    store.write(
        DatasetName.VVIX_SERIES,
        pl.DataFrame(
            [
                {
                    "ts": ts,
                    "date": ts.date(),
                    "underlying": "BTC",
                    "source": "dvol",
                    "vol_index_value": 57.0,
                    "rvvix_7d": 101.0,
                    "rvvix_14d": 99.0,
                    "rvvix_30d": 97.0,
                    "qvvix_30d": None,
                    "model_version": "rvvix-v1",
                    "quality_flag": "ok",
                }
            ]
        ),
    )
    incoming = pl.DataFrame(
        [
            {
                "ts": ts,
                "date": ts.date(),
                "underlying": "BTC",
                "source": "dvol",
                "vol_index_value": 57.0,
                "rvvix_7d": None,
                "rvvix_14d": None,
                "rvvix_30d": None,
                "qvvix_30d": 88.0,
                "model_version": "qvvix-var-v1",
                "quality_flag": "ok",
            }
        ]
    )

    merged = _merge_with_existing_vvix(store, incoming, source="dvol", underlying="BTC")

    assert merged.item(0, "rvvix_14d") == 99.0
    assert merged.item(0, "qvvix_30d") == 88.0
