"""
Pattern B Example — @infralo.tool Decorator with Real OpenAI Call and Dynamic Dispatch.

Demonstrates:
  1. Wrapping sync and async tools using the `@infralo.tool` decorator.
  2. Registering tools schema and calling OpenAI LLM via Infralo Gateway.
  3. Dynamically dispatching tool calls based on the LLM's response.
  4. Automatic extraction of trace contexts using contextvars during dispatch.
"""

import asyncio
import json
import logging
import os

from openai import OpenAI

from infralo import Infralo

# Configure standard logging
logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger("infralo.pattern_b")

# Load environment configuration
INFRALO_API_KEY = os.environ.get("INFRALO_API_KEY", "your_workspace_api_key")
INFRALO_ENDPOINT = os.environ.get("INFRALO_ENDPOINT", "http://localhost:8000")
MODEL_DEPLOYMENT = os.environ.get("INFRALO_MODEL_DEPLOYMENT", "your_model_name")

infralo = Infralo(
    api_key=INFRALO_API_KEY,
    endpoint=INFRALO_ENDPOINT,
    flush_interval=1.0,
)

openai_client = OpenAI(
    base_url=f"{INFRALO_ENDPOINT.rstrip('/')}/v1",
    api_key=INFRALO_API_KEY,
)


# ── DECORATED TOOLS ──────────────────────────────────────────────────────────


# 1. Sync Tool: captures inputs/outputs automatically
@infralo.tool
def add_numbers(a: int, b: int) -> int:
    print(f"Executing add_numbers({a}, {b})")
    return a + b


# 2. Async Tool: supports async execution, and overrides name
@infralo.tool(name="fetch_user_profile")
async def fetch_user(user_id: str) -> dict:
    print(f"Fetching profile for {user_id}...")
    await asyncio.sleep(0.1)
    return {"id": user_id, "name": "Jane Doe", "role": "admin"}


# 3. Tool with output payload capture disabled
@infralo.tool(capture_output=False)
def generate_secret_token(username: str) -> str:
    print(f"Generating token for {username}")
    return f"secret-token-for-{username}"


# Map tool names to the decorated functions for dynamic dispatching
tools_map = {
    "add_numbers": add_numbers,
    "fetch_user_profile": fetch_user,
    "generate_secret_token": generate_secret_token,
}

# OpenAI tools schema representing our functions
tools_schema = [
    {
        "type": "function",
        "function": {
            "name": "add_numbers",
            "description": "Add two integers together.",
            "parameters": {
                "type": "object",
                "properties": {
                    "a": {"type": "integer"},
                    "b": {"type": "integer"},
                },
                "required": ["a", "b"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "fetch_user_profile",
            "description": "Fetch user profile details.",
            "parameters": {
                "type": "object",
                "properties": {
                    "user_id": {"type": "string"},
                },
                "required": ["user_id"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "generate_secret_token",
            "description": "Generate a secret session token for a user.",
            "parameters": {
                "type": "object",
                "properties": {
                    "username": {"type": "string"},
                },
                "required": ["username"],
            },
        },
    },
]


async def main():
    print("=== Running Pattern B: Dynamic Dispatch with Decorators ===")

    with infralo.start_trace(session_id="session_decor_456") as trace:
        print(f"Started trace: {trace.trace_id}")

        # Call the LLM with a tuned prompt to trigger all three tools
        prompt = (
            "Please call the tools to add 20 and 30, fetch user profile for user-007, "
            "and generate a secret token for janesmith."
        )
        try:
            raw_response = openai_client.chat.completions.with_raw_response.create(
                model=MODEL_DEPLOYMENT,
                messages=[{"role": "user", "content": prompt}],
                tools=tools_schema,
                extra_headers=trace.headers(),
            )
        except Exception as exc:
            logger.error("Failed to connect to LLM gateway. Details: %s", exc)
            return

        # Read LLM span context
        trace.update_from_response(raw_response)
        response = raw_response.parse()

        # Process the tool calls returned by the LLM
        tool_calls = getattr(response.choices[0].message, "tool_calls", None) or []
        for tool_call in tool_calls:
            name = tool_call.function.name
            arguments = json.loads(tool_call.function.arguments)
            call_id = tool_call.id

            if name in tools_map:
                tool_func = tools_map[name]

                # Pass special infralo metadata and call identity using custom kwargs prefix.
                # These prefix parameters are stripped by the decorator before the real function is called.
                infralo_kwargs = {
                    "_infralo_tool_call_id": call_id,
                    "_infralo_metadata": {"invocation_mode": "agent_loop"},
                }

                # Execute dynamically
                if asyncio.iscoroutinefunction(tool_func):
                    result = await tool_func(**arguments, **infralo_kwargs)
                else:
                    result = tool_func(**arguments, **infralo_kwargs)
                print(f"Executed tool '{name}' successfully. Result: {result}")
            else:
                logger.warning("LLM requested unknown tool: %s", name)


if __name__ == "__main__":
    asyncio.run(main())
    infralo.shutdown()