# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Structured logging helpers for :mod:`weakincentives`."""

from __future__ import annotations

import json
import logging
import logging.config
import os
from collections.abc import Mapping, MutableMapping
from datetime import UTC, datetime
from typing import TYPE_CHECKING, Protocol, cast, override

from ..types import JSONValue

if TYPE_CHECKING:
    from .run_context import RunContext

type StructuredLogPayload = Mapping[str, JSONValue]


class SupportsLogMessage(Protocol):
    """Protocol describing values that logging can stringify."""

    @override
    def __str__(self) -> str: ...


__all__ = [
    "JSONValue",
    "StructuredLogPayload",
    "StructuredLogger",
    "SupportsLogMessage",
    "bind_run_context",
    "configure_logging",
    "get_logger",
]

_LOG_LEVEL_ENV = "WEAKINCENTIVES_LOG_LEVEL"
_LOG_FORMAT_ENV = "WEAKINCENTIVES_LOG_FORMAT"
_LEVEL_NAMES = {
    "CRITICAL": logging.CRITICAL,
    "ERROR": logging.ERROR,
    "WARNING": logging.WARNING,
    "INFO": logging.INFO,
    "DEBUG": logging.DEBUG,
    "NOTSET": logging.NOTSET,
}


class StructuredLogger(logging.LoggerAdapter[logging.Logger]):
    """Logger adapter enforcing a minimal structured event schema."""

    def __init__(
        self,
        logger: logging.Logger,
        *,
        context: StructuredLogPayload | None = None,
    ) -> None:
        base_context: dict[str, JSONValue] = (
            dict(context) if context is not None else {}
        )
        super().__init__(logger, base_context)
        self._context: dict[str, JSONValue] = base_context

    def bind(self, **context: JSONValue) -> StructuredLogger:
        """Return a new adapter with ``context`` merged into the baseline payload."""

        merged: dict[str, JSONValue] = {**dict(self._context), **context}
        return type(self)(self.logger, context=merged)

    @staticmethod
    def _get_extra_mapping(
        kwargs: MutableMapping[str, object],
    ) -> MutableMapping[str, JSONValue]:
        """Extract or create the extra mapping from kwargs."""
        extra_value = kwargs.setdefault("extra", {})
        if extra_value is None:
            extra_mapping: MutableMapping[str, JSONValue] = {}
            kwargs["extra"] = extra_mapping
            return extra_mapping
        if isinstance(extra_value, MutableMapping):
            return cast(MutableMapping[str, JSONValue], extra_value)
        raise TypeError(  # pragma: no cover - defensive guard
            "Structured logs require a mutable mapping for extra context."
        )

    @staticmethod
    def _extract_event(
        kwargs: MutableMapping[str, object],
        extra_mapping: MutableMapping[str, JSONValue],
    ) -> str:
        """Extract the event field from kwargs or extra_mapping."""
        event_obj = kwargs.pop("event", None)
        if event_obj is None:
            event_obj = extra_mapping.pop("event", None)
        if not isinstance(event_obj, str):
            raise TypeError("Structured logs require an 'event' field.")
        return event_obj

    @override
    def process(
        self, msg: SupportsLogMessage, kwargs: MutableMapping[str, object]
    ) -> tuple[SupportsLogMessage, MutableMapping[str, object]]:
        extra_mapping = self._get_extra_mapping(kwargs)
        context_payload: dict[str, JSONValue] = dict(self._context)

        inline_context = kwargs.pop("context", None)
        if inline_context is not None:
            if not isinstance(inline_context, Mapping):  # pragma: no cover - defensive
                raise TypeError("context must be a mapping when provided.")
            context_payload.update(cast(StructuredLogPayload, inline_context))

        for key in tuple(extra_mapping.keys()):
            if key != "event":
                context_payload[key] = extra_mapping.pop(key)

        event_obj = self._extract_event(kwargs, extra_mapping)

        extra_mapping.clear()
        extra_mapping.update({"event": event_obj, "context": context_payload})
        return msg, kwargs


def get_logger(
    name: str,
    *,
    logger_override: logging.Logger
    | logging.LoggerAdapter[logging.Logger]
    | None = None,
    context: StructuredLogPayload | None = None,
) -> StructuredLogger:
    """Return a :class:`StructuredLogger` scoped to ``name``.

    When ``logger_override`` is provided, the returned adapter reuses the supplied
    logger and merges its contextual ``extra`` payload when available.
    """

    base_context: dict[str, JSONValue] = dict(context or {})
    base_logger: logging.Logger

    if isinstance(logger_override, StructuredLogger):
        base_logger = logger_override.logger
        base_context = {
            **dict(cast(StructuredLogPayload, logger_override.extra)),
            **base_context,
        }
    elif isinstance(logger_override, logging.Logger):
        base_logger = logger_override
    elif isinstance(logger_override, logging.LoggerAdapter):
        base_logger = _unwrap_logger(cast(_SupportsNestedLogger, logger_override))
        adapter_extra = getattr(logger_override, "extra", None)
        if isinstance(adapter_extra, Mapping):
            base_context = {
                **dict(cast(StructuredLogPayload, adapter_extra)),
                **base_context,
            }
    else:
        base_logger = logging.getLogger(name)

    return StructuredLogger(base_logger, context=base_context)


