from __future__ import annotations

import contextvars
import json
import logging
from datetime import datetime, timezone

from celery import current_task

_request_id: contextvars.ContextVar[str] = contextvars.ContextVar('request_id', default='')


def set_request_id(value: str) -> None:
    _request_id.set(value or '')


def get_request_id() -> str:
    return _request_id.get('')


def clear_request_id() -> None:
    _request_id.set('')


class RequestContextFilter(logging.Filter):
    def filter(self, record: logging.LogRecord) -> bool:
        record.request_id = get_request_id()
        task = current_task
        request = getattr(task, 'request', None) if task else None
        record.task_id = getattr(request, 'id', '') if request else ''
        record.task_name = getattr(task, 'name', '') if task else ''
        return True


class JsonFormatter(logging.Formatter):
    def format(self, record: logging.LogRecord) -> str:
        payload = {
            'timestamp': datetime.fromtimestamp(record.created, tz=timezone.utc).isoformat(),
            'level': record.levelname,
            'logger': record.name,
            'message': record.getMessage(),
        }
        request_id = getattr(record, 'request_id', '') or ''
        task_id = getattr(record, 'task_id', '') or ''
        task_name = getattr(record, 'task_name', '') or ''
        if request_id:
            payload['request_id'] = request_id
        if task_id:
            payload['task_id'] = task_id
        if task_name:
            payload['task_name'] = task_name
        if record.exc_info:
            payload['exception'] = self.formatException(record.exc_info)
        return json.dumps(payload, ensure_ascii=False)