"""The plan-act-observe loop: tool calls, budgets, retries, and safe failure."""

from __future__ import annotations

import json
import uuid

from legate.agent.enums import RunStatus, StepType
from legate.agent.events import run_event_bus
from legate.agent.runner import AgentRunner
from legate.agent.tooling import ToolRegistry, ToolResult, ToolSpec
from legate.db.models.agent import Agent
from legate.db.models.run import Run
from legate.db.models.workspace import Workspace
from legate.llm.fake import FakeLLMClient

# -- builders --------------------------------------------------------------


def make_ws(db) -> Workspace:
    ws = Workspace(name="W", slug=f"w-{uuid.uuid4().hex[:8]}", settings_json={})
    db.add(ws)
    db.flush()
    return ws


def make_agent(db, ws, *, tools=None, guardrails=None) -> Agent:
    agent = Agent(
        workspace_id=ws.id,
        name="A",
        description="",
        system_prompt="You help.",
        provider="fake",
        model="fake-model",
        tools=tools or [],
        memory_config={},
        guardrails_json=guardrails or {},
    )
    db.add(agent)
    db.flush()
    return agent


def make_run(db, ws, agent, text="do it") -> Run:
    run = Run(
        workspace_id=ws.id, agent_id=agent.id, input={"input": text}, status=RunStatus.PENDING
    )
    db.add(run)
    db.flush()
    return run


def final(text: str) -> str:
    return json.dumps({"thought": "done", "action": "final", "args": {"output": text}})


def call(name: str, args: dict) -> str:
    return json.dumps({"thought": "act", "action": name, "args": args})


def steps_of(db, run):
    from legate.repositories import steps as steps_repo

    return steps_repo.list_for_run(db, workspace_id=run.workspace_id, run_id=run.id)


async def run_loop(db, agent, run, responses, **kwargs):
    runner = AgentRunner(db=db, llm=FakeLLMClient(responses), **kwargs)
    return await runner.execute(run, agent)


# -- tests -----------------------------------------------------------------


async def test_final_answer_succeeds(db):
    ws = make_ws(db)
    agent = make_agent(db, ws)
    run = make_run(db, ws, agent)

    await run_loop(db, agent, run, [final("hello")])

    assert run.status is RunStatus.SUCCEEDED
    assert run.output == {"output": "hello", "thought": "done"}
    assert run.total_steps == 1
    assert run.started_at is not None and run.ended_at is not None


async def test_tool_call_then_observation_then_final(db):
    ws = make_ws(db)
    agent = make_agent(db, ws, tools=["kb.search"])
    run = make_run(db, ws, agent)

    fake = FakeLLMClient([call("kb.search", {"query": "x"}), final("answer")])
    runner = AgentRunner(db=db, llm=fake)
    await runner.execute(run, agent)

    assert run.status is RunStatus.SUCCEEDED
    types = [s.type for s in steps_of(db, run)]
    assert types == [StepType.AGENT_REASON, StepType.TOOL_CALL, StepType.AGENT_REASON]

    # The observation is fed back as untrusted data on the next turn.
    second_turn = fake.calls[1]
    assert any("OBSERVATION (untrusted tool output" in m.content for m in second_turn)


async def test_unknown_tool_makes_agent_replan(db):
    ws = make_ws(db)
    agent = make_agent(db, ws, tools=["kb.search"])
    run = make_run(db, ws, agent)

    fake = FakeLLMClient([call("does.not.exist", {}), final("recovered")])
    runner = AgentRunner(db=db, llm=fake)
    await runner.execute(run, agent)

    assert run.status is RunStatus.SUCCEEDED
    assert any("Unknown tool 'does.not.exist'" in m.content for m in fake.calls[1])


