"""
Tests for data/store.py — Parquet-based raw message storage.
"""
import json
import os
import tempfile
import time

from data.store import RawMessageStore, read_range


def _make_parquet_dir():
    """Create a temporary directory for parquet storage."""
    return tempfile.mkdtemp()


def test_store_basic_write_read():
    """Write messages, flush, read them back."""
    tmpdir = tempfile.mkdtemp()
    store = RawMessageStore(data_dir=tmpdir, flush_interval_sec=0.5)
    store.start()

    for i in range(10):
        store.push(
            channel="l2book",
            coin="BTC",
            exchange_ts=1718000000000 + i * 1000,
            payload={"type": "snapshot", "levels": [[{"px": "50000", "sz": "1.0"}], [{"px": "50001", "sz": "0.5"}]]},
        )

    time.sleep(1.5)
    store.stop()
    assert store.total_written == 10

    from datetime import date
    today = date.today().isoformat()
    rows = read_range(tmpdir, "l2book", "BTC", today, today)
    assert len(rows) == 10
    assert rows[0]["payload"]["type"] == "snapshot"
    assert rows[0]["exchange_ts"] == 1718000000000
    assert rows[-1]["exchange_ts"] == 1718000009000


def test_store_multiple_channels():
    """Write to different channels and verify partitioning."""
    tmpdir = tempfile.mkdtemp()
    store = RawMessageStore(data_dir=tmpdir, flush_interval_sec=0.5)
    store.start()

    channels = ["l2book", "trades", "funding", "mark", "open_interest"]
    for ch in channels:
        for i in range(3):
            store.push(channel=ch, coin="BTC", exchange_ts=1718000000000 + i * 1000, payload={"ch": ch, "i": i})

    time.sleep(1.5)
    store.stop()
    assert store.total_written == 15

    from datetime import date
    today = date.today().isoformat()
    for ch in channels:
        rows = read_range(tmpdir, ch, "BTC", today, today)
        assert len(rows) == 3, f"Expected 3 rows for {ch}, got {len(rows)}"


def test_store_append_to_existing():
    """Write in two batches to same file — should append."""
    tmpdir = tempfile.mkdtemp()
    store = RawMessageStore(data_dir=tmpdir, flush_interval_sec=0.3)
    store.start()

    for i in range(5):
        store.push(channel="trades", coin="ETH", exchange_ts=1718000000000 + i * 1000, payload={"batch": 1, "i": i})

    time.sleep(1)
    store.stop()

    store2 = RawMessageStore(data_dir=tmpdir, flush_interval_sec=0.3)
    store2.start()
    for i in range(5):
        store2.push(channel="trades", coin="ETH", exchange_ts=1718000005000 + i * 1000, payload={"batch": 2, "i": i})
    time.sleep(1)
    store2.stop()

    from datetime import date
    today = date.today().isoformat()
    rows = read_range(tmpdir, "trades", "ETH", today, today)
    assert len(rows) == 10
    batches = [r["payload"]["batch"] for r in rows]
    assert batches == [1] * 5 + [2] * 5


def test_store_queue_full_does_not_crash():
    """Small queue — push many, ensure no crash."""
    tmpdir = tempfile.mkdtemp()
    store = RawMessageStore(data_dir=tmpdir, flush_interval_sec=60.0, max_queue_size=10)
    store.start()

    for i in range(1000):
        store.push(channel="trades", coin="BTC", exchange_ts=1718000000000 + i, payload={"i": i})

    store.stop()
    assert store.total_written >= 0  # some may be lost, but no crash


def test_read_range_empty():
    """Read a date range with no data returns empty list."""
    tmpdir = tempfile.mkdtemp()
    rows = read_range(tmpdir, "nonexistent", "BTC", "2020-01-01", "2020-01-02")
    assert rows == []