"""Minimal stdio Executa plugin scaffold (Python).

Speaks Executa JSON-RPC 2.0 over stdin/stdout. Implements both v1
methods (`describe`, `health`, `invoke`) and the v2 `initialize`
handshake so a host that knows sampling can attach reverse-RPC support.

Run locally:

    anna-app executa dev --describe
    anna-app executa dev --invoke ping
"""

from __future__ import annotations

import json
import sys
import uuid
from typing import Any

MANIFEST: dict[str, Any] = {
    # The manifest no longer declares a top-level `name`; the plugin's
    # identity is the server-assigned tool_id (resolved from the on-disk
    # shim name), not a self-declared manifest name.
    "display_name": "__TOOL_ID__",
    "version": "0.1.0",
    "description": "A minimal Executa plugin scaffold.",
    # Declaring `host_capabilities` here is informational for the
    # registry; `client_capabilities.sampling` in the initialize
    # response is what actually unlocks reverse-RPC sampling.
    "host_capabilities": ["llm.sample"],
    "tools": [
        {
            "name": "ping",
            "description": "Smoke-test method.",
            "parameters": [],
        },
        {
            "name": "summarize",
            "description": "Summarize a piece of text via host sampling.",
            "parameters": [
                {
                    "name": "text",
                    "type": "string",
                    "description": "Text to summarize.",
                    "required": True,
                }
            ],
        },
    ],
}

# Track in-flight reverse-RPC requests we initiate to the host.
_PENDING: dict[str, Any] = {}


def _write(envelope: dict[str, Any]) -> None:
    sys.stdout.write(json.dumps(envelope) + "\n")
    sys.stdout.flush()


def _request_sampling(prompt: str) -> dict[str, Any]:
    """Issue a reverse-RPC `sampling/createMessage` and block until response."""
    rid = str(uuid.uuid4())
    _write(
        {
            "jsonrpc": "2.0",
            "id": rid,
            "method": "sampling/createMessage",
            "params": {
                "messages": [
                    {"role": "user", "content": {"type": "text", "text": prompt}}
                ],
                "maxTokens": 400,
            },
        }
    )
    # Loop on stdin until our response arrives. Anything that's not the
    # expected response is dispatched normally.
    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue
        env = json.loads(line)
        if env.get("id") == rid and "method" not in env:
            if "error" in env:
                raise RuntimeError(env["error"].get("message", "sampling failed"))
            return env.get("result", {})
        # Re-queue other requests by handling them inline.
        _dispatch(env)
    raise RuntimeError("stdin closed before sampling response arrived")


def invoke(method: str, args: dict[str, Any]) -> dict[str, Any]:
    if method == "ping":
        return {"success": True, "data": {"pong": True}}
    if method == "summarize":
        text = args.get("text", "")
        if not text:
            return {"success": False, "error": "text is required"}
        try:
            sample = _request_sampling(f"Summarize in one sentence:\n{text}")
            summary = sample.get("content", {}).get("text", "")
            return {"success": True, "data": {"summary": summary}}
        except Exception as e:  # noqa: BLE001
            return {"success": False, "error": f"sampling failed: {e}"}
    return {"success": False, "error": f"unknown method: {method}"}


def _dispatch(env: dict[str, Any]) -> None:
    method = env.get("method")
    rid = env.get("id")
    try:
        if method == "initialize":
            result = {
                "protocolVersion": "2.0",
                "server_info": {"name": MANIFEST["display_name"], "version": MANIFEST["version"]},
                "capabilities": {"sampling": {}},
            }
        elif method == "describe":
            result = MANIFEST
        elif method == "health":
            result = {"status": "ready"}
        elif method == "invoke":
            params = env.get("params", {})
            result = invoke(params.get("tool", ""), params.get("arguments", {}))
        else:
            raise ValueError(f"unknown rpc: {method}")
        _write({"jsonrpc": "2.0", "id": rid, "result": result})
    except Exception as e:  # noqa: BLE001
        _write(
            {
                "jsonrpc": "2.0",
                "id": rid,
                "error": {"code": -32601, "message": str(e)},
            }
        )


def main() -> None:
    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue
        _dispatch(json.loads(line))


if __name__ == "__main__":
    main()