from __future__ import annotations

import json
import logging
import os
import sys
from datetime import datetime, timezone
from typing import Any, Dict

_DEFAULT_LOG_LEVEL = "INFO"
_RESERVED_LOG_KWARGS = {"exc_info", "stack_info", "stacklevel", "extra"}
_RESERVED_LOG_RECORD_FIELDS = {
    "args",
    "asctime",
    "created",
    "exc_info",
    "exc_text",
    "filename",
    "funcName",
    "levelname",
    "levelno",
    "lineno",
    "message",
    "module",
    "msecs",
    "msg",
    "name",
    "pathname",
    "process",
    "processName",
    "relativeCreated",
    "stack_info",
    "thread",
    "threadName",
}


class StructuredLoggerAdapter(logging.LoggerAdapter):
    """Logger adapter that converts arbitrary keyword args into structured fields."""

    def process(self, msg: str, kwargs: Dict[str, Any]) -> tuple[str, Dict[str, Any]]:
        extra = kwargs.get("extra")
        if extra is None:
            extra = {}
        else:
            extra = dict(extra)

        structured = {
            (f"log_{key}" if key in _RESERVED_LOG_RECORD_FIELDS else key): kwargs.pop(key)
            for key in list(kwargs.keys())
            if key not in _RESERVED_LOG_KWARGS
        }
        if structured:
            extra.update(structured)
            kwargs["extra"] = extra

        return msg, kwargs


class JsonFormatter(logging.Formatter):
    """Serialize log records as JSON lines."""

    def format(self, record: logging.LogRecord) -> str:  # noqa: D401
        payload: Dict[str, Any] = {
            "timestamp": datetime.fromtimestamp(record.created, tz=timezone.utc).isoformat(),
            "level": record.levelname,
            "logger": record.name,
            "message": record.getMessage(),
        }

        if record.exc_info:
            payload["exc_info"] = self.formatException(record.exc_info)

        if record.stack_info:
            payload["stack"] = record.stack_info

        for key, value in record.__dict__.items():
            if key.startswith("_"):
                continue
            if key in payload:
                continue
            if key in {
                "args",
                "msg",
                "message",
                "name",
                "levelno",
                "levelname",
                "pathname",
                "filename",
                "module",
                "lineno",
                "funcName",
                "created",
                "msecs",
                "relativeCreated",
                "thread",
                "threadName",
                "processName",
                "process",
            }:
                continue
            payload[key] = value

        return json.dumps(payload, ensure_ascii=False)


def configure_logging(level: str | None = None) -> None:
    """Configure root logging for the MCP server."""

    resolved_level = level or os.getenv("LOG_LEVEL", _DEFAULT_LOG_LEVEL)
    logging_level = getattr(logging, resolved_level.upper(), logging.INFO)

    handler = logging.StreamHandler(stream=sys.stderr)
    handler.setFormatter(JsonFormatter())

    root_logger = logging.getLogger()
    root_logger.setLevel(logging_level)
    root_logger.handlers = [handler]
    root_logger.propagate = False


def get_logger(name: str) -> logging.LoggerAdapter:
    """Return a logger configured for structured JSON output."""

    if not logging.getLogger().handlers:
        configure_logging()
    return StructuredLoggerAdapter(logging.getLogger(name), extra={})
