"""Running an agent over the API: execution, RBAC, tenant scoping, and WS streaming."""

from __future__ import annotations

import json

import pytest
from legate.policy.roles import Role
from starlette.websockets import WebSocketDisconnect


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 test_run_executes_and_persists_steps(client, helpers, set_agent_responses):
    owner = helpers.bootstrap_owner(client)
    token = owner["access_token"]
    agent = helpers.create_agent(client, token, tools=["kb.search"])
    set_agent_responses([_call("kb.search", {"query": "x"}), _final("the answer")])

    created = helpers.run_agent(client, token, agent["id"], input="help me")
    run_id = created["run"]["id"]
    assert created["stream_url"] == f"/runs/{run_id}/stream"

    run = client.get(f"/runs/{run_id}", headers=helpers.headers(token)).json()
    assert run["status"] == "succeeded"
    assert run["output"]["output"] == "the answer"
    assert run["total_steps"] == 3

    steps = client.get(f"/runs/{run_id}/steps", headers=helpers.headers(token)).json()
    assert [s["type"] for s in steps] == ["agent_reason", "tool_call", "agent_reason"]


def test_run_appears_in_list(client, helpers, set_agent_responses):
    owner = helpers.bootstrap_owner(client)
    token = owner["access_token"]
    agent = helpers.create_agent(client, token)
    set_agent_responses([_final("done")])
    helpers.run_agent(client, token, agent["id"])

    runs = client.get("/runs", headers=helpers.headers(token)).json()
    assert len(runs) == 1 and runs[0]["status"] == "succeeded"


def test_operator_can_run_viewer_cannot(client, helpers, set_agent_responses):
    owner = helpers.bootstrap_owner(client)
    agent = helpers.create_agent(client, owner["access_token"])
    operator = helpers.add_member(
        client, owner["access_token"], email="op@x.com", role=Role.OPERATOR
    )
    viewer = helpers.add_member(client, owner["access_token"], email="vw@x.com", role=Role.VIEWER)

    set_agent_responses([_final("ok"), _final("ok")])
    ok = client.post(
        f"/agents/{agent['id']}/run",
        json={"input": ""},
        headers=helpers.headers(operator["access_token"]),
    )
    assert ok.status_code == 201

    denied = client.post(
        f"/agents/{agent['id']}/run",
        json={"input": ""},
        headers=helpers.headers(viewer["access_token"]),
    )
    assert denied.status_code == 403


def test_runs_are_tenant_scoped(client, helpers, set_agent_responses):
    a = helpers.bootstrap_owner(client, email="a@x.com", name="A")
    b = helpers.bootstrap_owner(client, email="b@x.com", name="B")
    agent = helpers.create_agent(client, a["access_token"])
    set_agent_responses([_final("secret")])
    created = helpers.run_agent(client, a["access_token"], agent["id"])
    run_id = created["run"]["id"]

    assert (
        client.get(f"/runs/{run_id}", headers=helpers.headers(b["access_token"])).status_code == 404
    )


def test_websocket_streams_steps_then_terminal_status(client, helpers, set_agent_responses):
    owner = helpers.bootstrap_owner(client)
    token = owner["access_token"]
    agent = helpers.create_agent(client, token, tools=["kb.search"])
    set_agent_responses([_call("kb.search", {"query": "x"}), _final("answer")])
    created = helpers.run_agent(client, token, agent["id"])
    run_id = created["run"]["id"]

    events = []
    with client.websocket_connect(f"/runs/{run_id}/stream?token={token}") as ws:
        while True:
            msg = ws.receive_json()
            events.append(msg)
            if msg["type"] == "run.status" and msg["status"] == "succeeded":
                break

    step_events = [e for e in events if e["type"] == "step"]
    assert len(step_events) == 3
    assert events[-1]["status"] == "succeeded"


def test_websocket_rejects_missing_token(client, helpers, set_agent_responses):
    owner = helpers.bootstrap_owner(client)
    token = owner["access_token"]
    agent = helpers.create_agent(client, token)
    set_agent_responses([_final("done")])
    run_id = helpers.run_agent(client, token, agent["id"])["run"]["id"]

    with (
        pytest.raises(WebSocketDisconnect),
        client.websocket_connect(f"/runs/{run_id}/stream") as ws,
    ):
        ws.receive_json()
