"""Transactional Kafka producer for atomic multi-topic writes."""

from __future__ import annotations

import logging
import threading
import time
from typing import Optional

from confluent_kafka import KafkaException, Producer as ConfluentProducer

from neonlink.circuit_breaker import CircuitBreaker
from neonlink.config import Config
from neonlink.errors import PayloadTooLargeError, is_permanent
from neonlink.metrics import NoopMetrics
from neonlink.record import Record

logger = logging.getLogger(__name__)


class TransactionalProducer:
    """Publishes records atomically across one or more topics using Kafka transactions.

    Use this for fan-out flows that must be all-or-nothing
    (e.g., findata -> posting-input + sync-triggers).

    A TransactionalProducer requires ``Config.transactional_id`` to be set.
    The ID must be unique per producer instance but stable across restarts
    so Kafka can recover in-flight transactions. Typical pattern::

        {service}-{role}-{instance}   e.g. "findata-coordinator-0"
    """

    def __init__(self, cfg: Config) -> None:
        cfg.validate()
        if not cfg.transactional_id:
            raise ValueError("neonlink: transactional_id is required for transactional producer")

        self._producer = ConfluentProducer(cfg.to_transactional_config())
        self._producer.init_transactions()
        self._max_message_bytes = cfg.max_message_bytes
        self._metrics = cfg.metrics if cfg.metrics is not None else NoopMetrics()
        self._cb = CircuitBreaker(
            name="txn-producer",
            max_failures=cfg.cb_max_failures,
            reset_timeout_sec=cfg.cb_reset_timeout_sec,
            on_state_change=lambda n, f, t: self._metrics.circuit_breaker_state_change(n, f, t),
        )
        self._lock = threading.Lock()
        self._closed = False

    def publish_atomic(
        self,
        records: list[Record],
        *,
        cancel_event: Optional[threading.Event] = None,
        timeout_sec: float = 30.0,
    ) -> None:
        """Produce all records within a single Kafka transaction.

        Either all records are committed to their respective topics, or none are.
        Each record must have its ``topic`` set.

        Args:
            cancel_event: Optional threading.Event that, when set, aborts the transaction.
            timeout_sec: Maximum time to wait for delivery confirmation.

        Raises:
            RuntimeError: if the producer is closed or cancelled.
            KafkaException: on transaction failure.
            CircuitBreakerOpen: if the circuit breaker is open.
        """
        # Serialize publish paths through a single lock so one shared
        # transactional producer can safely handle normal fan-out and DLQ usage.
        with self._lock:
            if self._closed:
                raise RuntimeError("neonlink: transactional producer is closed")
            if cancel_event is not None and cancel_event.is_set():
                raise RuntimeError("neonlink: publish cancelled")
            if not records:
                return

            for i, rec in enumerate(records):
                if rec is None:
                    raise ValueError(f"neonlink: record[{i}] is None")
                if not rec.topic:
                    raise ValueError(f"neonlink: record[{i}] topic is required")

                record_size = _encoded_record_size(rec.key, rec.value, rec.headers)
                if self._max_message_bytes > 0 and record_size > self._max_message_bytes:
                    raise PayloadTooLargeError(
                        f"record {record_size} bytes exceeds limit {self._max_message_bytes}"
                    )

            def _do_transaction() -> None:
                self._producer.begin_transaction()

                errors: list[Exception] = []

                def _on_delivery(err, _msg):
                    if err is not None:
                        errors.append(KafkaException(err))

                try:
                    for rec in records:
                        kafka_headers = _headers_to_list(rec.headers) if rec.headers else None
                        self._producer.produce(
                            topic=rec.topic,
                            key=rec.key,
                            value=rec.value,
                            headers=kafka_headers,
                            on_delivery=_on_delivery,
                        )

                    # Flush to ensure all delivery callbacks fire.
                    self._producer.flush(timeout=timeout_sec)

                    if errors:
                        raise errors[0]

                    self._producer.commit_transaction()

                except Exception as exc:
                    try:
                        self._producer.abort_transaction()
                    except Exception as abort_exc:
                        logger.error(
                            "neonlink: transaction abort failed: %s (original: %s)",
                            abort_exc,
                            exc,
                        )
                    raise

            txn_topic = records[0].topic if records else ""
            start = time.monotonic()
            try:
                self._cb.execute_with_classifier(
                    _do_transaction,
                    should_trip=lambda err: not is_permanent(err),
                )
            except Exception as exc:
                self._metrics.publish_error(txn_topic, exc)
                raise
            self._metrics.publish_success(txn_topic, time.monotonic() - start)

    def publish(
        self,
        topic: str,
        key: Optional[bytes],
        value: Optional[bytes],
        headers: Optional[dict[str, str | bytes]] = None,
        *,
        cancel_event: Optional[threading.Event] = None,
        timeout_sec: float = 30.0,
    ) -> None:
        """Publish a single record transactionally."""
        self.publish_atomic(
            [Record(topic=topic, key=key, value=value, headers=headers or {})],
            cancel_event=cancel_event,
            timeout_sec=timeout_sec,
        )

    def publish_record(
        self,
        record: Record,
        *,
        cancel_event: Optional[threading.Event] = None,
        timeout_sec: float = 30.0,
    ) -> None:
        """Publish one Record transactionally."""
        self.publish_atomic([record], cancel_event=cancel_event, timeout_sec=timeout_sec)

    def publish_batch(
        self,
        records: list[Record],
        *,
        cancel_event: Optional[threading.Event] = None,
        timeout_sec: float = 30.0,
    ) -> None:
        """Publish multiple records in one transaction."""
        self.publish_atomic(records, cancel_event=cancel_event, timeout_sec=timeout_sec)

    def close(self) -> None:
        """Flush pending records and close the transactional producer."""
        with self._lock:
            if self._closed:
                return
            self._closed = True

        remaining = self._producer.flush(timeout=10.0)
        if remaining > 0:
            logger.warning(
                "neonlink: %d messages not delivered on transactional producer close",
                remaining,
            )


_HEADER_FRAMING_OVERHEAD = 6
_PROTOCOL_OVERHEAD = 128


def _encoded_record_size(key, value, headers) -> 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()]
