"""End-to-end: an agent grounds its answer in an ingested document with citations."""

from __future__ import annotations

import json

POLICY = (
    b"Our refund policy lets customers request a refund within 30 days of purchase. "
    b"Approved refunds are returned to the original payment method within five business days."
)


def _call(name, args):
    return json.dumps({"thought": "look it up", "action": name, "args": args})


def _final(text):
    return json.dumps({"thought": "answer", "action": "final", "args": {"output": text}})


def test_agent_answers_from_ingested_document_with_citations(client, helpers, set_agent_responses):
    owner = helpers.bootstrap_owner(client)
    token = owner["access_token"]

    kb = helpers.create_kb(client, token, name="Support KB")
    doc = helpers.upload_document(client, token, kb["id"], filename="refunds.txt", content=POLICY)
    assert doc["status"] == "ready"

    agent = helpers.create_agent(client, token, name="Support", tools=["kb.search"])
    set_agent_responses(
        [
            _call("kb.search", {"query": "refund timeframe"}),
            _final("Refunds are returned within five business days."),
        ]
    )

    created = helpers.run_agent(client, token, agent["id"], input="How long do refunds take?")
    run_id = created["run"]["id"]

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

    steps = client.get(f"/runs/{run_id}/steps", headers=helpers.headers(token)).json()
    tool_steps = [s for s in steps if s["type"] == "tool_call" and s["name"] == "kb.search"]
    assert tool_steps, "agent should have called kb.search"

    results = tool_steps[0]["output"]["data"]["results"]
    assert results, "kb.search should have retrieved grounded passages"
    assert results[0]["document_id"] == doc["id"]
    assert "refund" in results[0]["text"].lower()
