"""Tests for plan-and-execute orchestration."""
from __future__ import annotations

import asyncio
import time

import pytest

from logos_adk.agent import Agent
from logos_adk.planning import PlanAndExecute, ReWOOAgent
from logos_adk.providers.mock import MockProvider
from logos_adk.tools import tool
from logos_adk.types import StreamDone, TextChunk


@pytest.mark.asyncio
async def test_plan_and_execute_runs_structured_plan_with_dependency_substitution():
    planner = Agent(
        "planner",
        model=MockProvider(
            responses=[
                """{
                  "steps": [
                    {"id": "E1", "instruction": "Collect the source facts", "depends_on": []},
                    {"id": "E2", "instruction": "Summarize #E1 for the user", "depends_on": ["E1"]}
                  ]
                }"""
            ]
        ),
    )
    executor_provider = MockProvider(
        responses=[
            "source facts: release is Friday",
            "summary prepared",
            "The release is Friday.",
        ]
    )
    executor = Agent("executor", model=executor_provider)
    workflow = PlanAndExecute("planner-bot", planner=planner, executor=executor)

    response = await workflow.run("When is the release?", session_id="session-1")

    assert response.content == "The release is Friday."
    trace = response.raw["plan_and_execute"]
    assert trace["plan"]["format"] == "structured"
    assert [item["id"] for item in trace["executed_steps"]] == ["E1", "E2"]
    second_execution_prompt = executor_provider.call_history[1]["messages"][-1].content
    assert "source facts: release is Friday" in second_execution_prompt
    assert "#E1" not in second_execution_prompt
    assert executor_provider.call_history[0]["kwargs"] == {}


@pytest.mark.asyncio
async def test_plan_and_execute_supports_fallback_text_plan_and_callables():
    executor_prompts: list[str] = []

    async def planner(prompt: str) -> str:
        return "1. Inspect the problem\n2. Draft the answer"

    async def executor(prompt: str) -> str:
        executor_prompts.append(prompt)
        return f"step-{len(executor_prompts)}"

    async def solver(prompt: str) -> str:
        assert "step-1" in prompt
        assert "step-2" in prompt
        return "final synthesized answer"

    workflow = PlanAndExecute(
        "callable-bot",
        planner=planner,
        executor=executor,
        solver=solver,
    )

    response = await workflow.run("Solve the task")

    assert response.content == "final synthesized answer"
    assert response.raw["plan_and_execute"]["plan"]["format"] == "fallback"
    assert len(response.raw["plan_and_execute"]["executed_steps"]) == 2
    assert len(executor_prompts) == 2


@pytest.mark.asyncio
async def test_plan_and_execute_rejects_future_dependencies():
    workflow = PlanAndExecute(
        "invalid-plan",
        planner=Agent(
            "planner",
            model=MockProvider(
                responses=[
                    """{
                      "steps": [
                        {"id": "E1", "instruction": "Use #E2", "depends_on": ["E2"]},
                        {"id": "E2", "instruction": "Second", "depends_on": []}
                      ]
                    }"""
                ]
            ),
        ),
        executor=Agent("executor", model=MockProvider(responses=["ignored"])),
    )

    with pytest.raises(ValueError, match="unknown or future"):
        await workflow.run("Bad plan")


@pytest.mark.asyncio
async def test_plan_and_execute_requires_explicit_rewoo_step_ids():
    workflow = PlanAndExecute(
        "invalid-ids",
        planner=Agent(
            "planner",
            model=MockProvider(
                responses=[
                    """{
                      "steps": [
                        {"id": "step1", "instruction": "Do it", "depends_on": []}
                      ]
                    }"""
                ]
            ),
        ),
        executor=Agent("executor", model=MockProvider(responses=["ignored"])),
    )

    with pytest.raises(ValueError, match="E1..En"):
        await workflow.run("Bad ids")


