from __future__ import annotations

import asyncio

import pytest

from gfa_core.mas.agents.base import AgentMessage, AgentState
from gfa_core.mas.agents.intent_v2 import IntentAgentV2
from gfa_core.mas.core.fast_intent import fast_intent_recognition
from gfa_core.mas.stage_catalog import complete_dependencies
from gfa_core.mas.utils.intent_parser import parse_intent_from_query


FULL_TARGETS = ["gwas_plot", "gwas_emplot", "gwas_glm"]
FULL_CHAIN = complete_dependencies(FULL_TARGETS)

CASES = [
    ("使用 TASSEL 计算亲缘关系矩阵", "gwas_kinship"),
    ("准备 GWAS 表型数据", "gwas_trait"),
    ("根据 Q 矩阵和 FAM 准备 GWAS 协变量", "gwas_cov"),
    ("使用 TASSEL 执行 GWAS MLM 分析", "gwas_mlm"),
    ("绘制 GWAS MLM Manhattan 和 QQ 图", "gwas_plot"),
    ("执行 EMMAX GWAS 分析", "gwas_emmax"),
    ("绘制 EMMAX Manhattan 和 QQ 图", "gwas_emplot"),
    ("使用 TASSEL 执行 GWAS GLM 分析", "gwas_glm"),
]


def _public(query: str) -> dict:
    agent = IntentAgentV2(object())
    state = AgentState(session_id="gwas-chain-intent", job_id="job_test")
    message = AgentMessage(
        sender="router",
        receiver="intent",
        message_type="recognize",
        payload={"query": query, "workspace": "/tmp/job_test"},
    )
    result = asyncio.run(agent.execute(state, message))
    return result.results["results"]["intent"]


@pytest.mark.parametrize("query", ["gwas", "执行gwas分析", "运行完整 GWAS", "gwas流程"])
def test_full_gwas_intent_is_the_complete_dependency_dag(query: str) -> None:
    parsed = parse_intent_from_query(query)
    fast = fast_intent_recognition(query)
    public = _public(query)

    assert parsed["stages"] == FULL_CHAIN
    assert parsed["primary_stage"] == "gwas_glm"
    assert parsed["need_clarify"] is False
    assert fast is not None
    assert fast["stages"] == FULL_CHAIN
    assert fast["stage"] == "gwas_glm"
    assert public["stages"] == FULL_CHAIN
    assert public["primary_stage"] == "gwas_glm"
    assert "蛋白" not in str((parsed, fast, public))


@pytest.mark.parametrize(("query", "target"), CASES)
def test_specific_gwas_intents_share_one_target_and_dependency_closure(query: str, target: str) -> None:
    expected = complete_dependencies([target])
    parsed = parse_intent_from_query(query)
    fast = fast_intent_recognition(query)
    public = _public(query)

    assert parsed["stages"] == expected
    assert parsed["primary_stage"] == target
    assert fast is not None
    assert fast["stages"] == expected
    assert fast["stage"] == target
    assert public["stages"] == expected
    assert public["primary_stage"] == target