"""Logging configuration and setup for Alloy Runtime."""

import json
import logging
import re
import sys
import uuid
from collections.abc import Iterable
from dataclasses import dataclass, field
from decimal import Decimal
from datetime import datetime, timezone
from enum import Enum
from pathlib import Path
from typing import Any, TextIO, cast

import structlog

try:
    from celery.signals import setup_logging as celery_setup_logging  # pyright: ignore[reportUnknownVariableType]
except ImportError:
    celery_setup_logging = None


# Keys that commonly contain sensitive data
SENSITIVE_KEYS = frozenset(
    {
        "api_key",
        "password",
        "secret",
        "authorization",
        "auth",
        "credential",
        "key",
        "bearer",
    }
)

LITELLM_LOGGER_NAMES = (
    "LiteLLM",
    "LiteLLM Router",
    "LiteLLM Proxy",
    "litellm",
)


def _add_session_id(session_id: str) -> structlog.types.Processor:
    """Create a processor that adds session_id to all log entries.

    Args:
        session_id: The session ID to add to all logs

    Returns:
        A structlog processor function
    """

    def processor(
        logger: structlog.types.WrappedLogger,
        method_name: str,
        event_dict: structlog.types.EventDict,
    ) -> structlog.types.EventDict:
        _ = logger, method_name  # Required by structlog processor signature
        event_dict["session_id"] = session_id
        return event_dict

    return processor


def _add_service_context(
    service_name: str, environment: str
) -> structlog.types.Processor:
    """Create a processor that adds service name and environment to all log entries.

    Args:
        service_name: The service name to add to all logs
        environment: The environment (development, staging, production)

    Returns:
        A structlog processor function
    """

    def processor(
        logger: structlog.types.WrappedLogger,
        method_name: str,
        event_dict: structlog.types.EventDict,
    ) -> structlog.types.EventDict:
        _ = logger, method_name  # Required by structlog processor signature
        event_dict["service"] = service_name
        event_dict["environment"] = environment
        return event_dict

    return processor


def _sanitize_sensitive_data(
    logger: structlog.types.WrappedLogger,
    method_name: str,
    event_dict: structlog.types.EventDict,
) -> structlog.types.EventDict:
    """Remove or mask sensitive data from logs.

    Recursively searches through the event dict and masks any values
    where the key contains sensitive words.

    Args:
        logger: The wrapped logger
        method_name: The logging method name
        event_dict: The event dictionary to sanitize

    Returns:
        The sanitized event dictionary
    """
    _ = logger, method_name  # Required by structlog processor signature

    def _contains_sensitive_word(key: str) -> bool:
        """Check if key contains any sensitive word."""
        key_lower = key.lower()
        return any(sensitive in key_lower for sensitive in SENSITIVE_KEYS)

    def _sanitize_value(value: Any) -> Any:
        """Mask a sensitive value."""
        if isinstance(value, str):
            # Skip already-redacted values
            if "[REDACTED]" in value or "***REDACTED***" in value:
                return value
            if len(value) > 8:
                return value[:4] + "***REDACTED***"
            return "***REDACTED***"
        return "***REDACTED***"

    def _sanitize_dict(d: dict[str, Any]) -> dict[str, Any]:
        """Recursively sanitize a dictionary."""
        result: dict[str, Any] = {}
        for key, value in d.items():
            if _contains_sensitive_word(key):
                result[key] = _sanitize_value(value)
            elif isinstance(value, dict):
                result[key] = _sanitize_dict(cast(dict[str, Any], value))
            elif isinstance(value, list):
                sanitized_items: list[Any] = []
                for item in cast(list[Any], value):
                    if isinstance(item, dict):
                        sanitized_items.append(
                            _sanitize_dict(cast(dict[str, Any], item))
                        )
                    else:
                        sanitized_items.append(item)
                result[key] = sanitized_items
            else:
                result[key] = value
        return result

    return _sanitize_dict(dict(event_dict))


_ANSI_RE = re.compile(r"\x1B\[[0-?]*[ -/]*[@-~]")