@pytest.mark.asyncio
async def test_plan_and_execute_injects_available_tools_into_executor_prompt():
    @tool
    async def lookup_release_notes(topic: str) -> str:
        return f"notes for {topic}"

    planner = Agent(
        "planner",
        model=MockProvider(
            responses=[
                """{
                  "steps": [
                    {"id": "E1", "instruction": "Use the best tool to inspect release notes", "depends_on": []}
                  ]
                }"""
            ]
        ),
    )
    executor_provider = MockProvider(responses=["tool-assisted answer", "final answer"])
    executor = Agent("executor", model=executor_provider).tools([lookup_release_notes])
    workflow = PlanAndExecute("tool-aware", planner=planner, executor=executor)

    await workflow.run("Summarize release notes")

    prompt = executor_provider.call_history[0]["messages"][-1].content
    assert "Available tools for this executor" in prompt
    assert "lookup_release_notes" in prompt


@pytest.mark.asyncio
async def test_plan_and_execute_runs_independent_steps_in_parallel():
    planner = Agent(
        "planner",
        model=MockProvider(
            responses=[
                """{
                  "steps": [
                    {"id": "E1", "instruction": "Fetch source A", "depends_on": []},
                    {"id": "E2", "instruction": "Fetch source B", "depends_on": []},
                    {"id": "E3", "instruction": "Merge #E1 and #E2", "depends_on": ["E1", "E2"]}
                  ]
                }"""
            ]
        ),
    )
    started_at: dict[str, float] = {}

    async def executor(prompt: str) -> str:
        if "Current step (E1)" in prompt:
            started_at["E1"] = time.monotonic()
            await asyncio.sleep(0.06)
            return "source-a"
        if "Current step (E2)" in prompt:
            started_at["E2"] = time.monotonic()
            await asyncio.sleep(0.06)
            return "source-b"
        started_at["E3"] = time.monotonic()
        await asyncio.sleep(0.06)
        assert "source-a" in prompt
        assert "source-b" in prompt
        return "merged"

    workflow = PlanAndExecute(
        "parallel-bot",
        planner=planner,
        executor=executor,
        solver=lambda prompt: "done",
        parallel_execution=True,
        max_parallel_steps=2,
    )

    started = time.monotonic()
    response = await workflow.run("Combine two sources")
    elapsed = time.monotonic() - started

    assert response.content == "done"
    assert abs(started_at["E1"] - started_at["E2"]) < 0.03
    assert started_at["E3"] > started_at["E1"]
    assert elapsed < 0.16
    assert response.raw["plan_and_execute"]["parallel_execution"] is True


@pytest.mark.asyncio
async def test_plan_and_execute_runs_formal_rewoo_dsl_tool_steps():
    tool_calls: list[str] = []

    @tool
    async def lookup_release_notes(topic: str) -> str:
        tool_calls.append(topic)
        return "release notes: patch fixes and rollout timing"

    planner_provider = MockProvider(
        responses=[
            """{
              "steps": [
                {
                  "id": "E1",
                  "instruction": "Fetch the release notes",
                  "depends_on": [],
                  "tool": "lookup_release_notes",
                  "input_template": "{\\"topic\\": \\"{task}\\"}",
                  "output_key": "notes"
                },
                {
                  "id": "E2",
                  "instruction": "Summarize {notes} for the user",
                  "depends_on": ["E1"],
                  "input_template": "Summarize these notes: {notes}",
                  "output_key": "summary"
                }
              ]
            }"""
        ]
    )
    planner = Agent("planner", model=planner_provider)
    executor_provider = MockProvider(responses=["short summary for the user"])
    executor = Agent("executor", model=executor_provider).tools([lookup_release_notes])
    workflow = PlanAndExecute(
        "dsl-bot",
        planner=planner,
        executor=executor,
        solver=lambda prompt: "final answer",
    )

    response = await workflow.run("v1.2 release")

    assert response.content == "final answer"
    assert tool_calls == ["v1.2 release"]
    assert planner_provider.call_history[0]["output_schema"].__name__ == "ReWOOExecutionPlanSchema"
    execution_prompt = executor_provider.call_history[0]["messages"][-1].content
    assert "Summarize these notes: release notes: patch fixes and rollout timing" in execution_prompt
    assert "Named outputs available from prior steps" in execution_prompt
    trace = response.raw["plan_and_execute"]
    assert trace["planner_output_schema"] == "ReWOOExecutionPlanSchema"
    assert trace["executed_steps"][0]["tool"] == "lookup_release_notes"
    assert trace["executed_steps"][0]["resolved_input"] == '{"topic": "v1.2 release"}'
    assert trace["executed_steps"][0]["output_key"] == "notes"
    assert trace["executed_steps"][1]["output_key"] == "summary"


