"""Scripted stand-ins for an LLM, so the MCP examples run with no API key.

`--offline` swaps the model for one of these. Each makes the tool calls a
model given the examples' task would make, reading the session id and the
tickers from the tool results as a model would: open a session, observe,
buy, advance to the close, close. The tools, the MCP connection and the
server are all real; only the choosing is scripted.
"""

from __future__ import annotations

import json
import uuid
from typing import Any

RUN = uuid.uuid4().hex[:8]

# (tool, arguments from what the earlier results said)
PLAN = [
    ("open_session", lambda s: {"universe_size": 5, "seed": 7, "label": "scripted",
                                "idempotency_key": f"{RUN}-open"}),
    ("observe", lambda s: {"session_id": s["session_id"]}),
    ("place_order", lambda s: {"session_id": s["session_id"], "ticker": s["tickers"][0],
                               "side": "buy", "quantity": 10,
                               "idempotency_key": f"{RUN}-buy"}),
    ("advance", lambda s: {"session_id": s["session_id"], "until": "close",
                           "idempotency_key": f"{RUN}-advance"}),
    ("close_session", lambda s: {"session_id": s["session_id"],
                                 "idempotency_key": f"{RUN}-close"}),
]


def _as_dict(content: Any) -> dict[str, Any]:
    if isinstance(content, dict):
        return content
    if isinstance(content, list) and content:
        content = content[0]
        content = getattr(content, "text", content)
        if isinstance(content, dict):
            content = content.get("text", content)
    try:
        value = json.loads(content)
    except (TypeError, ValueError):
        return {}
    if isinstance(value, dict) and set(value) == {"type", "text"}:
        return _as_dict(value["text"])
    return value if isinstance(value, dict) else {}


def next_call(results: list[dict[str, Any]]) -> tuple[str, dict[str, Any]] | str:
    """The next tool call after `results` (one parsed result per call made so
    far), or the final answer once the plan is done."""
    state: dict[str, Any] = {}
    for r in results:
        if "tickers" in r and "session_id" in r:
            state.update(session_id=r["session_id"], tickers=r["tickers"])
    if len(results) < len(PLAN):
        if results and "code" in results[-1] and "message" in results[-1]:
            return f"Stopped: {results[-1]['code']}: {results[-1]['message']}"
        name, args = PLAN[len(results)]
        return name, args(state)
    report = results[-1]
    worth = report.get("account", {}).get("net_worth")
    caveats = report.get("caveats", [])
    return (f"Closed session {state.get('session_id')}: net worth {worth:,.2f} after "
            f"{report.get('days')} day(s). Caveats ({len(caveats)}): "
            + " | ".join(caveats[:2]))


# -- the OpenAI Agents SDK ------------------------------------------------------------


def openai_agents_model() -> Any:
    from agents.items import ModelResponse
    from agents.models.interface import Model
    from agents.usage import Usage
    from openai.types.responses import ResponseFunctionToolCall, ResponseOutputMessage, ResponseOutputText

    class ScriptedModel(Model):
        async def get_response(self, system_instructions, input, model_settings, tools,
                               output_schema, handoffs, tracing, *, previous_response_id=None,
                               conversation_id=None, prompt=None):
            items = input if isinstance(input, list) else []
            results = [_as_dict(i.get("output")) for i in items
                       if isinstance(i, dict) and i.get("type") == "function_call_output"]
            step = next_call(results)
            if isinstance(step, str):
                out = ResponseOutputMessage(
                    id=f"msg_{len(results)}", type="message", role="assistant", status="completed",
                    content=[ResponseOutputText(type="output_text", text=step, annotations=[])])
            else:
                name, args = step
                out = ResponseFunctionToolCall(type="function_call", id=f"fc_{len(results)}",
                                               call_id=f"call_{RUN}_{len(results)}", name=name,
                                               arguments=json.dumps(args), status="completed")
            return ModelResponse(output=[out], usage=Usage(), response_id=None)

        def stream_response(self, *args: Any, **kwargs: Any):
            raise NotImplementedError("the scripted model does not stream")

    return ScriptedModel()


# -- PydanticAI --------------------------------------------------------------------------


def pydantic_ai_model() -> Any:
    from pydantic_ai.messages import ModelRequest, ModelResponse, TextPart, ToolCallPart, ToolReturnPart
    from pydantic_ai.models.function import FunctionModel

    def choose(messages: list[Any], info: Any) -> Any:
        results = [_as_dict(part.content) for m in messages if isinstance(m, ModelRequest)
                   for part in m.parts if isinstance(part, ToolReturnPart)]
        step = next_call(results)
        if isinstance(step, str):
            return ModelResponse(parts=[TextPart(step)])
        name, args = step
        return ModelResponse(parts=[ToolCallPart(name, args)])

    return FunctionModel(choose)