import httpx
import pytest

import agent as framework
from agent.orchestration.api_v2 import create_orchestration_router
from agent.tools.builtin.knowledge import KnowledgeConfig
from agent.trace.api import router as trace_router
from agent.trace.api import set_trace_store

from test_mission_loop import MissionLoopLLM, ROOT_OBJECTIVE, ROOT_OBJECTIVE_V2


FastAPI = pytest.importorskip("fastapi").FastAPI


def _knowledge_off():
    return KnowledgeConfig(
        enable_extraction=False,
        enable_completion_extraction=False,
        enable_injection=False,
    )


def _all_keys(value):
    if isinstance(value, dict):
        return set(value).union(*(_all_keys(item) for item in value.values()))
    if isinstance(value, list):
        return set().union(*(_all_keys(item) for item in value))
    return set()


@pytest.mark.asyncio
async def test_public_framework_mission_is_fully_observable_without_ledger_access(
    tmp_path,
):
    root_trace_id = "framework-blackbox-root"
    trace_store = framework.FileSystemTraceStore(str(tmp_path))
    task_store = framework.FileSystemTaskStore(str(tmp_path))
    fake_llm = MissionLoopLLM(task_store)
    fake_llm.root_trace_id = root_trace_id
    runner = framework.AgentRunner(trace_store=trace_store, llm_call=fake_llm)
    coordinator = framework.wire_orchestration(
        runner,
        task_store,
        framework.FileSystemArtifactStore(str(tmp_path)),
    )

    result = await runner.run_result(
        [{"role": "user", "content": "run a generic observable mission"}],
        framework.RunConfig(
            agent_type="planner",
            completion_policy=framework.CompletionPolicy.EXPLICIT_VALIDATION,
            max_iterations=30,
            new_trace_id=root_trace_id,
            tools=["task_plan", "dispatch_tasks", "task_decide"],
            tool_groups=[],
            enable_memory=False,
            enable_research_flow=False,
            knowledge=_knowledge_off(),
            root_task_spec={
                "objective": ROOT_OBJECTIVE,
                "acceptance_criteria": [{
                    "criterion_id": "root-ready",
                    "description": "the root result is ready",
                    "hard": True,
                }],
                "context_refs": ["memory://input/root"],
            },
        ),
    )
    assert result["status"] == "completed"

    app = FastAPI()
    app.include_router(create_orchestration_router(coordinator))
    set_trace_store(trace_store)
    app.include_router(trace_router)
    http = httpx.AsyncClient(
        transport=httpx.ASGITransport(app=app), base_url="http://framework"
    )
    client = framework.OrchestrationClient("http://framework", client=http)

    mission = await client.get_mission_snapshot(root_trace_id)
    completion = await client.get_root_completion(root_trace_id)
    events = await client.list_events(root_trace_id, limit=1000)
    traces = await client.list_correlated_traces(
        root_trace_id=root_trace_id, limit=100
    )

    assert mission.root_objective == ROOT_OBJECTIVE_V2
    assert completion.root_objective == ROOT_OBJECTIVE_V2
    assert completion.status == framework.TaskStatus.COMPLETED
    assert completion.result_summary == f"submitted: {ROOT_OBJECTIVE_V2}"
    assert any(
        decision.action.value == "revise"
        and decision.task_id == mission.root_task_id
        for decision in mission.decisions
    )
    assert {event.schema_version for event in events.events} == {2}
    assert {trace["agent_role"] for trace in traces} == {
        "planner",
        "worker",
        "validator",
    }

    attempts = {item.attempt_id: item for item in mission.attempts}
    validations = {item.validation_id: item for item in mission.validations}
    trace_ids = {trace["trace_id"] for trace in traces}
    accepted = [
        decision for decision in mission.decisions
        if decision.action.value == "accept"
    ]
    assert len(accepted) == 2
    for decision in accepted:
        attempt = attempts[decision.attempt_id]
        validation = validations[decision.validation_id]
        artifact = await client.get_artifact_snapshot(
            root_trace_id, attempt.snapshot_id
        )
        assert attempt.task_id == decision.task_id
        assert validation.task_id == decision.task_id
        assert validation.attempt_id == attempt.attempt_id
        assert validation.snapshot_id == artifact.snapshot_id
        assert artifact.attempt_id == attempt.attempt_id
        assert attempt.worker_trace_id in trace_ids
        assert validation.validator_trace_id in trace_ids

    root_attempt = attempts[next(
        decision.attempt_id
        for decision in accepted
        if decision.task_id == mission.root_task_id
    )]
    child_accept = next(
        decision for decision in accepted
        if decision.task_id != mission.root_task_id
    )
    assert root_attempt.accepted_child_decision_ids == [child_accept.decision_id]

    forbidden = {"script", "workflow", "round", "branch", "page_nodes"}
    assert forbidden.isdisjoint(_all_keys(mission.model_dump(mode="json")))
    await http.aclose()