import os
import logging
from opentelemetry import trace, metrics
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from opentelemetry.sdk.metrics import MeterProvider
from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter
from opentelemetry.sdk.resources import Resource
from opentelemetry.instrumentation.logging import LoggingInstrumentor
from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor
from opentelemetry.instrumentation.requests import RequestsInstrumentor
from opentelemetry.instrumentation.sqlalchemy import SQLAlchemyInstrumentor

def parse_headers(header_str: str) -> dict:
    """Parse comma-separated headers into a dictionary, e.g. Authorization=Basic XYZ"""
    if not header_str:
        return {}
    headers = {}
    for pair in header_str.split(','):
        if '=' in pair:
            key, val = pair.split('=', 1)
            headers[key.strip()] = val.strip()
    return headers

def setup_telemetry():
    """Initialize OpenTelemetry Tracing, Metrics, and Logging"""
    service_name = os.getenv("OTEL_SERVICE_NAME", "vlta-dropshopping-engine")
    endpoint = os.getenv("OTEL_EXPORTER_OTLP_ENDPOINT")
    header_str = os.getenv("OTEL_EXPORTER_OTLP_HEADERS", "")
    
    headers = parse_headers(header_str)
    
    # Common Resource attributes
    resource = Resource.create({"service.name": service_name})
    
    if not endpoint:
        # If no endpoint configured, just use basic stdout logging
        logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
        return

    # 1. Tracing Setup
    trace_exporter = OTLPSpanExporter(
        endpoint=f"{endpoint}/v1/traces",
        headers=headers
    )
    tracer_provider = TracerProvider(resource=resource)
    tracer_provider.add_span_processor(BatchSpanProcessor(trace_exporter))
    trace.set_tracer_provider(tracer_provider)
    
    # 2. Metrics Setup
    metric_exporter = OTLPMetricExporter(
        endpoint=f"{endpoint}/v1/metrics",
        headers=headers
    )
    metric_reader = PeriodicExportingMetricReader(metric_exporter)
    meter_provider = MeterProvider(resource=resource, metric_readers=[metric_reader])
    metrics.set_meter_provider(meter_provider)
    
    # 3. Logging Setup (Injects Trace_ID into logs automatically)
    logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - [trace_id=%(otelTraceID)s span_id=%(otelSpanID)s] - %(message)s")
    LoggingInstrumentor().instrument(set_logging_format=True)
    
    # 4. Global Auto-Instrumentations (Catches outgoing API calls like LLMs & DB Queries)
    HTTPXClientInstrumentor().instrument()
    RequestsInstrumentor().instrument()
    try:
        # Automatically tracks any SQLAlchemy Engines created
        SQLAlchemyInstrumentor().instrument()
    except Exception as e:
        logging.getLogger(__name__).warning(f"Could not instrument SQLAlchemy: {e}")
    
    logging.getLogger(__name__).info("OpenTelemetry initialization complete. Connected to OTLP Endpoint.")