# src/sibi_flux/logger/_logger.py
from __future__ import annotations

import atexit
import logging
import re
import sys
import threading
import time
from collections.abc import Mapping
from contextlib import contextmanager, nullcontext, suppress
from logging import LoggerAdapter
from logging.handlers import QueueHandler, QueueListener, RotatingFileHandler
from pathlib import Path
from queue import Queue
from typing import (
    TYPE_CHECKING,
    Any,
)

if TYPE_CHECKING:
    from collections.abc import Iterable

# --- OpenTelemetry Imports (Defensive) ---
try:
    from opentelemetry import trace
    from opentelemetry import _logs as otel_logs
    from opentelemetry.exporter.otlp.proto.grpc.log_exporter import OTLPLogExporter
    from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
    from opentelemetry.sdk._logs import LoggerProvider, LoggingHandler
    from opentelemetry.sdk._logs.export import BatchLogRecordProcessor
    from opentelemetry.sdk.resources import Resource
    from opentelemetry.sdk.trace import TracerProvider
    from opentelemetry.sdk.trace.export import BatchSpanProcessor

    set_logger_provider = otel_logs.set_logger_provider
    get_logger_provider = otel_logs.get_logger_provider
    _OTEL_AVAILABLE = True
except (ImportError, Exception):
    trace = None
    set_logger_provider = None
    get_logger_provider = None
    OTLPLogExporter = None
    OTLPSpanExporter = None
    LoggerProvider = None
    LoggingHandler = None
    BatchLogRecordProcessor = None
    Resource = None
    TracerProvider = None
    BatchSpanProcessor = None
    _OTEL_AVAILABLE = False


class PIISecretFilter(logging.Filter):
    """
    Redacts PII and Secrets from log records with optimized pre-checks.
    """

    _SECRET_PATTERNS = [
        re.compile(
            r"(?i)(password|passwd|secret|token|api_key|access_key|auth_token|bearer)\s*(=|:|is)\s*([^\s]+)",
        ),
        re.compile(r"(?i)(Authorization)(\s*:\s*)(Bearer\s+[^\s]+)"),
    ]

    _SENSITIVE_KEYS = {
        "password",
        "passwd",
        "secret",
        "token",
        "api_key",
        "access_key",
        "auth_token",
        "authorization",
        "bearer",
    }

    def filter(self, record: logging.LogRecord) -> bool:
        # Optimized early exit
        msg_str = str(record.msg).lower()
        has_sensitive = any(sk in msg_str for sk in self._SENSITIVE_KEYS)

        # Check for sensitive keys in record __dict__ (which includes merged 'extra')
        record_keys = set(record.__dict__.keys())
        sensitive_in_dict = record_keys.intersection(self._SENSITIVE_KEYS)

        if not has_sensitive and not record.args and not sensitive_in_dict:
            return True

        if isinstance(record.msg, str) and has_sensitive:
            record.msg = self._redact_string(record.msg)

        if record.args:
            record.args = tuple(self._redact_arg(arg) for arg in record.args)

        # Redact any sensitive attributes merged from 'extra'
        for k in sensitive_in_dict:
            record.__dict__[k] = "[REDACTED]"

        return True

    def _redact_string(self, text: str) -> str:
        for pattern in self._SECRET_PATTERNS:
            text = pattern.sub(r"\1\2[REDACTED]", text)
        return text

    def _redact_arg(self, arg: Any, depth: int = 0) -> Any:
        if depth > 5:
            return "[DEPTH_LIMIT_REDACTED]"

        if isinstance(arg, str):
            if any(sk in arg.lower() for sk in self._SENSITIVE_KEYS):
                return self._redact_string(arg)
            return arg

        if isinstance(arg, dict):
            return {
                k: (
                    "[REDACTED]"
                    if str(k).lower() in self._SENSITIVE_KEYS
                    else self._redact_arg(v, depth + 1)
                )
                for k, v in arg.items()
            }

        if isinstance(arg, (list, tuple)):
            return type(arg)(self._redact_arg(item, depth + 1) for item in arg)

        return arg