def bind_run_context(
    logger: StructuredLogger,
    run_context: RunContext | None,
) -> StructuredLogger:
    """Bind RunContext fields to a logger for request-scoped tracing.

    When run_context is provided, returns a new logger with run_id, request_id,
    attempt, worker_id, and trace context bound. When run_context is None,
    returns the original logger unchanged.

    This enables consistent correlation of logs across the entire request
    lifecycle: mailbox receipt -> prompt render -> provider call -> tool calls
    -> completion -> reply/DLQ.

    Args:
        logger: The structured logger to bind context to.
        run_context: Optional execution context with correlation identifiers.

    Returns:
        Logger with run context fields bound, or original logger if no context.

    Example:
        >>> from weakincentives.runtime import RunContext, get_logger
        >>> from weakincentives.runtime.logging import bind_run_context
        >>> base_log = get_logger(__name__)
        >>> run_ctx = RunContext(worker_id="worker-1")
        >>> log = bind_run_context(base_log, run_ctx)
        >>> log.info("Processing request", event="request.start")
    """
    if run_context is None:
        return logger
    return logger.bind(**run_context.to_log_context())


def configure_logging(
    *,
    level: int | str | None = None,
    json_mode: bool | None = None,
    env: Mapping[str, str] | None = None,
    force: bool = False,
) -> None:
    """Configure the root logger with sensible defaults.

    ``level`` and ``json_mode`` can be supplied directly or via the
    ``WEAKINCENTIVES_LOG_LEVEL`` and ``WEAKINCENTIVES_LOG_FORMAT`` environment
    variables respectively (``json`` enables structured output, ``text`` keeps the
    plain formatter).

    The function avoids installing duplicate handlers when the host application has
    already configured logging unless ``force=True`` is supplied.
    """

    env = env or os.environ

    if level is not None:
        resolved_level = _coerce_level(level)
    else:
        resolved_level = _coerce_level(env.get(_LOG_LEVEL_ENV))

    if json_mode is None:
        format_value = env.get(_LOG_FORMAT_ENV)
        if format_value is not None:
            json_mode = format_value.lower() == "json"
        else:
            json_mode = False

    root_logger = logging.getLogger()

    if root_logger.handlers and not force:
        root_logger.setLevel(resolved_level)
        return

    formatter_key = "json" if json_mode else "text"
    logging.config.dictConfig(
        {
            "version": 1,
            "disable_existing_loggers": False,
            "formatters": {
                "text": {
                    "()": "weakincentives.runtime.logging._TextFormatter",
                },
                "json": {
                    "()": "weakincentives.runtime.logging._JsonFormatter",
                },
            },
            "handlers": {
                "stderr": {
                    "class": "logging.StreamHandler",
                    "stream": "ext://sys.stderr",
                    "formatter": formatter_key,
                }
            },
            "root": {
                "handlers": ["stderr"],
                "level": resolved_level,
            },
        }
    )


class _SupportsNestedLogger(Protocol):
    logger: logging.Logger | logging.LoggerAdapter[logging.Logger]


class _TextFormatter(logging.Formatter):
    """Formatter that renders structured records as human-readable text.

    Handles log records from third-party libraries that lack the ``event``
    and ``context`` fields expected by the structured logging schema.
    """

    _FMT = "%(asctime)s %(levelname)s %(name)s"
    _DATEFMT = "%Y-%m-%d %H:%M:%S"

    def __init__(self) -> None:
        super().__init__(fmt=self._FMT, datefmt=self._DATEFMT)

    @override
    def format(self, record: logging.LogRecord) -> str:
        # Prepare record for formatting (sets record.message and record.asctime)
        record.message = record.getMessage()
        record.asctime = self.formatTime(record, self.datefmt)

        # Format base fields only (timestamp, level, logger name) - no exception
        base = self.formatMessage(record)

        # Extract structured fields if present
        event = getattr(record, "event", None)
        context = getattr(record, "context", None)

        # Build output parts
        parts = [base]
        if event is not None:
            parts.append(str(event))
        parts.append(record.message)
        if context:
            parts.append(str(context))

        result = " ".join(parts)

        # Append exception info if present (cache to avoid repeated formatting)
        if record.exc_info and not record.exc_text:
            record.exc_text = self.formatException(record.exc_info)
        if record.exc_text:
            result = f"{result}\n{record.exc_text}"
        if record.stack_info:
            result = f"{result}\n{self.formatStack(record.stack_info)}"

        return result


class _JsonFormatter(logging.Formatter):
    """Formatter that renders structured records as compact JSON."""

    @override
    def format(self, record: logging.LogRecord) -> str:
        payload: dict[str, JSONValue] = {
            "timestamp": datetime.fromtimestamp(record.created, tz=UTC).isoformat(),
            "level": record.levelname,
            "logger": record.name,
            "message": record.getMessage(),
        }
        event = getattr(record, "event", None)
        if event is not None:
            payload["event"] = event
        context = getattr(record, "context", None)
        if context:
            payload["context"] = context
        if record.exc_info:
            payload["exc_info"] = self.formatException(record.exc_info)
        return json.dumps(payload, default=_json_default, separators=(",", ":"))


def _json_default(value: object) -> JSONValue:
    """Fallback serializer returning ``repr`` for unsupported values."""

    return repr(value)


_TEXT_FORMATTER_CLASS = _TextFormatter
_JSON_FORMATTER_CLASS = _JsonFormatter


def _unwrap_logger(adapter: _SupportsNestedLogger) -> logging.Logger:
    """Return the underlying :class:`logging.Logger` from an adapter."""

    logger_value = cast(object, adapter.logger)
    if isinstance(logger_value, logging.LoggerAdapter):
        return _unwrap_logger(cast(_SupportsNestedLogger, logger_value))
    if isinstance(logger_value, logging.Logger):
        return logger_value
    raise TypeError("LoggerAdapter.logger must be a logging.Logger instance.")


def _coerce_level(level: int | str | None) -> int:
    if isinstance(level, int):
        return level
    if isinstance(level, str):
        try:
            return _LEVEL_NAMES[level.upper()]
        except KeyError:  # pragma: no cover - defensive guard
            raise TypeError(f"Unknown log level: {level!r}") from None
    return logging.INFO
