"""Tests for all sink implementations (CSV, Parquet, NDJSON, streaming, etc.)."""
from __future__ import annotations

import csv
import io
import json
import tempfile
from pathlib import Path
from typing import Any

from universe_reads.sink import (
    CollectorSink,
    CountSink,
    CSVSink,
    NDJSONSink,
    StreamCSVSink,
    StreamJSONSink,
    _normalise_record,
)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _make_records() -> list[tuple[str, Any]]:
    """Return a small list of (key, record) tuples for testing."""
    return [
        ("K1", "hello world"),
        ("K2", {"name": "Alice", "score": 95}),
        ("K3", "another plain string"),
    ]


# ---------------------------------------------------------------------------
# _normalise_record
# ---------------------------------------------------------------------------

def test_normalise_bytes():
    assert _normalise_record(b"hello") == "hello"


def test_normalise_dict_passthrough():
    d = {"a": 1}
    assert _normalise_record(d) is d


def test_normalise_dynarray_list():
    """Objects with a .list attribute (like uopy DynArray) should return .list."""
    class FakeDyn:
        list = ["val1", "val2"]
    assert _normalise_record(FakeDyn()) == ["val1", "val2"]


def test_normalise_weird_type():
    """Unknown types should be stringified."""
    class Custom:
        def __str__(self):
            return "custom-str"
    assert _normalise_record(Custom()) == "custom-str"


# ---------------------------------------------------------------------------
# CountSink
# ---------------------------------------------------------------------------

def test_count_sink():
    s = CountSink()
    for k, r in _make_records():
        s.on_record(key=k, record=r)
    s.close()
    assert s.count == 3


# ---------------------------------------------------------------------------
# CollectorSink
# ---------------------------------------------------------------------------

def test_collector_sink():
    s = CollectorSink()
    for k, r in _make_records():
        s.on_record(key=k, record=r)
    s.close()
    assert s.count == 3
    assert set(s.data.keys()) == {"K1", "K2", "K3"}
    assert s.data["K2"] == {"name": "Alice", "score": 95}


# ---------------------------------------------------------------------------
# NDJSONSink
# ---------------------------------------------------------------------------

def test_ndjson_sink(tmp_path: Path):
    path = tmp_path / "out.ndjson"
    s = NDJSONSink(path)
    for k, r in _make_records():
        s.on_record(key=k, record=r)
    s.close()

    lines = path.read_text().strip().splitlines()
    assert len(lines) == 3

    first = json.loads(lines[0])
    assert first["key"] == "K1"
    assert first["record"] == "hello world"

    second = json.loads(lines[1])
    assert second["key"] == "K2"
    assert second["record"] == {"name": "Alice", "score": 95}


# ---------------------------------------------------------------------------
# CSVSink — dict records
# ---------------------------------------------------------------------------

def test_csv_sink_dict_records(tmp_path: Path):
    path = tmp_path / "out.csv"
    s = CSVSink(path)
    s.on_record(key="K1", record={"name": "Alice", "score": 95})
    s.on_record(key="K2", record={"name": "Bob", "score": 88})
    s.close()

    rows = list(csv.DictReader(path.open()))
    assert len(rows) == 2
    assert rows[0]["key"] == "K1"
    assert rows[0]["name"] == "Alice"
    assert rows[1]["score"] == "88"


# ---------------------------------------------------------------------------
# CSVSink — string records
# ---------------------------------------------------------------------------

def test_csv_sink_string_records(tmp_path: Path):
    path = tmp_path / "out.csv"
    s = CSVSink(path)
    s.on_record(key="K1", record="raw value 1")
    s.on_record(key="K2", record="raw value 2")
    s.close()

    reader = csv.reader(path.open())
    header = next(reader)
    assert header == ["key", "record"]
    row1 = next(reader)
    assert row1 == ["K1", "raw value 1"]


# ---------------------------------------------------------------------------
# StreamJSONSink
# ---------------------------------------------------------------------------

