import logging
from contextvars import ContextVar
from typing import Optional

# Глобальные счетчики
_connection_counter = 0
_message_counter = 0

# Контекстная переменная для хранения correlation ID
_correlation_id_var: ContextVar[Optional[str]] = ContextVar(
    "correlation_id", default=None
)

# Контекстная переменная для хранения ID подключения
_connection_id_var: ContextVar[Optional[int]] = ContextVar(
    "connection_id", default=None
)

# Контекстная переменная для хранения ID сообщения
_message_id_var: ContextVar[Optional[int]] = ContextVar(
    "message_id", default=None
)


def generate_connection_id() -> int:
    """Генерирует новый ID подключения (глобальный счетчик)."""
    global _connection_counter
    _connection_counter += 1
    return _connection_counter


def generate_message_id() -> int:
    """Генерирует новый ID сообщения (глобальный счетчик)."""
    global _message_counter
    _message_counter += 1
    return _message_counter


def set_connection_id(connection_id: int) -> None:
    """Устанавливает ID подключения."""
    _connection_id_var.set(connection_id)


def set_message_id(message_id: int) -> None:
    """Устанавливает ID сообщения."""
    _message_id_var.set(message_id)
    # Обновляем полный correlation ID
    _update_correlation_id()


def set_correlation_id(correlation_id: str) -> None:
    """Устанавливает полный correlation ID вручную."""
    _correlation_id_var.set(correlation_id)


def _update_correlation_id() -> None:
    """Обновляет полный correlation ID на основе connection_id и message_id."""
    connection_id = _connection_id_var.get()
    message_id = _message_id_var.get()

    if connection_id and message_id:
        correlation_id = f"CONN_{connection_id}+MSG_{message_id}"
    elif connection_id:
        correlation_id = f"CONN_{connection_id}"
    elif message_id:
        correlation_id = f"MSG_{message_id}"
    else:
        correlation_id = None

    if correlation_id:
        _correlation_id_var.set(correlation_id)


def get_correlation_id() -> Optional[str]:
    """Возвращает текущий correlation ID."""
    return _correlation_id_var.get()


def get_connection_id() -> Optional[int]:
    """Возвращает текущий ID подключения."""
    return _connection_id_var.get()


def get_message_id() -> Optional[int]:
    """Возвращает текущий ID сообщения."""
    return _message_id_var.get()


def clear_context() -> None:
    """Очищает все контекстные переменные."""
    _correlation_id_var.set(None)
    _connection_id_var.set(None)
    _message_id_var.set(None)


class CorrelationFilter(logging.Filter):
    """
    Фильтр логирования, который добавляет correlation_id в запись логирования.
    """

    def filter(self, record):
        correlation_id = get_correlation_id()
        record.correlation_id = correlation_id or "-"
        return True