"""Tests for CLI parsing and streaming profile execution."""

from __future__ import annotations

import json
from argparse import Namespace
from pathlib import Path

import pyarrow as pa
import pyarrow.parquet as pq

from swe_data_processing.cli import (
    _run_parallel_parquet_profile,
    _run_streaming_stage,
    build_parser,
)


def test_streaming_stage_writes_each_record(tmp_path: Path) -> None:
    input_path = tmp_path / "input.jsonl"
    output_path = tmp_path / "output.jsonl"
    input_path.write_text(
        "".join(json.dumps({"trajectory_id": f"s-{i}"}) + "\n" for i in range(10)),
        encoding="utf-8",
    )
    status = _run_streaming_stage(
        input_path=input_path,
        output_path=output_path,
        errors_path=None,
        limit=None,
        resume=False,
        workers=3,
        processor=lambda record: {"sample_id": record["trajectory_id"]},
    )
    assert status == 0
    assert len(output_path.read_text().splitlines()) == 10


def test_cli_has_only_deterministic_pipeline_commands() -> None:
    parser = build_parser()
    assert parser.parse_args(
        ["profile", "--input", "in.jsonl", "--output", "out.jsonl"]
    ).command == "profile"
    help_text = parser.format_help()
    assert "classify" not in help_text
    assert "repair" not in help_text
    assert "materialize" in help_text


def test_parallel_parquet_profile_uses_atomic_shard_parts(tmp_path: Path) -> None:
    dataset = tmp_path / "dataset"
    dataset.mkdir()
    for index in range(2):
        pq.write_table(
            pa.Table.from_pylist(
                [{"trajectory_id": f"sample-{index}", "resolved": 1}]
            ),
            dataset / f"part-{index}.parquet",
        )
    output = tmp_path / "metrics.jsonl"
    errors = tmp_path / "errors.jsonl"
    status = _run_parallel_parquet_profile(
        Namespace(
            input=dataset,
            output=output,
            errors=errors,
            tokenizer_json=None,
            workers=2,
            resume=False,
        )
    )

    assert status == 0
    assert [json.loads(line)["sample_id"] for line in output.read_text().splitlines()] == [
        "sample-0",
        "sample-1",
    ]
    assert errors.read_text() == ""