import logging
from typing import Optional
from .config import ObservabilityConfig

logger = logging.getLogger(__name__)

_initialized = False


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 .logs 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}")


_dead_providers = []

def _reset_otel_providers():
    """
    Reset OTel global providers so they are recreated with fresh batch threads.

    After fork(), the old providers' BatchSpanProcessor / BatchLogRecordProcessor
    background threads are dead. The OTel SDK uses a "set once" guard that
    prevents replacing global providers via set_tracer_provider / set_logger_provider.
    We must reset those guards and clear the cached providers.
    """
    def _reset_sdk_global(mod, provider_attr, guard_attr, label):
        """Reset a single OTel SDK global provider and its set-once guard."""
        logger.debug(f"[OTEL_RESET] {label}: {provider_attr}={getattr(mod, provider_attr, 'MISSING')}, guard={hasattr(mod, guard_attr)}")
        old_provider = getattr(mod, provider_attr, None)
        if old_provider is not None:
            _dead_providers.append(old_provider)
        setattr(mod, provider_attr, None)
        if hasattr(mod, guard_attr):
            getattr(mod, guard_attr)._done = False
        logger.debug(f"[OTEL_RESET] {label} reset OK")

    # 1. Reset module-level cached providers
    try:
        from . import traces, logs, metrics as _metrics_mod
        logger.debug(f"[OTEL_RESET] traces._tracer_provider={traces._tracer_provider}")
        logger.debug(f"[OTEL_RESET] logs._logger_provider={logs._logger_provider}")
        logger.debug(f"[OTEL_RESET] metrics._meter_provider={_metrics_mod._meter_provider}")
        if getattr(traces, '_tracer_provider', None) is not None: _dead_providers.append(traces._tracer_provider)
        if getattr(logs, '_logger_provider', None) is not None: _dead_providers.append(logs._logger_provider)
        if getattr(_metrics_mod, '_meter_provider', None) is not None: _dead_providers.append(_metrics_mod._meter_provider)
        traces._tracer_provider = None
        logs._logger_provider = None
        _metrics_mod._meter_provider = None
        logger.debug("[OTEL_RESET] module-level providers reset to None")
    except Exception as e:
        logger.debug(f"[OTEL_RESET] step 1 FAILED: {e}")

    # 2. Reset the global TracerProvider (+ its "set once" guard)
    #    Lives in opentelemetry.trace (opentelemetry/trace/__init__.py)
    try:
        import opentelemetry.trace as _trace_mod
        _reset_sdk_global(_trace_mod, '_TRACER_PROVIDER', '_TRACER_PROVIDER_SET_ONCE', 'trace')
    except Exception as e:
        logger.debug(f"[OTEL_RESET] step 2 FAILED: {e}")

    # 3. Reset the global LoggerProvider (+ its "set once" guard)
    #    Lives in opentelemetry._logs._internal (NOT opentelemetry._logs)
    try:
        import opentelemetry._logs._internal as _logs_internal
        _reset_sdk_global(_logs_internal, '_LOGGER_PROVIDER', '_LOGGER_PROVIDER_SET_ONCE', 'logs')
    except Exception as e:
        logger.debug(f"[OTEL_RESET] step 3 FAILED: {e}")

    # 4. Reset the global MeterProvider (+ its "set once" guard)
    try:
        import opentelemetry.metrics._internal as _metrics_internal
        _reset_sdk_global(_metrics_internal, '_METER_PROVIDER', '_METER_PROVIDER_SET_ONCE', 'metrics')
    except Exception as e:
        logger.debug(f"[OTEL_RESET] step 4 (metrics): {e}")

    # 5. Remove stale OTel LoggingHandlers from the root logger
    #    (they reference the old provider whose batch threads are dead)
    try:
        from opentelemetry.sdk._logs import LoggingHandler
        root = logging.getLogger()
        old_count = len([h for h in root.handlers if isinstance(h, LoggingHandler)])
        root.handlers = [h for h in root.handlers if not isinstance(h, LoggingHandler)]
        logger.debug(f"[OTEL_RESET] removed {old_count} stale LoggingHandler(s) from root logger")
    except Exception as e:
        logger.debug(f"[OTEL_RESET] step 5 FAILED: {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

    # Reset OTel global providers so setup_tracing / setup_logging
    # create fresh providers with live batch threads.
    _reset_otel_providers()

    # 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(
            f"Observability disabled/failed. "
            f"OTEL_ENABLED={ObservabilityConfig.OTEL_ENABLED}, "
            f"ENDPOINT={ObservabilityConfig.OTEL_EXPORTER_OTLP_ENDPOINT}"
        )
        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


