"""Tests for the streaming CLI runner."""

from __future__ import annotations

import json
import threading
import time
from pathlib import Path

from swe_data_processing.cli import _run_streaming_stage


def test_streaming_stage_processes_records_concurrently(tmp_path: Path) -> None:
    """Multiple workers run processors in parallel while one thread writes JSONL."""

    input_path = tmp_path / "input.jsonl"
    output_path = tmp_path / "output.jsonl"
    input_path.write_text(
        "".join(
            json.dumps({"trajectory_id": f"sample-{index}"}) + "\n" for index in range(20)
        ),
        encoding="utf-8",
    )
    lock = threading.Lock()
    active = 0
    peak_active = 0

    def processor(record: dict) -> dict:
        nonlocal active, peak_active
        with lock:
            active += 1
            peak_active = max(peak_active, active)
        time.sleep(0.01)
        with lock:
            active -= 1
        return {"sample_id": record["trajectory_id"]}

    status = _run_streaming_stage(
        input_path=input_path,
        output_path=output_path,
        errors_path=None,
        limit=None,
        resume=False,
        workers=4,
        stage_name="test",
        processor=processor,
    )

    written = [json.loads(line) for line in output_path.read_text().splitlines()]
    assert status == 0
    assert peak_active == 4
    assert {value["sample_id"] for value in written} == {
        f"sample-{index}" for index in range(20)
    }