def _normalize_log_value(value: Any) -> Any:
    if isinstance(value, (bool, str, int, float)):
        return value
    if value is None:
        return None
    if isinstance(value, Decimal):
        return str(value)
    if isinstance(value, Enum):
        return _normalize_log_value(value.value)
    if isinstance(value, (uuid.UUID, Path, datetime)):
        return str(value)
    if isinstance(value, bytes):
        return value.decode("utf-8", errors="replace")
    if isinstance(value, dict):
        mapping = cast(dict[object, Any], value)
        return {str(key): _normalize_log_value(item) for key, item in mapping.items()}
    if isinstance(value, (list, tuple, set, frozenset)):
        items = cast(Iterable[Any], value)
        return [_normalize_log_value(item) for item in items]
    return str(value)


def _normalize_otel_value(value: Any) -> Any:
    if isinstance(value, (bool, str, int, float, bytes)):
        return value
    if value is None:
        return "None"
    if isinstance(value, Decimal):
        return str(value)
    if isinstance(value, Enum):
        return _normalize_otel_value(value.value)
    if isinstance(value, (uuid.UUID, Path, datetime)):
        return str(value)
    if isinstance(value, dict):
        mapping = cast(dict[object, Any], value)
        return {str(key): _normalize_otel_value(item) for key, item in mapping.items()}
    if isinstance(value, (list, tuple, set, frozenset)):
        items = list(cast(Iterable[Any], value))
        return [_normalize_otel_value(item) for item in items]
    return str(value)


def _sanitize_foreign_log_event(
    logger: structlog.types.WrappedLogger,
    method_name: str,
    event_dict: structlog.types.EventDict,
) -> structlog.types.EventDict:
    """Strip ANSI codes and collapse whitespace in foreign log messages."""
    _ = logger, method_name
    event = event_dict.get("event", "")
    if isinstance(event, str):
        event = _ANSI_RE.sub("", event)
        event = " ".join(event.split())
        event_dict["event"] = event
    return event_dict


def _normalize_log_event_values(
    logger: structlog.types.WrappedLogger,
    method_name: str,
    event_dict: structlog.types.EventDict,
) -> structlog.types.EventDict:
    _ = logger, method_name
    normalized: structlog.types.EventDict = {}
    for key, value in event_dict.items():
        normalized[key] = _normalize_log_value(value)
    return normalized


class LogFormat(str, Enum):
    """Output format for log entries."""

    JSON = "json"  # Production: {"timestamp": "...", "level": "info", ...}
    CONSOLE = "console"  # Development: colored console output


class LogLevel(str, Enum):
    """Log level settings."""

    DEBUG = "debug"
    INFO = "info"
    WARNING = "warning"
    ERROR = "error"


def _generate_session_id() -> str:
    """Generate a session ID that identifies this process start.

    Format: YYYYMMDD-HHMMSS-<short-uuid>
    Example: 20251126-143052-a1b2c3d4

    This makes it easy to:
    - Find all logs from a specific server restart
    - Filter with jq: .session_id == "20251126-143052-a1b2c3d4"
    - grep for the date portion
    """
    now = datetime.now()
    timestamp = now.strftime("%Y%m%d-%H%M%S")
    short_uuid = uuid.uuid4().hex[:8]
    return f"{timestamp}-{short_uuid}"


@dataclass
class LogConfig:
    """Configuration for the unified logging system."""

    # Core settings
    level: LogLevel = LogLevel.INFO
    format: LogFormat = LogFormat.JSON

    # Output settings
    output_file: str | Path | None = None  # If set, write logs to file
    include_timestamp: bool = True
    include_caller: bool = False  # Include file:line info

    # Context settings
    service_name: str = "alloy-runtime"  # Identifies the service
    environment: str = "development"  # development, staging, production

    # Session tracking (for finding logs between restarts)
    session_id: str = field(default_factory=_generate_session_id)

    # Integration flags
    capture_stdlib_logging: bool = True  # Route stdlib logging through structlog
    configure_sqlalchemy: bool = False  # Configure SQLAlchemy logging
    configure_litellm: bool = False  # Configure LiteLLM logging
    configure_celery: bool = False  # Configure Celery logging to use structlog
    sqlalchemy_level: LogLevel = LogLevel.WARNING  # SQLAlchemy default level


