import os
import base64
from typing import Optional, Literal
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry import trace
from openinference.instrumentation.smolagents import SmolagentsInstrumentor


def setup_tracing(
    service_name: str = "agent-service",
    enabled: bool = True,
    provider: Literal["langfuse", "phoenix", "both"] = "langfuse",
) -> Optional[TracerProvider]:
    """
    Configure and set up OpenTelemetry tracing with Langfuse and/or Phoenix integration.

    Args:
        service_name: Name of the service for trace identification
        enabled: Whether tracing should be active
        provider: Which tracing provider(s) to use - "langfuse", "phoenix", or "both"

    Returns:
        Configured TracerProvider instance or None if disabled
    """
    if not enabled:
        return None

    trace_provider = None

    # Setup Phoenix if requested
    if provider in ["phoenix", "both"]:
        try:
            from phoenix.otel import register

            # Create Phoenix tracer provider

            trace_provider = register(
                project_name=service_name,
                endpoint="http://127.0.0.1:6006/v1/traces",
                auto_instrument=True,
            )
            print(f"✅ Phoenix tracing enabled for {service_name}")

        except ImportError:
            print("⚠️ Phoenix not installed. Run: pip install 'smolagents[telemetry]'")

    # Setup Langfuse if requested
    if provider in ["langfuse", "both"]:
        # Configure Langfuse host
        langfuse_host = os.environ.get("LANGFUSE_HOST", "https://cloud.langfuse.com")

        # Create auth token for Langfuse
        langfuse_public_key = os.environ.get("LANGFUSE_PUBLIC_KEY")
        langfuse_secret_key = os.environ.get("LANGFUSE_SECRET_KEY")

        if not langfuse_public_key or not langfuse_secret_key:
            print("⚠️ Langfuse keys not found in environment")

            # If Phoenix is not set up either, return None
            if trace_provider is None:
                return None
        else:
            langfuse_auth = base64.b64encode(
                f"{langfuse_public_key}:{langfuse_secret_key}".encode()
            ).decode()

            # Set environment variables for OTLP exporter
            os.environ["OTEL_EXPORTER_OTLP_ENDPOINT"] = (
                f"{langfuse_host}/api/public/otel"
            )
            os.environ["OTEL_EXPORTER_OTLP_HEADERS"] = (
                f"Authorization=Basic {langfuse_auth}"
            )

            if trace_provider is None:
                # Create new TracerProvider if Phoenix wasn't enabled
                trace_provider = TracerProvider()
                # Set as the global default tracer provider
                trace.set_tracer_provider(trace_provider)

            # Add Langfuse exporter to the tracer provider
            trace_provider.add_span_processor(SimpleSpanProcessor(OTLPSpanExporter()))
            print(f"✅ Langfuse tracing enabled for {service_name}")

    if trace_provider:
        # Instrument smolagents with the configured provider
        SmolagentsInstrumentor().instrument(tracer_provider=trace_provider)

        # Create a single tracer for the entire application
        trace.get_tracer(service_name)

    return trace_provider