import logging
from datetime import datetime, timezone
from pythonjsonlogger import jsonlogger
from .redactor import PIIRedactor


class ContextLoggerAdapter(logging.LoggerAdapter):
    def process(self, msg, kwargs):
        call_extra = kwargs.pop("extra", {}) or {}
        merged = dict(self.extra or {})
        merged.update(call_extra)
        kwargs["extra"] = merged
        return msg, kwargs


class JsonFormatter(jsonlogger.JsonFormatter):
    SAFE_RESERVED_FIELDS = {"timestamp", "level", "name", "message", "traceback"}

    def add_fields(self, log_record, record, message_dict):
        super().add_fields(log_record, record, message_dict)

        # Use the top-level package name of the caller, not this library
        log_record["name"] = record.name.split(".")[0]

        if not log_record.get("timestamp"):
            log_record["timestamp"] = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%fZ")

        log_record["level"] = (log_record.get("level") or record.levelname).upper()

        for key in list(log_record.keys()):
            if key not in self.SAFE_RESERVED_FIELDS:
                log_record[key] = PIIRedactor.redact(log_record[key], key=key)

        if "message" in log_record:
            log_record["message"] = PIIRedactor.redact_text(log_record["message"])
        if "traceback" in log_record:
            log_record["traceback"] = PIIRedactor.redact_text(log_record["traceback"])
            