"""Tests for CLI argument parsing — verifies the new flags are recognized."""
from __future__ import annotations

from universe_reads.cli import _parse_args


# ---------------------------------------------------------------------------
# New flags: --format, --stream, --where, --select, --partition
# ---------------------------------------------------------------------------

def test_format_ndjson():
    ns = _parse_args(["--dry-run", "10", "--format", "ndjson"])
    assert ns.output_format == "ndjson"


def test_format_csv():
    ns = _parse_args(["--dry-run", "10", "--format", "csv"])
    assert ns.output_format == "csv"


def test_format_parquet():
    ns = _parse_args(["--dry-run", "10", "--format", "parquet"])
    assert ns.output_format == "parquet"


def test_format_jsonl():
    ns = _parse_args(["--dry-run", "10", "--format", "jsonl"])
    assert ns.output_format == "jsonl"


def test_format_default():
    ns = _parse_args(["--dry-run", "10"])
    assert ns.output_format == "ndjson"


def test_stream_flag():
    ns = _parse_args(["--dry-run", "10", "--stream"])
    assert ns.stream is True


def test_stream_default():
    ns = _parse_args(["--dry-run", "10"])
    assert ns.stream is False


def test_where_clause():
    ns = _parse_args(["--dry-run", "10", "--where", "WITH STATUS = 'ACTIVE'"])
    assert ns.where_clause == "WITH STATUS = 'ACTIVE'"


def test_where_default():
    ns = _parse_args(["--dry-run", "10"])
    assert ns.where_clause is None


def test_select_override():
    ns = _parse_args(["--dry-run", "10", "--select", "SELECT STUDENTS WITH GPA >= '3.0'"])
    assert ns.select_override == "SELECT STUDENTS WITH GPA >= '3.0'"


def test_select_default():
    ns = _parse_args(["--dry-run", "10"])
    assert ns.select_override is None


def test_partition_prefix():
    ns = _parse_args(["--dry-run", "10", "--partition", "prefix"])
    assert ns.partition == "prefix"


def test_partition_hash():
    ns = _parse_args(["--dry-run", "10", "--partition", "hash"])
    assert ns.partition == "hash"


def test_partition_default():
    ns = _parse_args(["--dry-run", "10"])
    assert ns.partition is None


# ---------------------------------------------------------------------------
# Existing flags still work
# ---------------------------------------------------------------------------

def test_existing_retries():
    ns = _parse_args(["--dry-run", "10", "--retries", "5"])
    assert ns.retries == 5


def test_existing_sessions():
    ns = _parse_args(["--dry-run", "10", "--sessions", "3"])
    assert ns.sessions == 3


def test_existing_batch_size():
    ns = _parse_args(["--dry-run", "10", "--batch-size", "200"])
    assert ns.batch_size == 200


# ---------------------------------------------------------------------------
# Inspect sub-command
# ---------------------------------------------------------------------------

def test_inspect_subcommand():
    ns = _parse_args(["inspect", "ATTENDANCE", "--session-factory", "my.mod:factory"])
    assert ns.command == "inspect"
    assert ns.file_name == "ATTENDANCE"
    assert ns.session_factory == "my.mod:factory"


def test_inspect_samples():
    ns = _parse_args(["inspect", "ATTENDANCE", "--samples", "5"])
    assert ns.samples == 5


def test_inspect_json_flag():
    ns = _parse_args(["inspect", "FILE", "--json"])
    assert ns.json is True


# ---------------------------------------------------------------------------
# Combined flags
# ---------------------------------------------------------------------------

def test_combined_flags():
    ns = _parse_args([
        "--dry-run", "100",
        "--format", "csv",
        "--stream",
        "--partition", "hash",
        "--sessions", "3",
        "--retries", "2",
        "--where", "WITH CLASS = 'CS101'",
    ])
    assert ns.output_format == "csv"
    assert ns.stream is True
    assert ns.partition == "hash"
    assert ns.sessions == 3
    assert ns.retries == 2
    assert ns.where_clause == "WITH CLASS = 'CS101'"
    assert ns.dry_run == 100


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

def main() -> None:
    test_format_ndjson()
    test_format_csv()
    test_format_parquet()
    test_format_jsonl()
    test_format_default()
    print("PASS: --format variants")

    test_stream_flag()
    test_stream_default()
    print("PASS: --stream")

    test_where_clause()
    test_where_default()
    print("PASS: --where")

    test_select_override()
    test_select_default()
    print("PASS: --select")

    test_partition_prefix()
    test_partition_hash()
    test_partition_default()
    print("PASS: --partition")

    test_existing_retries()
    test_existing_sessions()
    test_existing_batch_size()
    print("PASS: existing flags")

    test_inspect_subcommand()
    test_inspect_samples()
    test_inspect_json_flag()
    print("PASS: inspect subcommand")

    test_combined_flags()
    print("PASS: combined flags")

    print("\nAll CLI arg tests passed.")


if __name__ == "__main__":
    main()
