#!/usr/bin/env python3
"""Run the DS4 token oracles, then the existing endpoint smoke corpus."""

import argparse
import pathlib
import subprocess
import sys


ROOT = pathlib.Path(__file__).resolve().parent.parent
HARDWARE_TESTS = (
    "flash_resident_and_ssd_streaming_choose_the_same_tokens",
    "flash_ssd_streaming_maps_batched_prefill_layers",
    "resident_multi_session_switching_preserves_each_kv_frontier",
    "legacy_mtp_runs_a_target_owned_greedy_cycle",
    "dspark_runs_a_target_owned_greedy_cycle",
    "ssd_streaming_supports_legacy_mtp_and_dspark",
    "directional_steering_matches_the_ds4_token_oracle",
    "resident_and_streamed_glm_match_ds4_decode_oracles",
    "streamed_glm_uses_ds4_indexed_prefill_for_long_prompts",
    "glm_mtp_preserves_target_tokens_and_drafts",
    "pro_ssd_streaming_full_layer_and_selected_batch_agree",
)
ENDPOINT_SCRIPTS = (
    "endpoint_parity.py",
    "endpoint_reasoning.py",
    "endpoint_continuation.py",
)


def run(command):
    print("+", " ".join(map(str, command)), flush=True)
    subprocess.run(command, cwd=ROOT, check=True)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--endpoint",
        help="also run the existing endpoint parity scripts against this URL",
    )
    parser.add_argument(
        "--skip-hardware",
        action="store_true",
        help="skip tests requiring the local Flash/MTP/DSpark GGUF fixtures",
    )
    args = parser.parse_args()

    if not args.skip_hardware:
        for test in HARDWARE_TESTS:
            run(
                [
                    "cargo",
                    "test",
                    "--all-features",
                    test,
                    "--",
                    "--ignored",
                    "--nocapture",
                ]
            )
    if args.endpoint:
        for script in ENDPOINT_SCRIPTS:
            run([sys.executable, ROOT / "scripts" / script, args.endpoint])


if __name__ == "__main__":
    main()
