"""Tests for streaming mode (--stream) via dry-run.

Streaming writes to stdout through a drain thread. We capture stdout to verify
records flow correctly.
"""
from __future__ import annotations

import io
import json
import sys
import tempfile
from pathlib import Path

from universe_reads import run
from universe_reads.parallel_read import RunOptions, run_parallel_read


# ---------------------------------------------------------------------------
# Streaming dry-run: NDJSON
# ---------------------------------------------------------------------------

def test_stream_ndjson_dry_run():
    """Streaming dry-run should write NDJSON lines to stdout."""
    n = 50

    old_stdout = sys.stdout
    buf = io.StringIO()
    sys.stdout = buf
    try:
        results = run(
            file_name="",
            dry_run=n,
            sessions=2,
            count_only=False,
            stream=True,
            output_format="ndjson",
        )
    finally:
        sys.stdout = old_stdout

    output = buf.getvalue()
    lines = [l for l in output.strip().splitlines() if l.strip()]

    # Each line should be valid JSON with "key" and "record".
    parsed = []
    for line in lines:
        obj = json.loads(line)
        assert "key" in obj, f"Missing 'key' in {line}"
        assert "record" in obj, f"Missing 'record' in {line}"
        parsed.append(obj)

    # Total records from results should match n.
    total = sum(r.records for r in results)
    assert total == n, f"Expected {n} worker records, got {total}"

    # We should have received lines on stdout for each record.
    assert len(parsed) == n, f"Expected {n} NDJSON lines, got {len(parsed)}"


# ---------------------------------------------------------------------------
# Streaming dry-run: CSV
# ---------------------------------------------------------------------------

def test_stream_csv_dry_run():
    """Streaming CSV dry-run should write CSV rows to stdout."""
    n = 30

    old_stdout = sys.stdout
    buf = io.StringIO()
    sys.stdout = buf
    try:
        results = run(
            file_name="",
            dry_run=n,
            sessions=2,
            count_only=False,
            stream=True,
            output_format="csv",
        )
    finally:
        sys.stdout = old_stdout

    output = buf.getvalue()
    lines = [l for l in output.strip().splitlines() if l.strip()]

    # First line should be header.
    assert lines[0].startswith("key"), f"Expected CSV header, got: {lines[0]}"

    # Should have header + n data rows.
    assert len(lines) == n + 1, f"Expected {n + 1} CSV lines (incl header), got {len(lines)}"


# ---------------------------------------------------------------------------
# Streaming with fixture data
# ---------------------------------------------------------------------------

def test_stream_with_fixture_data():
    """Streaming with dry_run_data should emit the actual fixture records."""
    data = {
        "REC-001": "field1\xfefield2\xfefield3",
        "REC-002": "a\xfeb\xfec",
        "REC-003": "x\xfey\xfez",
    }

    old_stdout = sys.stdout
    buf = io.StringIO()
    sys.stdout = buf
    try:
        results = run(
            file_name="TEST",
            dry_run=len(data),
            dry_run_data=data,
            sessions=1,
            count_only=False,
            stream=True,
            output_format="ndjson",
        )
    finally:
        sys.stdout = old_stdout

    lines = [l for l in buf.getvalue().strip().splitlines() if l.strip()]
    assert len(lines) == 3

    keys_seen = set()
    for line in lines:
        obj = json.loads(line)
        keys_seen.add(obj["key"])
        # The record should be the actual fixture value.
        assert obj["record"] == data[obj["key"]]

    assert keys_seen == set(data.keys())


# ---------------------------------------------------------------------------
# Non-streaming: output to files with format
# ---------------------------------------------------------------------------

def test_output_csv_files_dry_run():
    """Dry-run with output_dir + csv format should create CSV files."""
    n = 100
    with tempfile.TemporaryDirectory() as td:
        outdir = Path(td) / "output"
        results = run(
            file_name="",
            dry_run=n,
            sessions=2,
            count_only=False,
            output_dir=str(outdir),
            output_format="csv",
        )

        total = sum(r.records for r in results)
        assert total == n

        # Check that CSV files were created.
        csv_files = list(outdir.glob("*.csv"))
        assert len(csv_files) >= 1, f"Expected CSV files in {outdir}"

        # Each file should have a header + data rows.
        for f in csv_files:
            lines = f.read_text().strip().splitlines()
            assert len(lines) >= 2, f"Expected header + data in {f}"


def test_output_ndjson_files_dry_run():
    """Dry-run with output_dir + ndjson format should create NDJSON files."""
    n = 100
    with tempfile.TemporaryDirectory() as td:
        outdir = Path(td) / "output"
        results = run(
            file_name="",
            dry_run=n,
            sessions=2,
            count_only=False,
            output_dir=str(outdir),
            output_format="ndjson",
        )

        total = sum(r.records for r in results)
        assert total == n

        ndjson_files = list(outdir.glob("*.ndjson"))
        assert len(ndjson_files) >= 1

        for f in ndjson_files:
            for line in f.read_text().strip().splitlines():
                obj = json.loads(line)
                assert "key" in obj


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

def main() -> None:
    test_stream_ndjson_dry_run()
    print("PASS: stream NDJSON dry-run")

    test_stream_csv_dry_run()
    print("PASS: stream CSV dry-run")

    test_stream_with_fixture_data()
    print("PASS: stream with fixture data")

    test_output_csv_files_dry_run()
    print("PASS: output CSV files dry-run")

    test_output_ndjson_files_dry_run()
    print("PASS: output NDJSON files dry-run")

    print("\nAll streaming/output tests passed.")


if __name__ == "__main__":
    main()
