from gfa_core.mas.agents.planning_v2 import PlanningAgentV2


def test_ppi_species_does_not_change_unrelated_stage_biological_scope() -> None:
    planner = PlanningAgentV2.__new__(PlanningAgentV2)

    stages = planner._generate_default_stages(
        ["motif", "ppi"],
        "Ft.SPL",
        intent_params={
            "geneFamily": "SPL",
            "targetSpecies": ["Ft"],
            "scopeMode": "target_only",
            "analysisRoute": "single_species",
            "ppiSpecies": "ath",
        },
    )

    stage_by_id = {stage["id"]: stage for stage in stages}
    motif_params = stage_by_id["motif"]["params"]
    ppi_params = stage_by_id["ppi"]["params"]

    assert motif_params["targetSpecies"] == ["Ft"]
    assert motif_params["scopeMode"] == "target_only"
    assert motif_params["analysisRoute"] == "single_species"
    assert not motif_params.get("referenceSpecies")
    assert ppi_params["ppiSpecies"] == "ath"


def test_ppi_plan_is_target_only_while_kaks_keeps_project_reference() -> None:
    planner = PlanningAgentV2.__new__(PlanningAgentV2)

    stages = planner._generate_default_stages(
        ["ppi", "kaks"],
        "Ft.SPL",
        intent_params={
            "geneFamily": "SPL",
            "targetSpecies": ["Ft"],
            "referenceSpecies": ["At"],
            "scopeMode": "target_plus_reference",
            "analysisRoute": "multi_species",
        },
    )

    stage_by_id = {stage["id"]: stage for stage in stages}
    ppi_params = stage_by_id["ppi"]["params"]
    kaks_params = stage_by_id["kaks"]["params"]

    assert ppi_params["targetSpecies"] == ["Ft"]
    assert ppi_params["referenceSpecies"] == []
    assert ppi_params["scopeMode"] == "target_only"
    assert ppi_params["analysisRoute"] == "single_species"
    assert kaks_params["referenceSpecies"] == ["At"]
    assert kaks_params["scopeMode"] == "target_plus_reference"