"""Build router data from SWE-bench harness results."""

from __future__ import annotations

import argparse
from pathlib import Path

import yaml

from embergraphrouter.data_processing.swebench import (
    SWEbenchRouterDataBuilder,
    load_results_jsonl,
    load_swebench_examples,
)


def main() -> None:
    parser = argparse.ArgumentParser(description="Build router data from SWE-bench runs.")
    parser.add_argument(
        "--config",
        type=Path,
        default=Path("embergraphrouter/configs/swebench.yaml"),
        help="Path to SWE-bench config YAML.",
    )
    args = parser.parse_args()

    config = yaml.safe_load(args.config.read_text(encoding="utf-8"))
    dataset = config["dataset"]
    split = config.get("split", "dev")
    sample_size = config.get("sample_size")
    seed = int(config.get("seed", 0))

    examples = load_swebench_examples(
        dataset_name=dataset,
        split=split,
        sample_size=sample_size,
        shuffle=True,
        seed=seed,
    )
    examples_by_task = {example.task_id: example for example in examples}

    results_path = Path(config["results_path"])
    results = load_results_jsonl(results_path, examples_by_task)

    builder = SWEbenchRouterDataBuilder(
        llm_description_path=Path(config["llm_description_path"]),
        router_output_path=Path(config["router_output_path"]),
        llm_embedding_path=Path(config["llm_embedding_path"]),
    )
    output_path = builder.build(results)
    print(f"Wrote router dataset to {output_path}")


if __name__ == "__main__":
    main()
