"""
Pattern A Example — Explicit Context Manager with Real OpenAI Call.

Demonstrates:
  1. Starting a trace with session_id.
  2. Calling the LLM through the Infralo Gateway using real OpenAI client.
  3. Reading back the LLM span identity via `trace.update_from_response(response)`.
  4. Executing multiple tool calls returned by the LLM and tracing them using `with trace.tool(...)`.
  5. Automatic error capturing on exceptions.
  6. Subsequent LLM call propagating the trace parent automatically.
"""

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

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

# Initialize the Infralo SDK client.
infralo = Infralo(
    api_key=INFRALO_API_KEY,
    endpoint=INFRALO_ENDPOINT,
    flush_interval=1.0,
)

# Initialize the OpenAI client pointing to the Infralo API Gateway
openai_client = OpenAI(
    base_url=f"{INFRALO_ENDPOINT.rstrip('/')}/v1",
    api_key=INFRALO_API_KEY,
)


def execute_weather_tool(location: str) -> dict:
    """Mock execution of a local tool function."""
    if "Chicago" in location:
        return {"temperature": 72, "unit": "fahrenheit", "condition": "sunny"}
    elif "San Francisco" in location:
        raise ValueError("Weather service offline for San Francisco")
    return {"temperature": 60}


def run_pattern_a():
    print("=== Running Pattern A: Real LLM Gateway Call ===")

    # 1. Start the agentic chain trace
    with infralo.start_trace(session_id="user_session_abc") as trace:
        print(f"Started trace: {trace.trace_id}")

        # Declare the tools schema for the LLM
        tools_schema = [
            {
                "type": "function",
                "function": {
                    "name": "get_weather",
                    "description": "Get current weather for a given location",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "location": {
                                "type": "string",
                                "description": "City and state, e.g. San Francisco, CA",
                            }
                        },
                        "required": ["location"],
                    },
                },
            }
        ]

        try:
            # 2. First LLM Gateway call using trace propagation headers.
            # Use .with_raw_response to ensure we capture the HTTP headers.
            raw_response = openai_client.chat.completions.with_raw_response.create(
                model=MODEL_DEPLOYMENT,
                messages=[
                    {
                        "role": "user",
                        "content": "What is the weather like in Chicago and San Francisco?",
                    }
                ],
                tools=tools_schema,
                extra_headers=trace.headers(),  # ← Propagates trace_id and parent
            )
        except Exception as exc:
            logger.error(
                "Failed to connect to LLM gateway. (Ensure your Infralo backend and OpenAI API keys are active). Details: %s",
                exc,
            )
            return

        # 3. Read back the LLM span identity so tool calls are linked as children
        llm_span_id = trace.update_from_response(raw_response)
        logger.info("Gateway returned LLM span ID: %s", llm_span_id)

        # Parse the actual completion response model
        response = raw_response.parse()

        # 4. Iterate and execute the tool calls
        tool_calls = getattr(response.choices[0].message, "tool_calls", None) or []
        tool_messages = []

        for tool_call in tool_calls:
            name = tool_call.function.name
            arguments = json.loads(tool_call.function.arguments)
            call_id = tool_call.id

            # Wrap each tool call inside trace.tool() context manager
            try:
                with trace.tool(name=name, tool_call_id=call_id) as span:
                    span.set_input(arguments)
                    span.set_metadata(service_version="1.4.0")

                    # Run tool logic
                    result = execute_weather_tool(arguments["location"])

                    # Set output preview + payload
                    span.set_output(result)
                    print(f"Tool {call_id} succeeded: {result}")
                    tool_messages.append(
                        {
                            "role": "tool",
                            "tool_call_id": call_id,
                            "content": json.dumps(result),
                        }
                    )
            except Exception as e:
                # Note: trace.tool context manager automatically caught and recorded
                # this exception as status="error" and set error_message=str(e).
                print(f"Tool {call_id} failed as expected: {e}")
                tool_messages.append(
                    {
                        "role": "tool",
                        "tool_call_id": call_id,
                        "content": f"Error: {e}",
                    }
                )

        # 5. Subsequent LLM call. trace.headers() automatically propagates
        # the trace_id and trace.last_tool_span_id (the parent link).
        if tool_calls:
            messages = [
                {
                    "role": "user",
                    "content": "What is the weather like in Chicago and San Francisco?",
                },
                response.choices[0].message,
                *tool_messages,
            ]
            print("\nSending tool outputs back to LLM...")
            try:
                response2 = openai_client.chat.completions.create(
                    model=MODEL_DEPLOYMENT,
                    messages=messages,
                    extra_headers=trace.headers(),  # ← Propagates trace_id and last_tool_span_id
                )
                print(f"Final LLM Response: {response2.choices[0].message.content}")
            except Exception as exc:
                logger.error("Failed to connect to LLM gateway for second call: %s", exc)


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