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 .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 _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.
    """
    _print = lambda msg: print(f"[OTEL_RESET] {msg}", flush=True)

    def _reset_sdk_global(mod, provider_attr, guard_attr, label):
        """Reset a single OTel SDK global provider and its set-once guard."""
        old = getattr(mod, provider_attr, 'MISSING')
        has_guard = hasattr(mod, guard_attr)
        old_done = getattr(getattr(mod, guard_attr, None), '_done', 'MISSING')
        _print(f"{label}: {provider_attr}={old}, guard={has_guard}, _done={old_done}")
        setattr(mod, provider_attr, None)
        if has_guard:
            getattr(mod, guard_attr)._done = False
        _print(f"{label} reset OK")

    # 1. Reset module-level cached providers
    try:
        from . import traces, logs, metrics as _metrics_mod
        _print(f"module traces._tracer_provider={traces._tracer_provider}")
        _print(f"module logs._logger_provider={logs._logger_provider}")
        _print(f"module metrics._meter_provider={_metrics_mod._meter_provider}")
        traces._tracer_provider = None
        logs._logger_provider = None
        _metrics_mod._meter_provider = None
        _print("module-level providers reset to None")
    except Exception as e:
        _print(f"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:
        _print(f"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:
        _print(f"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:
        _print(f"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)]
        _print(f"removed {old_count} stale LoggingHandler(s) from root logger")
    except Exception as e:
        _print(f"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:
        print(f"[CELERY_OTEL] Observability disabled/failed. "
              f"OTEL_ENABLED={ObservabilityConfig.OTEL_ENABLED}, "
              f"ENDPOINT={ObservabilityConfig.OTEL_EXPORTER_OTLP_ENDPOINT}", flush=True)
        return False

    # Verify that the new batch threads are actually alive
    _verify_batch_threads()

    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


def _verify_batch_threads():
    """Print diagnostic info about batch processor threads after re-init."""
    _print = lambda msg: print(f"[OTEL_VERIFY] {msg}", flush=True)

    # Check log batch processor
    try:
        from .logs import _logger_provider
        if _logger_provider is None:
            _print("WARN: _logger_provider is None after setup!")
        else:
            mp = getattr(_logger_provider, '_multi_log_record_processor', None)
            if mp:
                processors = getattr(mp, '_log_record_processors', [])
                for p in processors:
                    cls_name = type(p).__name__
                    worker = getattr(p, '_worker_thread', None)
                    alive = worker.is_alive() if worker else 'N/A'
                    _print(f"  LogProcessor: {cls_name}, worker_thread alive={alive}")
            else:
                _print(f"  LoggerProvider processors: {mp}")
    except Exception as e:
        _print(f"  Log check failed: {e}")

    # Check trace batch processor
    try:
        from .traces import _tracer_provider
        if _tracer_provider is None:
            _print("WARN: _tracer_provider is None after setup!")
        else:
            processor = getattr(_tracer_provider, '_active_span_processor', None)
            if processor:
                processors = getattr(processor, '_span_processors', [])
                for p in processors:
                    cls_name = type(p).__name__
                    # DatabaseSpanFilter wraps the real batch processor
                    inner = getattr(p, '_batch_processor', p)
                    worker = getattr(inner, 'worker_thread', None) or getattr(inner, '_worker_thread', None)
                    alive = worker.is_alive() if worker else 'N/A'
                    _print(f"  SpanProcessor: {cls_name}, worker_thread alive={alive}")
    except Exception as e:
        _print(f"  Trace check failed: {e}")

    # Check OTel LoggingHandler on root logger
    try:
        from opentelemetry.sdk._logs import LoggingHandler
        root = logging.getLogger()
        otel_count = sum(1 for h in root.handlers if isinstance(h, LoggingHandler))
        _print(f"  Root logger OTel handlers: {otel_count}, total handlers: {len(root.handlers)}")
    except Exception as e:
        _print(f"  Handler check failed: {e}")
