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, ensure_all_logs_captured
    except ImportError as e:
        logger.warning(f"OpenTelemetry not installed: {e}. Observability disabled.")
        return False
    
    traces_ok = True
    if enable_traces:
        traces_ok = setup_tracing(filter_database_spans=ObservabilityConfig.OTEL_FILTER_DATABASE_SPANS)

    if enable_metrics:
        if not setup_metrics():
            logger.warning("Metrics setup failed — traces and logs still active")

    if enable_logs:
        if setup_logging():
            attach_log_handlers()
            ensure_all_logs_captured()

    if instrument_framework:
        _instrument_framework()

    _initialized = True
    return traces_ok

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_web_observability(
    service_name: str = "integration-service",
) -> bool:
    """
    Re-initialize observability for ASGI web servers (FastAPI/uvicorn).

    Call this from the ASGI app startup event — NOT at module import time.
    When uvicorn uses --reload it spawns a subprocess via SpawnProcess and
    imports the app inside asyncio.run(), before the event loop is fully
    running.  The BatchLogRecordProcessor threads and gRPC channel created at
    that point are unreliable.  Calling this function from the startup event
    guarantees a fresh provider, exporter, and export thread after the event
    loop is healthy — mirroring the pattern used by setup_celery_observability()
    in worker_process_init.
    """
    global _initialized

    # Force re-initialization so setup_observability() doesn't short-circuit.
    _initialized = False

    # Clear the existing logger provider so setup_logging() creates a new one
    # with a fresh BatchLogRecordProcessor (new export thread + gRPC channel).
    try:
        from . import logs as _obs_logs
        _obs_logs._logger_provider = None
    except Exception:
        pass

    # Clear tracer provider as well so traces also get fresh exporters.
    try:
        from . import traces as _obs_traces
        _obs_traces._tracer_provider = None
    except Exception:
        pass

    return setup_observability(
        service_name=service_name,
        enable_traces=True,
        enable_metrics=True,
        enable_logs=True,
        instrument_framework=False,  # Already instrumented at import time
    )


def force_flush(timeout_millis: int = 5000):
    """
    Force-flush all pending spans and logs to SigNoz.

    Required for AWS Lambda — the container freezes immediately after the
    handler returns, killing BatchSpanProcessor/BatchLogRecordProcessor
    background threads before they can export. Call this before returning
    from the Lambda handler to ensure all telemetry reaches SigNoz.
    """
    try:
        from .traces import _tracer_provider
        if _tracer_provider:
            _tracer_provider.force_flush(timeout_millis=timeout_millis)
    except Exception as e:
        logger.debug("force_flush traces failed: %s", e)

    try:
        from .logs import _logger_provider
        if _logger_provider:
            _logger_provider.force_flush(timeout_millis=timeout_millis)
    except Exception as e:
        logger.debug("force_flush logs failed: %s", 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

