"""Brizz + Segment + OpenAI example.

Brizz traces the OpenAI call, Segment sends a track event.
"""

import os
import sys
import uuid

import segment.analytics as analytics
from dotenv import load_dotenv
from openai import OpenAI
from openai.types.chat import ChatCompletionUserMessageParam

from brizz import Brizz, start_session

load_dotenv()

# Check for API key
if not os.getenv("OPENAI_API_KEY"):
    print("Set OPENAI_API_KEY environment variable")
    sys.exit(1)

# Initialize Brizz (tracing) and Segment (events)
SERVICE_NAME = "segment-example"
Brizz.initialize(disable_batch=True, app_name=SERVICE_NAME)
analytics.write_key = os.getenv("SEGMENT_WRITE_KEY", "")

client = OpenAI()
session_id = uuid.uuid4().hex

with start_session(session_id):
    # OpenAI call is auto-instrumented by Brizz
    response = client.chat.completions.create(
        model="gpt-5.1",
        messages=[ChatCompletionUserMessageParam(role="user", content="What is 2+2?")],
    )
    print(f"Agent: {response.choices[0].message.content}")

    # Track the interaction via Segment
    analytics.track(
        user_id="test-user-001",
        event="Agent Response",
        properties={
            "model": "gpt-5.1",
            "brizzSessionId": session_id,
            "brizzServiceName": SERVICE_NAME,
            "brizzEnvironment": "development",
        },
    )

analytics.flush()
