# Requirements:
#   pip install logfire openai

import os
import logfire
from openai import OpenAI

# Configure logfire first (before setting up our own tracer)
logfire.configure(send_to_logfire=False)

from opentelemetry import trace
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter

# Add our exporter to capture spans
exporter = InMemorySpanExporter()
provider = trace.get_tracer_provider()
if hasattr(provider, 'add_span_processor'):
    provider.add_span_processor(SimpleSpanProcessor(exporter))

client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))

# Instrument OpenAI with logfire
logfire.instrument_openai(client)

def test_streaming():
    """Test OpenAI Responses API streaming and check for token usage."""
    exporter.clear()

    # Make a streaming request using Responses API
    with client.responses.stream(
        model="gpt-4o-mini",
        input="Write a haiku about programming."
    ) as stream:
        stream.until_done()

    # Check the captured spans for token usage
    spans = exporter.get_finished_spans()
    for span in spans:
        attrs = dict(span.attributes) if span.attributes else {}
        input_tokens = attrs.get("gen_ai.usage.input_tokens")
        output_tokens = attrs.get("gen_ai.usage.output_tokens")
        if "streaming" in span.name.lower():
            print(f"Span: {span.name}")
            print(f"  gen_ai.usage.input_tokens: {input_tokens}")
            print(f"  gen_ai.usage.output_tokens: {output_tokens}")
            return input_tokens, output_tokens

    return None, None

print("="*60)
print("Testing WITHOUT patch")
print("="*60)

input_tokens, output_tokens = test_streaming()

if input_tokens is None and output_tokens is None:
    print("\nBUG: Token usage not captured for streaming!")
else:
    print(f"\nToken usage captured: input={input_tokens}, output={output_tokens}")

# Now apply the patch
print("\n" + "="*60)
print("Applying patch: adding token usage extraction")
print("="*60)

from logfire._internal.integrations.llm_providers.openai import (
    OpenaiResponsesStreamState,
    responses_output_events,
)

def patched_get_attributes(self, span_data: dict) -> dict:
    response = self.get_response_data()
    span_data["events"] = span_data["events"] + responses_output_events(response)

    # FIX: Extract token usage from the response
    usage = getattr(response, "usage", None)
    input_tokens = getattr(usage, "input_tokens", None)
    output_tokens = getattr(usage, "output_tokens", None)
    if isinstance(input_tokens, int):
        span_data["gen_ai.usage.input_tokens"] = input_tokens
    if isinstance(output_tokens, int):
        span_data["gen_ai.usage.output_tokens"] = output_tokens

    return span_data

OpenaiResponsesStreamState.get_attributes = patched_get_attributes

print("\n" + "="*60)
print("Testing WITH patch")
print("="*60)

input_tokens, output_tokens = test_streaming()

if input_tokens is None and output_tokens is None:
    print("\nToken usage still not captured!")
else:
    print(f"\nFIXED: Token usage captured: input={input_tokens}, output={output_tokens}")

class OpenaiResponsesStreamState:
    def get_attributes(self, span_data: dict[str, Any]) -> dict[str, Any]:
        response = self.get_response_data()
        span_data['events'] = span_data['events'] + responses_output_events(response)
        return span_data  # Missing: token usage extraction

def get_attributes(self, span_data: dict[str, Any]) -> dict[str, Any]:
    response = self.get_response_data()
    span_data['events'] = span_data['events'] + responses_output_events(response)

    # Extract token usage from the response
    usage = getattr(response, "usage", None)
    input_tokens = getattr(usage, "input_tokens", None)
    output_tokens = getattr(usage, "output_tokens", None)
    if isinstance(input_tokens, int):
        span_data["gen_ai.usage.input_tokens"] = input_tokens
    if isinstance(output_tokens, int):
        span_data["gen_ai.usage.output_tokens"] = output_tokens

    return span_data