#!/usr/bin/env python3
"""Test script to verify distributed tracing between webapp and API."""

import sys
import time
from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
import requests

# Setup OpenTelemetry (use Docker network hostname)
otlp_exporter = OTLPSpanExporter(
    endpoint="http://jaeger:4318/v1/traces",
    timeout=2,
)
provider = TracerProvider()
processor = BatchSpanProcessor(otlp_exporter)
provider.add_span_processor(processor)
trace.set_tracer_provider(provider)

tracer = trace.get_tracer("test-distributed-tracing")

# Create a parent span and make HTTP request to API
print("Creating parent span and calling API...")
with tracer.start_as_current_span("test-webapp-request") as parent_span:
    parent_span.set_attribute("test.type", "distributed_tracing")
    
    # Inject trace context into headers
    headers = {
        "Authorization": "Bearer dev-token-change-in-production",
        "X-User-ID": "test-user",
        "Content-Type": "application/json",
    }
    TraceContextTextMapPropagator().inject(headers)
    
    print(f"Injected headers: {headers}")
    print(f"Parent span trace_id: {format(parent_span.get_span_context().trace_id, '032x')}")
    print(f"Parent span span_id: {format(parent_span.get_span_context().span_id, '016x')}")
    
    # Make API request (use Docker network hostname)
    with tracer.start_as_current_span("http.request"):
        try:
            response = requests.get(
                "http://api:4004/sessions",
                headers=headers,
                timeout=5,
            )
            print(f"\nAPI response status: {response.status_code}")
            print(f"Sessions returned: {len(response.json())}")
        except Exception as e:
            print(f"Error calling API: {e}", file=sys.stderr)
            sys.exit(1)

print("\nWaiting 2 seconds for spans to flush to Jaeger...")
time.sleep(2)

trace_id = format(parent_span.get_span_context().trace_id, '032x')
print(f"\n✅ Test complete!")
print(f"View trace in Jaeger UI:")
print(f"  http://localhost:16686/trace/{trace_id}")
print(f"\nExpected structure:")
print(f"  test-webapp-request (this script)")
print(f"    └─ http.request (HTTP client)")
print(f"       └─ GET /sessions (API middleware) <-- Should be child of http.request")
print(f"          └─ SessionManager.ListSessions")
print(f"             └─ PostgresStorage.ListSessions")