"""Batched multi-row upsert in TimescaleFeatureStore: the write paths must issue ONE multi-row
`INSERT ... ON CONFLICT` per batch (≈1000x fewer round-trips than per-row executemany) while keeping the
EXACT idempotent semantics (re-upsert overwrites, distinct features at same ts coexist, no dupes). These
run against in-memory SQLite (which, like Postgres, supports multi-row VALUES + ON CONFLICT) so they
exercise the real batching path without a live Postgres."""
from __future__ import annotations

from fxhnt.adapters.warehouse import TimescaleFeatureStore


def _count(store) -> int:
    from sqlalchemy import text

    with store._engine.connect() as c:
        return int(c.execute(text("SELECT count(*) FROM features")).scalar_one())


def test_large_upsert_correct_and_idempotent() -> None:
    """50k synthetic rows via one upsert_feature_rows call: all present, exact count (no dupes), and a
    re-upsert overwrites in place (idempotent), still no dupes."""
    store = TimescaleFeatureStore("sqlite://")
    n = 50_000
    rows = [("SYM", ts, "close", float(ts)) for ts in range(n)]
    store.upsert_feature_rows(rows)
    assert _count(store) == n

    # spot-check a few values round-trip
    read = store.read_features("SYM", start_ts=0, end_ts=2)
    assert dict(read) == {0: {"close": 0.0}, 1: {"close": 1.0}, 2: {"close": 2.0}}

    # re-upsert the SAME keys with new values -> overwrite, never duplicate
    rows2 = [("SYM", ts, "close", float(ts) + 1000.0) for ts in range(n)]
    store.upsert_feature_rows(rows2)
    assert _count(store) == n  # still exactly n, no dupes
    read2 = store.read_features("SYM", start_ts=0, end_ts=1)
    assert dict(read2) == {0: {"close": 1000.0}, 1: {"close": 1001.0}}
    store.close()


def test_distinct_features_same_ts_coexist() -> None:
    """The PK is (feature, symbol, ts): distinct features at the same ts must NOT collide through batching."""
    store = TimescaleFeatureStore("sqlite://")
    ts = 19000 * 86400
    rows = [("SYM", ts, "close", 1.0), ("SYM", ts, "volume", 2.0), ("SYM", ts, "funding", 3.0)]
    store.upsert_feature_rows(rows)
    assert _count(store) == 3
    got = store.as_of("SYM", ts)
    assert got == {"close": 1.0, "volume": 2.0, "funding": 3.0}
    store.close()


def test_batch_split_into_multiple_inserts(monkeypatch) -> None:
    """A call whose row count exceeds the param-safe batch size must be split into MULTIPLE INSERTs (not a
    single oversized statement), while results stay correct. Each batch is further split into power-of-two
    sub-INSERTs so the compiled-statement cache stays bounded (the OOM fix), so we assert >1 INSERT and
    correct results rather than an exact count."""
    store = TimescaleFeatureStore("sqlite://")

    calls = {"n": 0}

    # Wrap Connection.execute to count multi-row INSERT statements issued by the bulk path.
    from sqlalchemy.engine import Connection

    orig_execute = Connection.execute

    def counting_execute(self, statement, *args, **kwargs):
        try:
            sql = str(statement)
        except Exception:
            sql = ""
        if sql.startswith("INSERT INTO features"):
            calls["n"] += 1
        return orig_execute(self, statement, *args, **kwargs)

    monkeypatch.setattr(Connection, "execute", counting_execute)

    # Use a small batch_size so 3x batch_size rows guarantee multiple INSERTs.
    bs = 500
    n = bs * 3 + 7  # 1507 rows -> multiple power-of-two sub-INSERTs across 4 chunks
    rows = [("SYM", ts, "close", float(ts)) for ts in range(n)]
    with store._engine.begin() as conn:
        store._bulk_upsert(conn, rows, batch_size=bs)

    assert calls["n"] > 1, f"expected the oversized load to split into multiple INSERTs, got {calls['n']}"
    assert _count(store) == n  # every row landed exactly once (no padding, no dupes)
    store.close()


def test_default_batch_size_param_safe() -> None:
    """Default batch_size must keep bound params under the Postgres 65535 cap (4 cols x batch_size)."""
    import inspect

    sig = inspect.signature(TimescaleFeatureStore._bulk_upsert)
    default_bs = sig.parameters["batch_size"].default
    assert isinstance(default_bs, int)
    assert 4 * default_bs <= 65_535, "4 cols x batch_size must stay under PG 65535 bind-param cap"


def test_compiled_statement_cache_bounded_over_many_varied_size_upserts() -> None:
    """THE OOM REGRESSION GUARD. Upserting many symbols whose row counts (and therefore remainder-batch
    sizes) all differ must NOT grow the engine's compiled-statement cache without bound. A multi-row VALUES
    INSERT compiles a distinct prepared statement per distinct tuple-count; the old variable-size flush cached
    one per distinct size → O(num distinct row counts) compiled statements → ~2Gi OOM at ~683 symbols.

    The power-of-two flush bounds the tuple-counts to powers of two, so the cache must stay O(log batch_size)
    (~a dozen entries) regardless of how many differently-sized symbols are written."""
    store = TimescaleFeatureStore("sqlite://")
    # 400 symbols, every one a DIFFERENT row count from 1..400 -> 400 distinct remainder sizes.
    for sym in range(400):
        n = sym + 1
        store.upsert_feature_rows([(f"S{sym}", ts, "close", float(ts)) for ts in range(n)])

    cache = store._engine._compiled_cache
    distinct = len(cache) if cache is not None else 0
    # Powers of two up to 400 = {1,2,4,8,16,32,64,128,256} = 9 INSERT shapes; allow slack for the
    # store's other (read/DDL) statements. The point: O(1) in #symbols, NOT the 400 of the old path.
    assert distinct <= 32, f"compiled-statement cache grew to {distinct} (expected O(log n), ~a dozen)"

    # And the data is fully and correctly present (bounded memory did not cost correctness).
    assert _count(store) == sum(range(1, 401))
    store.close()


def test_symbols_with_feature_returns_distinct_symbols_per_feature() -> None:
    """symbols_with_feature(feature) is the ingest resume probe: the SET of symbols holding >=1 row for that
    feature, per feature (a symbol with only feature X is NOT reported for feature Y)."""
    store = TimescaleFeatureStore("sqlite://", table="bybit_features")
    store.upsert_feature_rows([
        ("AUSDT", 19000 * 86400, "funding", 0.1),
        ("AUSDT", 19001 * 86400, "funding", 0.2),   # same symbol, 2 rows -> still one entry in the set
        ("BUSDT", 19000 * 86400, "funding", 0.3),
        ("CUSDT", 19000 * 86400, "close", 100.0),   # close only -> not in the funding set
        ("AUSDT", 19000 * 86400, "close", 50.0),    # A also has close
    ])
    assert store.symbols_with_feature("funding") == {"AUSDT", "BUSDT"}
    assert store.symbols_with_feature("close") == {"AUSDT", "CUSDT"}
    assert store.symbols_with_feature("spot_close") == set()   # no rows -> empty set
    store.close()
