from aitop.domain.job_topology import annotate_job_topology, summarize_job_topology


def test_annotate_job_topology_detects_distributed_group_and_straggler() -> None:
    processes = [
        {
            "pid": 1001,
            "ppid": 900,
            "name": "python",
            "cmdline": "torchrun --rdzv-id exp-42 --rank 0 train.py",
            "cpu_percent": 82.0,
        },
        {
            "pid": 1002,
            "ppid": 900,
            "name": "python",
            "cmdline": "torchrun --rdzv-id exp-42 --rank 1 train.py",
            "cpu_percent": 79.0,
        },
        {
            "pid": 1003,
            "ppid": 900,
            "name": "python",
            "cmdline": "torchrun --rdzv-id exp-42 --rank 2 train.py",
            "cpu_percent": 4.0,
        },
    ]
    gpu_processes = [
        {"pid": 1001, "sm_util": 93.0},
        {"pid": 1002, "sm_util": 88.0},
        {"pid": 1003, "sm_util": 1.0},
    ]

    enriched = annotate_job_topology(processes, gpu_processes)

    assert {proc["job_id"] for proc in enriched} == {"run-exp-42"}
    assert {proc["job_size"] for proc in enriched} == {3}
    assert {proc["is_distributed_job"] for proc in enriched} == {True}
    by_rank = {proc["job_rank"]: proc for proc in enriched}
    assert by_rank[2]["is_job_straggler"] is True
    assert by_rank[0]["is_job_straggler"] is False
    assert by_rank[2]["job_straggler_score"] > by_rank[0]["job_straggler_score"]


def test_annotate_job_topology_falls_back_to_parent_and_script_identity() -> None:
    processes = [
        {
            "pid": 2101,
            "ppid": 2000,
            "name": "python",
            "cmdline": "python /workspace/train.py --epochs 5",
            "cpu_percent": 22.0,
        },
        {
            "pid": 2102,
            "ppid": 2000,
            "name": "python",
            "cmdline": "python /workspace/train.py --epochs 5",
            "cpu_percent": 25.0,
        },
    ]

    enriched = annotate_job_topology(processes, gpu_processes=[])

    job_ids = {proc["job_id"] for proc in enriched}
    assert len(job_ids) == 1
    only_job_id = next(iter(job_ids))
    assert only_job_id.startswith("ppid-2000-train.py")


def test_summarize_job_topology_prioritizes_jobs_with_stragglers() -> None:
    processes = annotate_job_topology(
        [
            {
                "pid": 3101,
                "ppid": 3000,
                "name": "python",
                "cmdline": "torchrun --rdzv-id hot-job --rank 0 train.py",
                "cpu_percent": 70.0,
            },
            {
                "pid": 3102,
                "ppid": 3000,
                "name": "python",
                "cmdline": "torchrun --rdzv-id hot-job --rank 1 train.py",
                "cpu_percent": 2.0,
            },
            {
                "pid": 3201,
                "ppid": 3001,
                "name": "python",
                "cmdline": "python /srv/eval.py",
                "cpu_percent": 95.0,
            },
        ],
        gpu_processes=[
            {"pid": 3101, "sm_util": 90.0},
            {"pid": 3102, "sm_util": 1.0},
            {"pid": 3201, "sm_util": 60.0},
        ],
    )

    summaries = summarize_job_topology(processes, max_jobs=3)

    assert summaries
    assert summaries[0]["job_id"] == "run-hot-job"
    assert summaries[0]["stragglers"] == 1
    assert summaries[0]["is_distributed"] is True
