import pyarrow.parquet as pq
import pytest

from app.events import COLUMNS, ClientEvent, to_row
from app.sink import HubParquetSink, LocalParquetSink, rows_to_parquet_bytes


def rows(n):
    return [to_row(ClientEvent(event_type="session", model_family="Qwen", quantization="4-bit")) for _ in range(n)]


def test_batch_becomes_one_shard(tmp_path):
    s = LocalParquetSink(tmp_path, flush_seconds=999)
    s.add(rows(7))
    s.add(rows(5))
    assert s.flush()
    shards = list((tmp_path / "data/incoming").rglob("*.parquet"))
    assert len(shards) == 1
    t = pq.read_table(shards[0])
    assert t.num_rows == 12 and t.column_names == list(COLUMNS)
    assert s.flush() and len(list((tmp_path / "data/incoming").rglob("*.parquet"))) == 1  # empty flush = no shard
    s.add(rows(1))
    s.flush()
    assert len(list((tmp_path / "data/incoming").rglob("*.parquet"))) == 2  # appends, never rewrites
    assert len(s.read_existing()) == 13


class FlakyApi:
    def __init__(self, fail_times):
        self.fail_times = fail_times
        self.commits = []
        self.token = None

    def create_commit(self, repo_id, repo_type, operations, commit_message):
        if self.fail_times > 0:
            self.fail_times -= 1
            raise ConnectionError("hub down")
        assert repo_type == "dataset"
        self.commits.append([op.path_in_repo for op in operations])


def test_hub_failure_keeps_spool_and_retries(tmp_path):
    api = FlakyApi(fail_times=2)
    s = HubParquetSink("org/data", None, spool_dir=tmp_path, api=api, flush_seconds=999)
    s.add(rows(3))
    assert not s.flush() and s.pending == 3 and not s.status()["healthy"]
    # a restart in between must not lose rows
    s2 = HubParquetSink("org/data", None, spool_dir=tmp_path, api=api, flush_seconds=999)
    assert s2.pending == 3
    assert not s2.flush()
    assert s2.flush() and s2.pending == 0 and s2.status()["healthy"]
    assert len(api.commits) == 1 and len(api.commits[0]) == 1
    assert api.commits[0][0].startswith("data/incoming/") and api.commits[0][0].endswith(".parquet")
    s3 = HubParquetSink("org/data", None, spool_dir=tmp_path, api=api, flush_seconds=999)
    assert s3.pending == 0


def test_local_refuses_overwrite(tmp_path):
    s = LocalParquetSink(tmp_path, flush_seconds=999)
    data = rows_to_parquet_bytes(rows(1))
    s._write_shard("data/incoming/x.parquet", data)
    with pytest.raises(FileExistsError):
        s._write_shard("data/incoming/x.parquet", data)


def test_dataset_schema_is_deliberate():
    """Changing columns changes the public dataset schema: bump this list on purpose, and
    migrate or document existing shards (see the dataset card's versioning section)."""
    assert len(COLUMNS) == 67
    assert {"perplexity", "perplexity_stderr", "eval_dataset", "eval_tokens"} <= set(COLUMNS)


def test_size_triggered_flush(tmp_path, monkeypatch):
    import time
    import app.sink as sink_mod
    monkeypatch.setattr(sink_mod, "FLUSH_ROWS", 5)
    s = LocalParquetSink(tmp_path, flush_seconds=999)
    s.add(rows(3))
    assert s.pending == 3
    s.add(rows(3))  # crosses FLUSH_ROWS -> flushes in the background
    for _ in range(50):
        if s.pending == 0:
            break
        time.sleep(0.05)
    assert s.pending == 0 and len(list((tmp_path / "data/incoming").rglob("*.parquet"))) == 1


def test_reads_monthly_and_incoming(tmp_path):
    from app.sink import rows_to_parquet_bytes as rb
    (tmp_path / "data/events").mkdir(parents=True)
    (tmp_path / "data/events/2026-08.parquet").write_bytes(rb(rows(4)))
    s = LocalParquetSink(tmp_path, flush_seconds=999)
    s.add(rows(2))
    s.flush()
    assert len(s.read_existing()) == 6


def test_reload_neither_misses_nor_double_counts(tmp_path):
    import time
    s = LocalParquetSink(tmp_path, flush_seconds=999)
    s.add(rows(3))
    stored = s.read_existing()
    assert len(stored) + len(s.live_rows()) == 3          # accepted, not yet written
    s.flush()
    assert len(s.read_existing()) + len(s.live_rows()) == 6  # written but not yet reloaded: still visible (overlap is fine before a reload)
    started = time.time()
    stored = s.read_existing()
    s.forget_flushed_before(started)
    s.add(rows(2))
    assert len(stored) + len(s.live_rows()) == 5          # after reload: 3 stored + 2 new, no duplicates