"""Tests for the in-memory HII gap-fill in the ML data loader."""

import datetime

import numpy as np
import pandas as pd
import pytest

from src.hii_collector import HiiStore
from src.ml.data import fill_from_hii

START = datetime.datetime(2024, 9, 1, 0, 0)


def _hours(n, offset=0):
    return [START + datetime.timedelta(hours=offset + i) for i in range(n)]


def _base_frame(code="P.20", n=200, level=2.0):
    return pd.DataFrame(
        {
            "timestamp": _hours(n),
            "station_code": code,
            "water_level": level,
            "discharge": 100.0,
        }
    )


@pytest.fixture
def hii_db(tmp_path):
    """SQLite DB with hii_* tables; returns (db_url, store)."""
    db_url = f"sqlite:///{tmp_path}/hii_fill.db"
    store = HiiStore(db_url, "sqlite")
    assert store.connect()
    return db_url, store

def _seed_mirror(store, station_id, rid_code, hours, wl_msl, discharge=250.0):
    from sqlalchemy import text

    with store.engine.begin() as conn:
        conn.execute(
            text("INSERT INTO hii_wl_stations (id, rid_code) VALUES (:i, :c)"),
            {"i": station_id, "c": rid_code},
        )
        conn.execute(
            text(
                "INSERT INTO hii_waterlevel (station_id, timestamp, wl_msl, discharge) "
                "VALUES (:i, :t, :w, :d)"
            ),
            [
                {"i": station_id, "t": t, "w": w, "d": discharge}
                for t, w in zip(hours, wl_msl)
            ],
        )


class TestFillFromHii:
    def test_exact_mirror_fills_missing_hours(self, hii_db):
        db_url, store = hii_db
        base = _base_frame("P.20", n=200, level=2.0)
        # Mirror covers the base window plus 48 extra hours, at MSL offset +300
        _seed_mirror(store, 1, "P.20", _hours(248), [302.0] * 248)

        filled = fill_from_hii(base, db_url, min_overlap_hours=168)
        p20 = filled[filled["station_code"] == "P.20"]
        assert len(p20) == 248
        new_rows = p20[p20["timestamp"] >= START + datetime.timedelta(hours=200)]
        assert len(new_rows) == 48
        # MSL converted back to gauge datum via the overlap-derived offset
        assert new_rows["water_level"].round(6).eq(2.0).all()
        # Exact mirrors copy discharge
        assert new_rows["discharge"].eq(250.0).all()

    def test_p81_bias_corrected_without_discharge(self, hii_db):
        db_url, store = hii_db
        base = _base_frame("P.81", n=200, level=1.5)
        # Biased mirror: offset 310.2 (not a published offset — derived only)
        _seed_mirror(store, 2, "P.81", _hours(230), [311.7] * 230)

        filled = fill_from_hii(base, db_url, min_overlap_hours=168)
        p81 = filled[filled["station_code"] == "P.81"]
        new_rows = p81[p81["timestamp"] >= START + datetime.timedelta(hours=200)]
        assert len(new_rows) == 30
        assert new_rows["water_level"].round(6).eq(1.5).all()
        assert new_rows["discharge"].isna().all()

    def test_insufficient_overlap_skips_station(self, hii_db):
        db_url, store = hii_db
        base = _base_frame("P.20", n=50)  # only 50 shared hours
        _seed_mirror(store, 1, "P.20", _hours(100), [302.0] * 100)

        filled = fill_from_hii(base, db_url, min_overlap_hours=168)
        assert len(filled) == len(base)

    def test_nan_level_rows_are_replaced(self, hii_db):
        db_url, store = hii_db
        base = _base_frame("P.20", n=200, level=2.0)
        base.loc[10, "water_level"] = np.nan
        _seed_mirror(store, 1, "P.20", _hours(200), [302.5] * 200)

        filled = fill_from_hii(base, db_url, min_overlap_hours=168)
        p20 = filled[filled["station_code"] == "P.20"]
        assert len(p20) == 200  # no duplicate hour
        replaced = p20[p20["timestamp"] == START + datetime.timedelta(hours=10)]
        assert replaced["water_level"].round(6).eq(2.0).all()

    def test_never_merged_station_untouched(self, hii_db):
        db_url, store = hii_db
        base = _base_frame("P.76", n=200, level=1.0)
        _seed_mirror(store, 3, "P.76", _hours(300), [301.0] * 300)

        filled = fill_from_hii(base, db_url, min_overlap_hours=1)
        assert len(filled) == len(base)  # P.76 not in HII_FILL_STATIONS

    def test_missing_tables_degrade_gracefully(self, tmp_path):
        base = _base_frame("P.20")
        filled = fill_from_hii(base, f"sqlite:///{tmp_path}/empty.db")
        assert filled.equals(base)

    def test_no_station_overlap_returns_input(self, hii_db):
        db_url, _ = hii_db
        base = _base_frame("P.5")  # not fillable
        assert fill_from_hii(base, db_url).equals(base)