from langchain_core.messages import HumanMessage

import app.agents.graph as graph_module
from app.agents.orchestration.scheduler import build_subtask_specs
from app.agents.orchestration.task_graph import build_bounded_task_graph
from app.agents.schemas.orchestration import TaskGraph, assess_parallel_worthiness
from app.agents.state import initial_state


def test_build_subtask_specs_keeps_dependencies_and_contract_fields():
    worthiness = assess_parallel_worthiness(
        "先查资料、再分析风险、再安排下周计划",
        retrospective_count=2,
        skill_count=1,
    )
    task_graph = build_bounded_task_graph(
        query_text="先查资料、再分析风险、再安排下周计划",
        parallel_worthiness=worthiness,
    )

    specs = build_subtask_specs(task_graph, query_text="先查资料、再分析风险、再安排下周计划")

    assert specs
    assert all(spec.parent_run_id == task_graph.graph_id for spec in specs)
    assert all(isinstance(spec.context_slice, dict) for spec in specs)
    assert all(spec.expected_output_schema for spec in specs)
    assert any(spec.dependencies for spec in specs)


async def test_run_collaboration_flow_uses_task_graph_plan_and_records_subtask_events(monkeypatch):
    async def fake_run_sub_commander(
        state,
        assigned_role,
        _system_prompt,
        task_goal,
        **_kwargs,
    ):
        state["final_response"] = f"{assigned_role.value} handled: {task_goal}"
        return state

    monkeypatch.setattr(graph_module, "_run_sub_commander", fake_run_sub_commander)

    state = initial_state("u1", "c1")
    state["messages"] = [HumanMessage(content="先查资料、再分析风险、再安排下周计划")]
    state["current_datetime_context"] = "CURRENT_TIME: 2026-03-28T12:00:00+08:00"
    state["task_graph"] = TaskGraph.model_validate(
        build_bounded_task_graph(
            query_text="先查资料、再分析风险、再安排下周计划",
            parallel_worthiness=assess_parallel_worthiness(
                "先查资料、再分析风险、再安排下周计划",
                retrospective_count=2,
                skill_count=1,
            ),
        ).model_dump(mode="json")
    ).model_dump(mode="json")

    result = await graph_module._run_collaboration_flow(
        state,
        "先查资料、再分析风险、再安排下周计划",
    )

    assert result["scheduled_subtasks"]
    event_types = [item["event_type"] for item in result["event_trace"]]
    assert "agent.subtask.started" in event_types
    assert "agent.subtask.completed" in event_types
    assert result["task_results"]
