"""Record type wrapping a Kafka message with typed header accessors."""

from __future__ import annotations

from dataclasses import dataclass, field
from datetime import datetime
from typing import Optional


# Header key constants.
HEADER_MESSAGE_TYPE = "message_type"
HEADER_SOURCE_SERVICE = "source_service"
HEADER_TARGET_SERVICE = "target_service"
HEADER_MESSAGE_ID = "message_id"
HEADER_CORRELATION_ID = "correlation_id"
HEADER_IDENTITY_CTX = "identity_context"
HEADER_RETRY_COUNT = "retry_count"
HEADER_TRACEPARENT = "traceparent"
HEADER_TRACESTATE = "tracestate"

Headers = dict[str, str | bytes]


@dataclass
class Record:
    """A message consumed from or produced to a broker topic."""

    topic: str = ""
    key: Optional[bytes] = None
    value: Optional[bytes] = None
    headers: Headers = field(default_factory=dict)
    partition: int = 0
    offset: int = 0
    timestamp: Optional[datetime] = None

    def get_header(self, key: str) -> str:
        """Return header value as string, or empty string if not found.

        Bytes values are decoded as UTF-8 with replacement characters for
        backward compatibility with callers that expect str.
        """
        val = self.headers.get(key)
        if val is None:
            return ""
        if isinstance(val, bytes):
            return val.decode("utf-8", errors="replace")
        return val

    def get_header_bytes(self, key: str) -> bytes | None:
        """Return header value as raw bytes, or None if not found."""
        val = self.headers.get(key)
        if val is None:
            return None
        if isinstance(val, str):
            return val.encode("utf-8")
        return val

    def set_header(self, key: str, value: str) -> None:
        """Set a header key-value pair (str, for backward compatibility)."""
        self.headers[key] = value

    def set_header_bytes(self, key: str, value: bytes) -> None:
        """Set a header key-value pair with raw bytes."""
        self.headers[key] = value

    @property
    def message_type(self) -> str:
        return self.get_header(HEADER_MESSAGE_TYPE)

    @property
    def correlation_id(self) -> str:
        return self.get_header(HEADER_CORRELATION_ID)

    @property
    def message_id(self) -> str:
        return self.get_header(HEADER_MESSAGE_ID)


# ---------------------------------------------------------------------------
# Wire-size estimation (shared by producer and transactional producer)
# ---------------------------------------------------------------------------

HEADER_FRAMING_OVERHEAD = 6
"""Per-header framing cost in the Kafka wire protocol."""

PROTOCOL_OVERHEAD = 128
"""Conservative estimate of Kafka record protocol overhead."""


def encoded_record_size(
    key: Optional[bytes],
    value: Optional[bytes],
    headers: Optional[dict[str, str | bytes]],
) -> int:
    """Estimate the total wire size of a Kafka record."""
    size = (len(key) if key else 0) + (len(value) if value else 0) + PROTOCOL_OVERHEAD
    if headers:
        for k, v in headers.items():
            size += len(k) + (len(v) if v else 0) + HEADER_FRAMING_OVERHEAD
    return size


def headers_to_list(headers: dict[str, str | bytes]) -> list[tuple[str, bytes]]:
    """Convert header dict to confluent-kafka format."""
    return [(k, v if isinstance(v, bytes) else v.encode("utf-8")) for k, v in headers.items()]
