"""
Parallel Execution Example with Real OpenAI Call and Dynamic Dispatch.

Demonstrates:
  1. Registering multiple tool schemas and calling OpenAI via Infralo Gateway.
  2. Dynamically dispatching tool calls based on the LLM's response.
  3. Concurrently executing async tool calls in a `with trace.parallel()` block.
  4. Verifying sequence number behavior during concurrent execution.
"""

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.parallel")

# 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,
)


# ── TOOLS DEFINITION ─────────────────────────────────────────────────────────


@infralo.tool(name="fetch_product_details")
async def get_product_details(product_id: str) -> dict:
    print(f"Executing fetch_product_details for product '{product_id}'...")
    await asyncio.sleep(0.15)
    return {"product_id": product_id, "price": 49.99, "stock": 120}


@infralo.tool(name="fetch_user_reviews")
async def get_user_reviews(product_id: str) -> list:
    print(f"Executing fetch_user_reviews for product '{product_id}'...")
    await asyncio.sleep(0.1)
    return [
        {"user": "Alice", "rating": 5, "comment": "Great product!"},
        {"user": "Bob", "rating": 4, "comment": "Worth the money."},
    ]


@infralo.tool(name="log_event")
def log_analytics_event(product_id: str, action: str):
    print(f"Executing log_event: '{action}' on '{product_id}'")
    return {"status": "event_logged"}


tools_map = {
    "fetch_product_details": get_product_details,
    "fetch_user_reviews": get_user_reviews,
    "log_event": log_analytics_event,
}

tools_schema = [
    {
        "type": "function",
        "function": {
            "name": "fetch_product_details",
            "description": "Fetch current price and stock details for a product.",
            "parameters": {
                "type": "object",
                "properties": {
                    "product_id": {"type": "string"},
                },
                "required": ["product_id"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "fetch_user_reviews",
            "description": "Fetch user ratings and reviews for a product.",
            "parameters": {
                "type": "object",
                "properties": {
                    "product_id": {"type": "string"},
                },
                "required": ["product_id"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "log_event",
            "description": "Log an analytics action for a product.",
            "parameters": {
                "type": "object",
                "properties": {
                    "product_id": {"type": "string"},
                    "action": {"type": "string"},
                },
                "required": ["product_id", "action"],
            },
        },
    },
]


async def run_parallel_example():
    print("=== Running Parallel Execution Example with Real LLM ===")

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

        # Prompt tuned to trigger parallel operations: fetching product details + reviews
        prompt = "Please call the tools to fetch product details and user reviews for product prod_abc123."
        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 from header
        trace.update_from_response(raw_response)
        response = raw_response.parse()

        tool_calls = getattr(response.choices[0].message, "tool_calls", None) or []
        async_coros = []

        print("\nStarting parallel API fetch block (Shared Sequence position)...")

        # Gather parallel async coroutines to execute inside the parallel context
        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]
                infralo_kwargs = {
                    "_infralo_tool_call_id": call_id,
                    "_infralo_metadata": {"invocation_mode": "async_parallel"},
                }

                # Construct coroutine (each will be executed dynamically)
                if asyncio.iscoroutinefunction(tool_func):
                    async_coros.append(tool_func(**arguments, **infralo_kwargs))
                else:
                    # For demonstration, log warning if sync tool returned for parallel block
                    logger.warning("Expected async tool for parallel block, got sync: %s", name)

        # Run concurrent tools within the parallel sequence context
        if async_coros:
            async with trace.parallel():
                # Both tools are awaited concurrently and share sequence_number = 1
                results = await asyncio.gather(*async_coros)
                print(f"Parallel fetch complete. Results gathered: {results}")

        print("\nStarting subsequent sequential step (Sequence should increment to 2)...")

        # This sequential step runs after the parallel block and gets sequence_number = 2
        log_analytics_event("prod_abc123", "view_details")


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