from datetime import UTC, datetime
from logging import Formatter, LogRecord
from typing import Any, Self

import orjson

from snakestack.logging.encoders import safe_jsonable_encoder


class JsonFormatter(Formatter):
    def format(
        self: Self,
        record: LogRecord,
        *args: Any,
        **kwargs: Any
    ) -> str:
        Formatter.format(self, record)
        return orjson.dumps(
            self._create_log(record),
            default=safe_jsonable_encoder
        ).decode("utf-8")

    def _create_log(self: Self, record: LogRecord) -> dict[str, Any]:
        message: dict[str, Any] = {
            "time": datetime.fromtimestamp(record.created, tz=UTC).isoformat(),
            "level": record.levelname,
            "pid": record.process,
            "name": f"{record.name}:{record.lineno}",
            "message": record.getMessage()
        }

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

        if hasattr(record, "request_id") and record.request_id:
            message.update({
                "request": {
                    "id": record.request_id
                }
            })

        if hasattr(record, "otelSpanID") and hasattr(record, "otelTraceID"):
            message.update({
                "open_telemetry": {
                    "trace_id": record.otelTraceID,
                    "span_id": record.otelSpanID
                }
            })

        if hasattr(record, "extra_data") and record.extra_data:
            message.update({
                "extra": record.extra_data
            })

        return message