def test_stream_json_sink():
    buf = io.StringIO()
    s = StreamJSONSink(fp=buf)
    s.on_record(key="K1", record="val1")
    s.on_record(key="K2", record={"x": 1})
    s.close()

    lines = buf.getvalue().strip().splitlines()
    assert len(lines) == 2
    assert json.loads(lines[0]) == {"key": "K1", "record": "val1"}
    assert json.loads(lines[1]) == {"key": "K2", "record": {"x": 1}}


# ---------------------------------------------------------------------------
# StreamCSVSink — string records
# ---------------------------------------------------------------------------

def test_stream_csv_sink_strings():
    buf = io.StringIO()
    s = StreamCSVSink(fp=buf)
    s.on_record(key="K1", record="val1")
    s.on_record(key="K2", record="val2")
    s.close()

    buf.seek(0)
    reader = csv.reader(buf)
    header = next(reader)
    assert header == ["key", "record"]
    assert next(reader) == ["K1", "val1"]
    assert next(reader) == ["K2", "val2"]


# ---------------------------------------------------------------------------
# StreamCSVSink — dict records
# ---------------------------------------------------------------------------

def test_stream_csv_sink_dicts():
    buf = io.StringIO()
    s = StreamCSVSink(fp=buf)
    s.on_record(key="K1", record={"a": 1, "b": 2})
    s.on_record(key="K2", record={"a": 3, "b": 4})
    s.close()

    buf.seek(0)
    rows = list(csv.DictReader(buf))
    assert len(rows) == 2
    assert rows[0]["key"] == "K1"
    assert rows[0]["a"] == "1"
    assert rows[1]["b"] == "4"


# ---------------------------------------------------------------------------
# ParquetSink  (only if pyarrow is installed)
# ---------------------------------------------------------------------------

def test_parquet_sink_dict_records(tmp_path: Path):
    try:
        import pyarrow.parquet as pq  # noqa: F401
    except ImportError:
        print("SKIP: pyarrow not installed")
        return

    from universe_reads.sink import ParquetSink

    path = tmp_path / "out.parquet"
    s = ParquetSink(path)
    s.on_record(key="K1", record={"name": "Alice", "score": 95})
    s.on_record(key="K2", record={"name": "Bob", "score": 88})
    s.close()

    table = pq.read_table(str(path))
    assert table.num_rows == 2
    assert "key" in table.column_names
    assert "name" in table.column_names


def test_parquet_sink_string_records(tmp_path: Path):
    try:
        import pyarrow.parquet as pq  # noqa: F401
    except ImportError:
        print("SKIP: pyarrow not installed")
        return

    from universe_reads.sink import ParquetSink

    path = tmp_path / "out.parquet"
    s = ParquetSink(path)
    s.on_record(key="K1", record="raw1")
    s.on_record(key="K2", record="raw2")
    s.close()

    table = pq.read_table(str(path))
    assert table.num_rows == 2
    assert set(table.column_names) == {"key", "record"}


# ---------------------------------------------------------------------------
# Entry point  (for running directly: PYTHONPATH=src:. python tests/test_sinks.py)
# ---------------------------------------------------------------------------

def main() -> None:
    with tempfile.TemporaryDirectory() as td:
        tp = Path(td)

        test_normalise_bytes()
        test_normalise_dict_passthrough()
        test_normalise_dynarray_list()
        test_normalise_weird_type()
        print("PASS: _normalise_record")

        test_count_sink()
        print("PASS: CountSink")

        test_collector_sink()
        print("PASS: CollectorSink")

        test_ndjson_sink(tp / "ndjson")
        print("PASS: NDJSONSink")

        test_csv_sink_dict_records(tp / "csv_dict")
        print("PASS: CSVSink (dict records)")

        test_csv_sink_string_records(tp / "csv_str")
        print("PASS: CSVSink (string records)")

        test_stream_json_sink()
        print("PASS: StreamJSONSink")

        test_stream_csv_sink_strings()
        print("PASS: StreamCSVSink (strings)")

        test_stream_csv_sink_dicts()
        print("PASS: StreamCSVSink (dicts)")

        test_parquet_sink_dict_records(tp / "pq_dict")
        print("PASS: ParquetSink (dict records)")

        test_parquet_sink_string_records(tp / "pq_str")
        print("PASS: ParquetSink (string records)")

    print("\nAll sink tests passed.")


if __name__ == "__main__":
    main()