async def test_tool_outside_allowlist_is_denied(db):
    ws = make_ws(db)
    # kb.search exists in the registry but is NOT granted to this agent.
    agent = make_agent(db, ws, tools=[])
    run = make_run(db, ws, agent)

    fake = FakeLLMClient([call("kb.search", {"query": "x"}), final("ok")])
    runner = AgentRunner(db=db, llm=fake)
    await runner.execute(run, agent)

    denied = [s for s in steps_of(db, run) if s.status == "denied"]
    assert denied and denied[0].name == "kb.search"
    assert any("denied by policy" in m.content for m in fake.calls[1])


async def test_sensitive_tool_suspends_for_approval(db):
    ws = make_ws(db)
    agent = make_agent(db, ws, tools=["danger.do"])
    run = make_run(db, ws, agent)

    called = {"hit": False}

    async def handler(args, ctx):
        called["hit"] = True
        return ToolResult(ok=True)

    registry = ToolRegistry()
    registry.register(
        ToolSpec(
            "danger.do", "does danger", {"type": "object", "properties": {}}, is_sensitive=True
        ),
        handler,
    )

    runner = AgentRunner(db=db, llm=FakeLLMClient([call("danger.do", {})]), registry=registry)
    await runner.execute(run, agent)

    assert run.status is RunStatus.WAITING_APPROVAL
    assert called["hit"] is False  # the sensitive action never executed
    assert any(s.type is StepType.APPROVAL for s in steps_of(db, run))


async def test_max_steps_budget_enforced(db):
    ws = make_ws(db)
    agent = make_agent(db, ws, tools=["kb.search"], guardrails={"max_steps": 2})
    run = make_run(db, ws, agent)

    await run_loop(
        db, agent, run, [call("kb.search", {"query": "a"}), call("kb.search", {"query": "b"})]
    )

    assert run.status is RunStatus.FAILED
    assert "max steps" in run.error


async def test_token_budget_enforced(db):
    ws = make_ws(db)
    agent = make_agent(db, ws, guardrails={"max_tokens": 1})
    run = make_run(db, ws, agent)

    await run_loop(db, agent, run, [final("hello")])

    assert run.status is RunStatus.FAILED
    assert "token budget" in run.error


async def test_wallclock_budget_enforced(db):
    ws = make_ws(db)
    agent = make_agent(db, ws, guardrails={"max_wallclock_sec": 0})
    run = make_run(db, ws, agent)

    fake = FakeLLMClient([final("hello")])
    await AgentRunner(db=db, llm=fake).execute(run, agent)

    assert run.status is RunStatus.FAILED
    assert "wall-clock" in run.error
    assert fake.calls == []  # failed before any model call


async def test_malformed_output_retries_then_fails_safe(db):
    ws = make_ws(db)
    agent = make_agent(db, ws)
    run = make_run(db, ws, agent)

    # MAX_PARSE_RETRIES = 2 -> three attempts, all malformed.
    await run_loop(db, agent, run, ["nope", "still bad", "{not json"])

    assert run.status is RunStatus.FAILED
    assert "malformed" in run.error
    reason_steps = [s for s in steps_of(db, run) if s.type is StepType.AGENT_REASON]
    assert len(reason_steps) == 3
    assert all(s.status == "error" for s in reason_steps)


async def test_malformed_output_then_recovers(db):
    ws = make_ws(db)
    agent = make_agent(db, ws)
    run = make_run(db, ws, agent)

    await run_loop(db, agent, run, ["garbage", final("recovered")])

    assert run.status is RunStatus.SUCCEEDED
    assert run.output["output"] == "recovered"


async def test_runner_publishes_live_events(db):
    ws = make_ws(db)
    agent = make_agent(db, ws, tools=["kb.search"])
    run = make_run(db, ws, agent)

    queue = run_event_bus.subscribe(run.id)
    try:
        await run_loop(db, agent, run, [call("kb.search", {"query": "x"}), final("done")])
    finally:
        run_event_bus.unsubscribe(run.id, queue)

    events = []
    while not queue.empty():
        events.append(queue.get_nowait())

    assert any(e["type"] == "step" for e in events)
    assert events[-1] == {
        "type": "run.status",
        "status": "succeeded",
        "output": {"output": "done", "thought": "done"},
    }
