"""End-to-end demo: three scripted models over one toolset.

These are not the test suite -- tests/ is.  They are the readable proof that the
engine does something: a model that answers, a model that misuses its tools and
recovers, and a model that never stops and is cut off by the budget.

    python3 demo.py            # runs all three, writes demo-*.jsonl, prints a summary
"""

from __future__ import annotations

import json
import os
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, ".."))

from agentloop import AgentLoop, ModelReply, StopReason, Tool, ToolCall, ToolRegistry  # noqa: E402
from agentloop.transcript import verify_file  # noqa: E402


def make_registry():
    def add(a, b):
        return {"sum": a + b}

    def greet(name, excited=False):
        return "hello %s%s" % (name, "!" if excited else ".")

    def fail():
        raise RuntimeError("the tool is broken on purpose")

    def echo(text):
        return text

    return ToolRegistry(
        [
            Tool("add", "Add two integers.", {
                "type": "object",
                "properties": {
                    "a": {"type": "integer"},
                    "b": {"type": "integer"},
                },
                "required": ["a", "b"],
                "additionalProperties": False,
            }, add),
            Tool("greet", "Greet someone by name.", {
                "type": "object",
                "properties": {
                    "name": {"type": "string", "minLength": 1},
                    "excited": {"type": "boolean"},
                },
                "required": ["name"],
            }, greet),
            Tool("fail", "Always raises, to show a tool error in the transcript.", {
                "type": "object", "properties": {}, "additionalProperties": False,
            }, fail),
            Tool("echo", "Return the text unchanged.", {
                "type": "object",
                "properties": {"text": {"type": "string"}},
                "required": ["text"],
            }, echo),
        ]
    )


# -- the three models -----------------------------------------------------


def answering_model(messages, manifest):
    """Adds 2 and 3, then answers without asking for anything else."""
    if not any(m.get("name") == "add" for m in messages):
        return ModelReply("Let me add those.", [ToolCall("add", {"a": 2, "b": 3})])
    return ModelReply("2 + 3 = 5.")


def recovering_model(messages, manifest):
    """First turn: a tool that does not exist, then a bad argument, then a
    crashing tool, then a valid call.  Every failure is reported to the model as
    a tool message and the loop continues."""
    n = sum(1 for m in messages if m["role"] == "tool")
    if n == 0:
        return ModelReply("First try a name I made up.", [ToolCall("multiply", {"a": 1})])
    if n == 1:
        return ModelReply("Now with the wrong argument type.", [ToolCall("add", {"a": "2", "b": 3})])
    if n == 2:
        return ModelReply("Now one that raises.", [ToolCall("fail", {})])
    if n == 3:
        return ModelReply("Now correctly.", [ToolCall("add", {"a": 20, "b": 22})])
    return ModelReply("20 + 22 = 42. DONE", [])


def runaway_model(messages, manifest):
    """Never stops on its own.  Only the budget stops it."""
    i = sum(1 for m in messages if m["role"] == "assistant") + 1
    return ModelReply("still going (step %d) " % i * 12, [ToolCall("echo", {"text": "tick"})])


# -- runner ---------------------------------------------------------------


def main():
    out_dir = HERE
    runs = []

    loop = AgentLoop(make_registry(), max_steps=6)
    result = loop.run("What is 2 + 3?", answering_model)
    path = os.path.join(out_dir, "demo-answer.jsonl")
    result.transcript.save(path)
    runs.append(("answering", result, path))

    loop = AgentLoop(make_registry(), max_steps=6, stop_sentinel="DONE")
    result = loop.run("Add 20 and 22, tolerating my mistakes.", recovering_model)
    path = os.path.join(out_dir, "demo-recover.jsonl")
    result.transcript.save(path)
    runs.append(("recovering", result, path))

    loop = AgentLoop(make_registry(), max_steps=4, max_tokens=200)
    result = loop.run("Keep going forever.", runaway_model)
    path = os.path.join(out_dir, "demo-budget.jsonl")
    result.transcript.save(path)
    runs.append(("runaway", result, path))

    print("demo runs")
    for name, result, path in runs:
        ok, violations, count = verify_file(path)
        print(
            "  %-11s stop=%-13s steps=%d tokens=%-4d records=%d chain=%s"
            % (name, result.stop_reason, result.steps, result.tokens, count, "OK" if ok else violations)
        )

    errors = sum(1 for r in runs[1][1].tool_results() if not r.ok)
    print("  recovering  tool failures surfaced to the model: %d" % errors)
    print("  runaway     stopped by budget, not by the model: %s"
          % (runs[2][1].stop_reason in (StopReason.STEP_BUDGET, StopReason.TOKEN_BUDGET)))
    print("  transcripts written next to demo.py")

    # Exit 0 only if every claim above actually held.
    if not all(verify_file(p)[0] for _, _, p in runs):
        print("FAIL: a transcript failed chain verification")
        return 1
    if runs[0][1].stop_reason != StopReason.NO_TOOL_CALLS:
        print("FAIL: the answering model did not stop on no_tool_calls")
        return 1
    if runs[1][1].stop_reason != StopReason.SENTINEL:
        print("FAIL: the recovering model did not stop on the sentinel")
        return 1
    if errors != 3:
        print("FAIL: expected 3 surfaced tool failures, saw %d" % errors)
        return 1
    if runs[2][1].stop_reason not in (StopReason.STEP_BUDGET, StopReason.TOKEN_BUDGET):
        print("FAIL: the runaway model was not cut off by a budget")
        return 1
    print("ALL DEMO CLAIMS HELD")
    return 0


if __name__ == "__main__":
    sys.exit(main())