# Module-level state
_configured = False
_current_session_id: str | None = None
_output_file_handle: TextIO | None = None
_our_handler: logging.Handler | None = None
_our_filter: logging.Filter | None = None
_stamped_handlers: list[logging.Handler] = []


class _OtelAttributeStamper(logging.Filter):
    """Stamp structlog attributes onto LogRecord for OpenTelemetry visibility.

    When structlog uses wrap_for_formatter, structured attributes are stored
    inside a _FormatArgs wrapper on record.msg. This filter extracts them to
    record.__dict__ so OTel's LoggingHandler can access them via _get_attributes().

    For foreign (non-structlog) entries, stamps basic service context.
    """

    def __init__(self, session_id: str, service_name: str, environment: str):
        super().__init__()
        self._session_id = session_id
        self._service_name = service_name
        self._environment = environment

    def filter(self, record: logging.LogRecord) -> bool:
        msg = record.msg
        event_dict = getattr(msg, "event_dict", None)
        if isinstance(event_dict, dict):
            for key, value in cast(dict[str, Any], event_dict).items():
                if not key.startswith("_"):
                    record.__dict__[key] = _normalize_otel_value(value)
        elif isinstance(msg, dict):
            for key, value in cast(dict[str, Any], msg).items():
                if not key.startswith("_"):
                    record.__dict__[key] = _normalize_otel_value(value)
        else:
            for key, value in structlog.contextvars.get_contextvars().items():
                if not key.startswith("_") and key not in record.__dict__:
                    record.__dict__[key] = _normalize_otel_value(value)
            if "session_id" not in record.__dict__:
                record.__dict__["session_id"] = self._session_id
            if "service" not in record.__dict__:
                record.__dict__["service"] = self._service_name
            if "environment" not in record.__dict__:
                record.__dict__["environment"] = self._environment
        return True


class _StructlogOnlyFilter(logging.Filter):
    def filter(self, record: logging.LogRecord) -> bool:
        msg = record.msg
        return bool(getattr(msg, "event_dict", None) or isinstance(msg, dict))


def is_configured() -> bool:
    """Check if logging has been configured."""
    return _configured


def get_session_id() -> str | None:
    """Get the current logging session ID."""
    return _current_session_id


def get_logger(name: str | None = None) -> structlog.stdlib.BoundLogger:
    """Get a logger instance.

    Args:
        name: Optional logger name for context

    Returns:
        A structlog bound logger instance
    """
    logger = structlog.get_logger(name) if name else structlog.get_logger()
    return logger


def get_or_configure_logger(
    name: str | None = None,
    *,
    service_name: str = "alloy-runtime-sdk",
) -> structlog.stdlib.BoundLogger:
    if not is_configured():
        configure(LogConfig(format=LogFormat.JSON, service_name=service_name))
    return get_logger(name)


