"""
Pattern C Example — Manual Span Control with Real LLM and Dynamic Dispatch.

Demonstrates:
  1. Registering manual tool schema and calling OpenAI via Infralo Gateway.
  2. Dynamically dispatching tool calls based on the LLM's response.
  3. Starting spans manually via `trace.start_tool_span(...)` and ending them
     with success, timeout, or error statuses and error codes.
"""

import json
import logging
import os
import time

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_c")

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


def execute_database_query(query: str, query_timeout: float = 2.0) -> list:
    """Mock database query execution that can timeout or fail."""
    print(f"Executing Query: '{query}' (Timeout: {query_timeout}s)")
    if "SELECT" in query:
        return [{"id": 1, "value": "A"}, {"id": 2, "value": "B"}]
    elif "SLEEP" in query:
        time.sleep(2.5)
        raise TimeoutError("Database query timed out after exceeding threshold")
    else:
        raise ValueError("Syntax error: Near 'CRASH' (line 1)")


tools_schema = [
    {
        "type": "function",
        "function": {
            "name": "db_query",
            "description": "Execute a raw SQL query on the database.",
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "SQL statement to run, e.g. SELECT * FROM items",
                    }
                },
                "required": ["query"],
            },
        },
    }
]


def run_pattern_c():
    print("=== Running Pattern C: Dynamic Dispatch with Manual Spans ===")

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

        # Call the LLM with a tuned prompt that requests three different SQL statements
        prompt = (
            "Please call the db_query tool three times to execute the following SQL statements:\n"
            "1. 'SELECT id, value FROM items'\n"
            "2. 'SELECT SLEEP(10)'\n"
            "3. 'CRASH DATABASE'"
        )
        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()

        # Iterate and execute the tool calls
        tool_calls = getattr(response.choices[0].message, "tool_calls", None) or []
        for tool_call in tool_calls:
            arguments = json.loads(tool_call.function.arguments)
            query = arguments["query"]
            call_id = tool_call.id

            # Start tool span manually (Pattern C)
            span = trace.start_tool_span(name="db_query", tool_call_id=call_id)
            span.set_input({"query": query})

            try:
                # Run query
                results = execute_database_query(query)

                # End with successful status
                span.end(output=results, status="success")
                print(f"Query '{query}' succeeded and manual span recorded.")

            except TimeoutError as te:
                # End with explicit timeout status
                span.end(status="timeout", error=str(te), error_code="DB_TIMEOUT")
                print(f"Query '{query}' timeout handled manually: {te}")

            except Exception as e:
                # End with explicit error status
                span.end(status="error", error=str(e), error_code="SQL_SYNTAX_ERROR")
                print(f"Query '{query}' exception handled manually: {e}")


if __name__ == "__main__":
    run_pattern_c()
    infralo.shutdown()