@pytest.mark.asyncio
async def test_plan_and_execute_uses_arg_mapping_for_multi_argument_tools():
    observed_args: list[tuple[str, str]] = []

    @tool
    async def summarize_for_audience(topic: str, audience: str) -> str:
        observed_args.append((topic, audience))
        return f"{topic} for {audience}"

    planner = Agent(
        "planner",
        model=MockProvider(
            responses=[
                """{
                  "steps": [
                    {
                      "id": "E1",
                      "instruction": "Prepare a targeted summary",
                      "depends_on": [],
                      "tool": "summarize_for_audience",
                      "tool_input_schema": {
                        "type": "object",
                        "properties": {
                          "topic": {"type": "string"},
                          "audience": {"type": "string"}
                        },
                        "required": ["topic", "audience"]
                      },
                      "arg_mapping": {
                        "topic": "{task}",
                        "audience": "executive-team"
                      },
                      "output_key": "summary"
                    }
                  ]
                }"""
            ]
        ),
    )
    workflow = PlanAndExecute(
        "strict-tool-args",
        planner=planner,
        executor=Agent("executor", model=MockProvider(responses=["ignored"])).tools([summarize_for_audience]),
        solver=lambda prompt: "done",
    )

    response = await workflow.run("release readiness")

    assert response.content == "done"
    assert observed_args == [("release readiness", "executive-team")]
    executed = response.raw["plan_and_execute"]["executed_steps"][0]
    assert executed["resolved_arguments"] == {
        "topic": "release readiness",
        "audience": "executive-team",
    }
    assert executed["tool_input_schema"]["required"] == ["topic", "audience"]


@pytest.mark.asyncio
async def test_plan_and_execute_raises_for_unknown_rewoo_dsl_tool():
    planner = Agent(
        "planner",
        model=MockProvider(
            responses=[
                """{
                  "steps": [
                    {
                      "id": "E1",
                      "instruction": "Run a missing tool",
                      "depends_on": [],
                      "tool": "missing_tool",
                      "input_template": "payload"
                    }
                  ]
                }"""
            ]
        ),
    )
    workflow = PlanAndExecute(
        "missing-tool",
        planner=planner,
        executor=Agent("executor", model=MockProvider(responses=["ignored"])),
    )

    with pytest.raises(ValueError, match="missing_tool"):
        await workflow.run("Bad tool plan")


@pytest.mark.asyncio
async def test_rewoo_alias_streams_final_answer():
    workflow = ReWOOAgent(
        "rewoo-bot",
        planner=lambda prompt: """{"steps": [{"id": "E1", "instruction": "Answer directly", "depends_on": []}]}""",
        executor=lambda prompt: "draft result",
        solver=lambda prompt: "streamed final answer",
    )

    chunks = [chunk async for chunk in workflow.stream("Do it")]

    assert isinstance(chunks[0], TextChunk)
    assert chunks[0].text == "streamed final answer"
    assert isinstance(chunks[-1], StreamDone)
    assert chunks[-1].response.content == "streamed final answer"