def configure(
    config: LogConfig | None = None,
    output: TextIO | None = None,
) -> str:
    """Configure the unified logging system.

    Args:
        config: Logging configuration. If None, uses defaults.
        output: Output stream. Defaults to sys.stderr for console format.
                For JSON format with file output, this is ignored.

    Returns:
        The session_id for this logging session.

    Example:
        # Server startup
        session_id = configure(LogConfig(
            level=LogLevel.INFO,
            format=LogFormat.JSON,
            service_name="alloy-runtime-server",
            configure_sqlalchemy=True,
            configure_litellm=True,
        ))

        # CLI usage (logs to file)
        configure(LogConfig(
            level=LogLevel.WARNING,
            format=LogFormat.JSON,
            service_name="alloy-runtime-cli",
            output_file="logs/cli.log",
        ))
    """
    global \
        _configured, \
        _current_session_id, \
        _output_file_handle, \
        _our_handler, \
        _our_filter, \
        _stamped_handlers

    config = config or LogConfig()
    _current_session_id = config.session_id

    # Close any existing file handle to prevent leaks on reconfigure
    if _output_file_handle is not None:
        _output_file_handle.close()
        _output_file_handle = None

    # Determine output stream
    if config.output_file:
        # Ensure parent directory exists
        output_path = Path(config.output_file)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        # Open with line buffering for immediate writes
        _output_file_handle = open(output_path, "a", encoding="utf-8", buffering=1)
        log_output = _output_file_handle
    elif output:
        log_output = output
    else:
        log_output = sys.stderr

    shared_processors: list[structlog.types.Processor] = [
        structlog.contextvars.merge_contextvars,
        structlog.stdlib.add_log_level,
        _add_session_id(config.session_id),
        _add_service_context(config.service_name, config.environment),
        _sanitize_sensitive_data,
        _normalize_log_event_values,
        structlog.processors.StackInfoRenderer(),
        structlog.dev.set_exc_info,
    ]

    if config.include_timestamp:
        shared_processors.insert(2, structlog.processors.TimeStamper(fmt="iso"))

    if config.include_caller:
        shared_processors.append(
            structlog.processors.CallsiteParameterAdder(
                [
                    structlog.processors.CallsiteParameter.FILENAME,
                    structlog.processors.CallsiteParameter.LINENO,
                ]
            )
        )

    structlog.configure(
        processors=shared_processors
        + [
            structlog.stdlib.ProcessorFormatter.wrap_for_formatter,
        ],
        wrapper_class=structlog.make_filtering_bound_logger(
            getattr(logging, config.level.value.upper())
        ),
        context_class=dict,
        logger_factory=structlog.stdlib.LoggerFactory(),
        cache_logger_on_first_use=True,
    )

    if config.format == LogFormat.JSON:
        renderer: structlog.types.Processor = structlog.processors.JSONRenderer(
            sort_keys=False,
            ensure_ascii=False,
        )
    else:
        renderer = structlog.dev.ConsoleRenderer(
            colors=True,
            exception_formatter=structlog.dev.plain_traceback,
        )

    formatter = structlog.stdlib.ProcessorFormatter(
        foreign_pre_chain=[_sanitize_foreign_log_event] + shared_processors,
        processors=[
            structlog.stdlib.ProcessorFormatter.remove_processors_meta,
            renderer,
        ],
    )

    root_logger = logging.getLogger()

    if _our_filter is not None:
        for stamped_handler in _stamped_handlers:
            if _our_filter in stamped_handler.filters:
                stamped_handler.removeFilter(_our_filter)
        _stamped_handlers.clear()

    if _our_handler is not None:
        root_logger.removeHandler(_our_handler)

    handler = logging.StreamHandler(log_output)
    handler.setFormatter(formatter)
    if not config.capture_stdlib_logging:
        handler.addFilter(_StructlogOnlyFilter())
    root_logger.addHandler(handler)
    root_logger.setLevel(getattr(logging, config.level.value.upper()))
    _our_handler = handler

    otel_filter = _OtelAttributeStamper(
        session_id=config.session_id,
        service_name=config.service_name,
        environment=config.environment,
    )
    for root_handler in root_logger.handlers:
        root_handler.addFilter(otel_filter)
        _stamped_handlers.append(root_handler)
    _our_filter = otel_filter

    # Configure stdlib logging bridge if requested
    if config.capture_stdlib_logging:
        _configure_stdlib_bridge(config.level)

    # Configure SQLAlchemy logging if requested
    if config.configure_sqlalchemy:
        _configure_sqlalchemy_logging(config.sqlalchemy_level)

    # Configure LiteLLM logging if requested
    if config.configure_litellm:
        _configure_litellm_logging(config.level)

    # Configure Celery logging if requested
    if config.configure_celery:
        _configure_celery_logging(config.level)

    # Ensure OTel LoggingHandler is on root logger for log export.
    # opentelemetry-instrument may have attached it before configure() ran,
    # but handler cleanup could have removed it. This is a safety net.
    _ensure_otel_logging_handler(root_logger)
    _configured = True
    return config.session_id


