# aiva/core/tracing.py
from opentelemetry import trace
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from opentelemetry.exporter.jaeger.thrift import JaegerExporter
from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
from fastapi import FastAPI
import os

def configure_tracing(app: FastAPI):
    """Configures tracing for the FastAPI application."""

    service_name = os.getenv("TRACING_SERVICE_NAME", "aiva-service")
    jaeger_host = os.getenv("JAEGER_HOST", "localhost")
    jaeger_port = int(os.getenv("JAEGER_PORT", 6831))
    aiva_version = os.getenv("AIVA_VERSION", "unknown")

    resource = Resource.create({
        "service.name": service_name,
        "service.version": aiva_version,
    })

    tracer_provider = TracerProvider(resource=resource)
    jaeger_exporter = JaegerExporter(
        collector_endpoint=f"http://{jaeger_host}:14268/api/traces", # Updated for Jaeger HTTP endpoint
        service_name=service_name,
    )

    tracer_provider.add_span_processor(BatchSpanProcessor(jaeger_exporter))
    trace.set_tracer_provider(tracer_provider)

    FastAPIInstrumentor.instrument_app(app, tracer_provider=tracer_provider)

    # Add AIVA configuration parameters as span attributes (example)
    @app.middleware("http")
    async def add_aiva_config_to_trace(request, call_next):
        with trace.get_tracer(__name__).start_as_current_span("middleware") as span:
            span.set_attribute("aiva.version", aiva_version)
            span.set_attribute("aiva.config.param1", "value1") # Example
            span.set_attribute("aiva.config.param2", "value2") # Example

            response = await call_next(request)
            return response