"""An ordinary external tool loop, called by the pinned Lightning local runner."""

from __future__ import annotations

import asyncio
import json
import os
import re
import urllib.request
from collections.abc import Awaitable, Callable


TOOLS = [{"type": "function", "function": {
    "name": "multiply", "description": "Multiply two integers.",
    "parameters": {"type": "object", "properties": {
        "a": {"type": "integer"}, "b": {"type": "integer"}},
        "required": ["a", "b"], "additionalProperties": False},
}}]


async def external_harness(question: str, policy: Callable[[list[dict]], Awaitable[dict]]) -> str:
    """Keep tools and orchestration outside the weight-training framework."""
    messages = [{"role": "user", "content": (
        question + " Use the multiply tool. Finish with exactly <answer>integer</answer>."
    )}]
    for _ in range(3):
        message = await policy(messages)
        messages.append(message)
        calls = message.get("tool_calls") or []
        if not calls:
            return message.get("content") or ""
        if len(calls) > 4:
            return ""
        for call in calls:
            try:
                function = call["function"]
                arguments = json.loads(function["arguments"])
                if function["name"] != "multiply" or not isinstance(arguments, dict) or set(arguments) != {"a", "b"}:
                    raise ValueError("Expected multiply(a, b)")
                if any(type(v) is not int or abs(v) > 100_000 for v in arguments.values()):
                    raise ValueError("Arguments must be bounded integers")
                observation = str(arguments["a"] * arguments["b"])
            except (ValueError, TypeError, KeyError):
                observation = "Tool error: use multiply with two bounded integer arguments a and b."
            messages.append({"role": "tool", "tool_call_id": call["id"], "content": observation})
    return ""  # A bounded task failure, not an infrastructure failure.


def _post(url: str, key: str, payload: dict) -> dict:
    request = urllib.request.Request(
        url, data=json.dumps(payload).encode(), method="POST",
        headers={"Authorization": f"Bearer {key}", "Content-Type": "application/json"},
    )
    with urllib.request.urlopen(request, timeout=45) as response:
        return json.loads(response.read())


class Agent:
    """`async run()` and these environment variables match upstream local mode."""

    async def run(self) -> None:
        key = os.environ["AGL_KEY"]

        async def policy(messages: list[dict]) -> dict:
            response = await asyncio.to_thread(
                _post, os.environ["AGL_OPENAI_BASE_URL"].rstrip("/") + "/chat/completions", key,
                {"model": "auto", "messages": messages, "tools": TOOLS,
                 "tool_choice": "auto", "temperature": 1.0, "max_tokens": 128},
            )
            return response["choices"][0]["message"]

        answer = await asyncio.wait_for(external_harness(os.environ["QUESTION"], policy), timeout=150)
        match = re.fullmatch(r"\s*<answer>\s*(-?\d+)\s*</answer>\s*", answer)
        reward = float(match is not None and match.group(1) == os.environ["RESULT"].strip())
        await asyncio.to_thread(_post, os.environ["AGL_EVENT_URL"], key,
                                {"event_type": "reward", "data": {"value": reward}})