def _ensure_otel_logging_handler(root_logger: logging.Logger) -> None:
    """Ensure an OTel LoggingHandler is attached to the root logger.

    opentelemetry-instrument's auto-instrumentation is supposed to attach a
    LoggingHandler when OTEL_PYTHON_LOGGING_AUTO_INSTRUMENTATION_ENABLED=true,
    but it may not survive our configure() handler setup, or may never have
    been attached. This function explicitly creates one if missing.
    """
    import os

    if (
        os.environ.get("OTEL_PYTHON_LOGGING_AUTO_INSTRUMENTATION_ENABLED", "").lower()
        != "true"
    ):
        return

    # Check if an OTel LoggingHandler is already present
    for h in root_logger.handlers:
        if (
            type(h).__name__ == "LoggingHandler"
            and "opentelemetry" in type(h).__module__
        ):
            return  # Already attached

    try:
        from opentelemetry._logs import get_logger_provider
        from opentelemetry.sdk._logs import LoggerProvider

        logger_provider = get_logger_provider()
        # Only attach if the SDK was actually configured (not a no-op proxy)
        if not isinstance(logger_provider, LoggerProvider):
            return

        from opentelemetry.sdk._logs import LoggingHandler

        otel_handler = LoggingHandler(logger_provider=logger_provider)
        otel_handler.setLevel(logging.NOTSET)  # Let root logger level filter
        root_logger.addHandler(otel_handler)
    except Exception:
        pass  # OTel SDK not available or not configured — skip silently


def _configure_stdlib_bridge(level: LogLevel) -> None:
    """Suppress noisy third-party stdlib loggers."""
    _ = level
    logging.getLogger("httpx").setLevel(logging.WARNING)
    logging.getLogger("httpcore").setLevel(logging.WARNING)
    logging.getLogger("asyncio").setLevel(logging.WARNING)


def _configure_sqlalchemy_logging(level: LogLevel) -> None:
    """Configure SQLAlchemy to use structured logging."""
    sqlalchemy_logger = logging.getLogger("sqlalchemy.engine")
    sqlalchemy_logger.setLevel(getattr(logging, level.value.upper()))


def _configure_litellm_logging(level: LogLevel) -> None:
    """Configure LiteLLM to use structured logging via propagation to root handler."""
    log_level = getattr(logging, level.value.upper())

    for logger_name in LITELLM_LOGGER_NAMES:
        lg = logging.getLogger(logger_name)
        lg.handlers.clear()
        lg.setLevel(log_level)
        lg.propagate = True


def _configure_celery_logging(level: LogLevel) -> None:
    if celery_setup_logging is None:
        return

    log_level = getattr(logging, level.value.upper())

    def on_celery_setup_logging(**kwargs: Any) -> None:
        _ = kwargs
        for logger_name in [
            "celery",
            "celery.task",
            "celery.worker",
            "celery.beat",
            "celery.redirected",
        ]:
            lg = logging.getLogger(logger_name)
            lg.handlers.clear()
            lg.setLevel(log_level)
            lg.propagate = True

    setup_logging_signal = celery_setup_logging
    setup_logging_signal.connect(  # pyright: ignore[reportUnknownMemberType]
        on_celery_setup_logging,
        weak=False,
        dispatch_uid="alloy_runtime_celery_structlog",
    )


def cleanup() -> None:
    """Clean up logging resources (close file handles)."""
    global _output_file_handle, _our_handler, _our_filter, _stamped_handlers
    if _output_file_handle:
        _output_file_handle.close()
        _output_file_handle = None
    if _our_filter:
        for stamped_handler in _stamped_handlers:
            if _our_filter in stamped_handler.filters:
                stamped_handler.removeFilter(_our_filter)
        _stamped_handlers = []
    if _our_handler:
        logging.getLogger().removeHandler(_our_handler)
        _our_handler = None
    _our_filter = None


class UvicornJsonFormatter(logging.Formatter):
    """JSON formatter for uvicorn logs.

    Used by uvicorn's log config to output JSON-structured logs that match
    our structlog format. This is loaded by uvicorn BEFORE our main app
    starts, ensuring consistent JSON output from the very first log message.
    """

    def format(self, record: logging.LogRecord) -> str:
        """Format a log record as JSON."""
        log_entry = {
            "event": record.getMessage(),
            "level": record.levelname.lower(),
            "timestamp": self.formatTime(record, self.datefmt),
            "logger": record.name,
        }

        # Add exception info if present
        if record.exc_info:
            log_entry["exception"] = self.formatException(record.exc_info)

        return json.dumps(log_entry, ensure_ascii=False)

    def formatTime(self, record: logging.LogRecord, datefmt: str | None = None) -> str:
        """Format time in ISO 8601 format with Z suffix."""
        dt = datetime.fromtimestamp(record.created, tz=timezone.utc)
        return dt.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "Z"
