import logging
from typing import Optional, Dict, Any
from .config import ObservabilityConfig

logger = logging.getLogger(__name__)

_initialized = False


def diagnose() -> Dict[str, Any]:
    """
    Run diagnostics on the OTel setup and print results to stdout.
    Returns a dict with all diagnostic info.

    Usage:
        from varicon_observability import diagnose
        diagnose()
    """
    results = {}
    _print = lambda k, v: print(f"[OTEL_DIAG] {k}: {v}", flush=True)

    # 1. Package version
    try:
        from . import __version__
        results["package_version"] = __version__
    except Exception as e:
        results["package_version"] = f"ERROR: {e}"
    _print("package_version", results["package_version"])

    # 2. Config
    results["OTEL_ENABLED"] = ObservabilityConfig.OTEL_ENABLED
    results["OTEL_EXPORTER_OTLP_ENDPOINT"] = ObservabilityConfig.OTEL_EXPORTER_OTLP_ENDPOINT
    results["OTEL_EXPORTER_OTLP_PROTOCOL"] = ObservabilityConfig.OTEL_EXPORTER_OTLP_PROTOCOL
    results["OTEL_SERVICE_NAME"] = ObservabilityConfig.OTEL_SERVICE_NAME
    results["_initialized"] = _initialized
    for k in ["OTEL_ENABLED", "OTEL_EXPORTER_OTLP_ENDPOINT", "OTEL_EXPORTER_OTLP_PROTOCOL",
              "OTEL_SERVICE_NAME", "_initialized"]:
        _print(k, results[k])

    # 3. TracerProvider
    try:
        from opentelemetry import trace
        tp = trace.get_tracer_provider()
        tp_type = type(tp).__name__
        is_proxy = isinstance(tp, trace.ProxyTracerProvider)
        results["tracer_provider_type"] = tp_type
        results["tracer_provider_is_proxy"] = is_proxy
        _print("tracer_provider_type", tp_type)
        _print("tracer_provider_is_proxy (should be False)", is_proxy)

        # Test creating a real span
        tracer = trace.get_tracer("varicon_observability.diag")
        with tracer.start_as_current_span("diag_test") as span:
            sc = span.get_span_context()
            results["test_span_type"] = type(span).__name__
            results["test_span_valid"] = sc.is_valid
            results["test_trace_id"] = format(sc.trace_id, "032x")
            _print("test_span_type", type(span).__name__)
            _print("test_span_valid (should be True)", sc.is_valid)
            _print("test_trace_id (should be non-zero)", results["test_trace_id"])
    except Exception as e:
        results["tracer_provider_error"] = str(e)
        _print("tracer_provider_error", e)

    # 4. LoggerProvider / OTel LoggingHandler
    try:
        root = logging.getLogger()
        handler_types = [type(h).__name__ for h in root.handlers]
        results["root_handlers"] = handler_types
        _print("root_logger_handlers", handler_types)

        from opentelemetry.sdk._logs import LoggingHandler
        has_otel = any(isinstance(h, LoggingHandler) for h in root.handlers)
        results["has_otel_logging_handler"] = has_otel
        _print("has_otel_logging_handler (should be True)", has_otel)
    except Exception as e:
        results["logging_handler_error"] = str(e)
        _print("logging_handler_error", e)

    # 5. Module-level providers
    try:
        from .traces import _tracer_provider
        results["traces._tracer_provider"] = type(_tracer_provider).__name__ if _tracer_provider else "None"
        _print("traces._tracer_provider", results["traces._tracer_provider"])
    except Exception as e:
        _print("traces._tracer_provider", f"ERROR: {e}")

    try:
        from .logs import _logger_provider
        results["logs._logger_provider"] = type(_logger_provider).__name__ if _logger_provider else "None"
        _print("logs._logger_provider", results["logs._logger_provider"])
    except Exception as e:
        _print("logs._logger_provider", f"ERROR: {e}")

    return results

def setup_observability(
    service_name: Optional[str] = None,
    enable_traces: bool = True,
    enable_metrics: bool = True,
    enable_logs: bool = True,
    instrument_framework: bool = True,
) -> bool:
    global _initialized
    
    if _initialized:
        return True
    
    if service_name:
        ObservabilityConfig.OTEL_SERVICE_NAME = service_name
    
    if not ObservabilityConfig.OTEL_ENABLED:
        logger.info("OpenTelemetry disabled")
        return False
    
    if not ObservabilityConfig.OTEL_EXPORTER_OTLP_ENDPOINT:
        logger.warning("OTEL_EXPORTER_OTLP_ENDPOINT not set")
        return False
    
    # Try importing OpenTelemetry modules
    try:
        from .traces import setup_tracing
        from .metrics import setup_metrics
        from .logs import setup_logging, attach_log_handlers
        from .capture import ensure_all_logs_captured
    except ImportError as e:
        logger.warning(f"OpenTelemetry not installed: {e}. Observability disabled.")
        return False
    
    success = True
    
    if enable_traces:
        success &= setup_tracing(filter_database_spans=ObservabilityConfig.OTEL_FILTER_DATABASE_SPANS)
    
    if enable_metrics:
        success &= setup_metrics()
    
    if enable_logs:
        if setup_logging():
            attach_log_handlers()
            ensure_all_logs_captured()
    
    if instrument_framework:
        _instrument_framework()
    
    _initialized = True
    return success

def _instrument_framework():
    try:
        from .integrations import setup_django, setup_fastapi, setup_celery
        
        if setup_django():
            return
        
        if setup_fastapi():
            return
        
        setup_celery()
    except ImportError as e:
        logger.warning(f"Failed to instrument framework: {e}")

def setup_celery_observability(
    service_name: str = "integration-service-celery",
    logger_class=None,
) -> bool:
    """
    Complete observability setup for Celery workers.
    
    This function combines:
    1. General observability setup (traces, metrics, logs)
    2. Celery-specific instrumentation for trace context propagation
    
    Args:
        service_name: Name of the service for OpenTelemetry
        logger_class: Optional logger class with set_tenant_id method
        
    Returns:
        True if setup was successful, False otherwise
        
    Example:
        from integration_management.varicon_observability import setup_celery_observability
        from integration_management.common.logger import IntegrationLogger
        
        setup_celery_observability(logger_class=IntegrationLogger)
    """
    global _initialized
    
    # Force re-initialization for Celery workers.
    # Celery prefork model forks from the main process which already ran
    # setup_observability() (setting _initialized=True). But in the forked
    # child process, the OTel batch processor threads are dead (killed by fork),
    # so we MUST re-initialize to create new providers, exporters, and threads.
    _initialized = False
    
    # First, set up general observability
    observability_enabled = setup_observability(
        service_name=service_name,
        enable_traces=True,
        enable_metrics=True,
        enable_logs=True,
        instrument_framework=True,
    )
    
    if not observability_enabled:
        logger.warning("Varicon Observability disabled or failed to initialize for Celery worker")
        return False
    
    logger.info("Varicon Observability initialized for Celery worker")
    
    # Then, set up Celery-specific instrumentation
    try:
        from .celery_instrumentation import setup_celery_instrumentation
        
        setup_celery_instrumentation(logger_class=logger_class)
        logger.info("Celery trace context propagation configured")
        return True
    except Exception as e:
        logger.warning(f"Failed to setup Celery instrumentation: {e}")
        return False