class Logger:
    """
    Thread-safe, non-blocking Logger with optional OpenTelemetry integration.
    """

    DEBUG = logging.DEBUG
    INFO = logging.INFO
    WARNING = logging.WARNING
    ERROR = logging.ERROR
    CRITICAL = logging.CRITICAL

    _lock = threading.RLock()
    _attached_keys: set[tuple[str, str]] = set()
    _otel_initialized_names: set[str] = set()
    
    # Static queue shared across instances for unified non-blocking logging
    _log_queue: Queue[logging.LogRecord | None] = Queue(-1)
    _listener: QueueListener | None = None

    def __init__(
        self,
        log_dir: str | Path,
        logger_name: str,
        log_file: str,
        log_level: int = logging.INFO,
        enable_otel: bool = False,
        otel_service_name: str | None = None,
        otel_stream_name: str | None = None,
        otel_endpoint: str = "0.0.0.0:4317",
        otel_insecure: bool = False,
    ) -> None:
        self.log_dir = Path(log_dir).resolve()
        self.logger_name = logger_name
        self.log_file = log_file
        self.log_level = log_level

        self.enable_otel = bool(enable_otel and _OTEL_AVAILABLE)
        self.otel_service_name = (
            otel_service_name or logger_name or "unknown_service"
        ).strip()
        self.otel_stream_name = (otel_stream_name or "").strip() or None
        self.otel_endpoint = otel_endpoint
        self.otel_insecure = otel_insecure

        self.logger_provider = None
        self.tracer_provider = None
        self.tracer = None

        self._core = logging.getLogger(self.logger_name)
        self._core.setLevel(self.log_level)
        self._core.propagate = False

        # Public logger ref
        self.public_logger: logging.Logger | LoggerAdapter = self._core

        self._setup_handlers()

        if self.enable_otel:
            self._setup_otel()
            if self.otel_stream_name:
                attrs = {
                    "log.stream": self.otel_stream_name,
                    "service.name": self.otel_service_name,
                }
                self.public_logger = LoggerAdapter(self._core, extra=attrs)

    @classmethod
    def default_logger(
        cls,
        log_dir: str | Path = "logs",
        logger_name: str | None = None,
        log_file: str | None = None,
        log_level: int = logging.INFO,
        enable_otel: bool = False,
        otel_service_name: str | None = None,
        otel_stream_name: str | None = None,
        otel_endpoint: str = "0.0.0.0:4317",
        otel_insecure: bool = False,
    ) -> Logger:
        """Factory for quick instantiation with absolute path defaults."""
        try:
            caller_frame = sys._getframe(1)
            caller_name = caller_frame.f_globals.get("__name__", "default_logger")
        except (ValueError, AttributeError):
            caller_name = "default_logger"

        logger_name = logger_name or caller_name
        log_file = log_file or logger_name
        
        # Ensure log_dir is resolved relative to CWD if relative
        log_dir = Path(log_dir)
        if not log_dir.is_absolute():
            log_dir = Path.cwd() / log_dir

        return cls(
            log_dir=log_dir,
            logger_name=logger_name,
            log_file=log_file,
            log_level=log_level,
            enable_otel=enable_otel,
            otel_service_name=otel_service_name,
            otel_stream_name=otel_stream_name,
            otel_endpoint=otel_endpoint,
            otel_insecure=otel_insecure,
        )

    def set_level(self, level: int) -> None:
        self._core.setLevel(level)

    def setLevel(self, level: int) -> None:
        self.set_level(level)

    def flush(self) -> None:
        """
        Force flush all handlers. For QueueHandler, this just flushes destination handlers, 
        but doesn't guarantee the queue is empty. Use with care.
        """
        for h in self._core.handlers:
            h.flush()
        if Logger._listener:
            # There isn't a native 'flush' for listener, but we can wait a bit
            # or put a sentinel and wait for it. For now, simple sleep to avoid complexity.
            time.sleep(0.05)

    @property
    def handlers(self) -> list[logging.Handler]:
        return list(self._core.handlers)

    def _log(self, level: int, msg: str, *args: Any, **kwargs: Any) -> None:
        """
        Optimized logged that avoids redundant Adapter creation.
        """
        extra = kwargs.get("extra")
        
        # If we have a public adapter and NO override extra, use the adapter directly
        if isinstance(self.public_logger, LoggerAdapter) and not extra:
            self.public_logger.log(level, msg, *args, **kwargs)
            return

        # If we have extra overrides or are using the core logger
        self._core.log(level, msg, *args, **kwargs)

    def debug(self, msg: str, *a: Any, **k: Any) -> None:
        self._log(logging.DEBUG, msg, *a, **k)

    def info(self, msg: str, *a: Any, **k: Any) -> None:
        self._log(logging.INFO, msg, *a, **k)

    def warning(self, msg: str, *a: Any, **k: Any) -> None:
        self._log(logging.WARNING, msg, *a, **k)

    def error(self, msg: str, *a: Any, **k: Any) -> None:
        self._log(logging.ERROR, msg, *a, **k)

    def critical(self, msg: str, *a: Any, **k: Any) -> None:
        self._log(logging.CRITICAL, msg, *a, **k)

    def bind(self, **extra: Any) -> LoggerAdapter:
        if isinstance(self.public_logger, LoggerAdapter):
            return LoggerAdapter(self.public_logger.logger, {**self.public_logger.extra, **extra})
        return LoggerAdapter(self._core, extra)

    @contextmanager
    def bound(self, **extra: Any) -> Any:
        yield self.bind(**extra)

    def start_span(self, name: str, attributes: dict[str, Any] | None = None) -> Any:
        if not (self.enable_otel and self.tracer):
            return nullcontext()
        return self.tracer.start_as_current_span(name, attributes=attributes)

    def trace_function(self, span_name: str | None = None) -> Any:
        def deco(func: Any) -> Any:
            def wrapper(*a: Any, **k: Any) -> Any:
                name = span_name or func.__name__
                with self.start_span(name):
                    return func(*a, **k)
            return wrapper
        return deco

    def shutdown(self) -> None:
        """Shutdown OTel and flush handlers."""
        if self.enable_otel:
            with suppress(Exception):
                if self.logger_provider:
                    self.logger_provider.force_flush()
                    self.logger_provider.shutdown()
                if self.tracer_provider:
                    self.tracer_provider.force_flush()
                    self.tracer_provider.shutdown()

        for h in list(self._core.handlers):
            with suppress(Exception):
                h.flush()
                h.close()
                self._core.removeHandler(h)

    def _setup_handlers(self) -> None:
        self.log_dir.mkdir(parents=True, exist_ok=True)

        try:
            script_name = sys.argv[0] or "interactive"
        except IndexError:
            script_name = "unknown"

        calling_script = Path(script_name).stem
        log_file_path = self.log_dir / f"{self.log_file}_{calling_script}.log"

        file_key = (self.logger_name, str(log_file_path.resolve()))
        console_key = (self.logger_name, "__console__")

        fmt = logging.Formatter(
            "[%(asctime)s][%(levelname)s][%(name)s] %(message)s",
            datefmt="%Y-%m-%d %H:%M:%S",
        )
        fmt.converter = time.localtime

        with self._lock:
            # Initialize QueueListener if not already present
            if Logger._listener is None:
                # Actual destination handlers (not attached to loggers directly)
                # We use internal lists for the listener
                Logger._init_global_listener()

            # Attach non-blocking QueueHandler to this instance's core logger
            # but only if not already mixed in
            if not any(isinstance(h, QueueHandler) for h in self._core.handlers):
                qh = QueueHandler(Logger._log_queue)
                qh.addFilter(PIISecretFilter())
                self._core.addHandler(qh)

            # Ensure the specific destinations are added to the global listener
            Logger._add_destination(file_key, RotatingFileHandler(
                log_file_path, maxBytes=5 * 1024 * 1024, backupCount=5, delay=True
            ), fmt)
            
            Logger._add_destination(console_key, logging.StreamHandler(sys.stdout), fmt)

    @staticmethod
    def _init_global_listener() -> None:
        # Placeholder handlers to start the listener; we update them dynamically
        Logger._listener = QueueListener(Logger._log_queue, respect_handler_level=True)
        Logger._listener.start()
        atexit.register(Logger._stop_listener)

    @staticmethod
    def _add_destination(key: tuple[str, str], handler: logging.Handler, formatter: logging.Formatter) -> None:
        if key in Logger._attached_keys:
            return
        
        logger_name, _ = key
        # Ensure this handler only processes records for its intended logger
        handler.addFilter(logging.Filter(name=logger_name))
        
        handler.setFormatter(formatter)
        if Logger._listener:
            Logger._listener.handlers = Logger._listener.handlers + (handler,)
        
        Logger._attached_keys.add(key)

    @staticmethod
    def _stop_listener() -> None:
        if Logger._listener:
            Logger._log_queue.put_nowait(None)
            Logger._listener.stop()

    def _setup_otel(self) -> None:
        if not _OTEL_AVAILABLE:
            return

        with self._lock:
            if self.logger_name in self._otel_initialized_names:
                self.tracer = trace.get_tracer(self.logger_name)
                return

            resource = Resource.create({
                "service.name": self.otel_service_name,
                "logger.name": self.logger_name,
                **({"log.stream": self.otel_stream_name} if self.otel_stream_name else {})
            })
            endpoint = self._normalize_otlp_endpoint(self.otel_endpoint)

            # Provider Setup
            existing_lp = get_logger_provider()
            self.logger_provider = existing_lp if hasattr(existing_lp, "add_log_record_processor") else LoggerProvider(resource=resource)
            if self.logger_provider != existing_lp:
                with suppress(Exception): set_logger_provider(self.logger_provider)

            existing_tp = trace.get_tracer_provider()
            self.tracer_provider = existing_tp if hasattr(existing_tp, "add_span_processor") else TracerProvider(resource=resource)
            if self.tracer_provider != existing_tp:
                with suppress(Exception): trace.set_tracer_provider(self.tracer_provider)

            # Exporters
            with suppress(Exception):
                log_exp = OTLPLogExporter(endpoint=endpoint, insecure=self.otel_insecure)
                self.logger_provider.add_log_record_processor(BatchLogRecordProcessor(log_exp))
                
                span_exp = OTLPSpanExporter(endpoint=endpoint, insecure=self.otel_insecure)
                self.tracer_provider.add_span_processor(BatchSpanProcessor(span_exp))

            # Logging integration
            if not any(isinstance(h, LoggingHandler) for h in self._core.handlers):
                self._core.addHandler(LoggingHandler(level=logging.NOTSET, logger_provider=self.logger_provider))

            self.tracer = trace.get_tracer(self.logger_name)
            self._otel_initialized_names.add(self.logger_name)

    def _normalize_otlp_endpoint(self, ep: str) -> str:
        if "://" not in ep:
            return ("http://" if self.otel_insecure else "https://") + ep
        return ep
