"""Kafka consumer with manual offset commit, failure tracking, and panic recovery."""

from __future__ import annotations

import logging
import math
import random
import threading
import time
from datetime import datetime, timezone
from typing import Callable, Optional, Protocol

from confluent_kafka import (
    Consumer as ConfluentConsumer,
    KafkaError,
    KafkaException,
    Message,
    TopicPartition,
    TIMESTAMP_CREATE_TIME,
    TIMESTAMP_LOG_APPEND_TIME,
)

from neonlink.circuit_breaker import CircuitBreaker
from neonlink.config import Config
from neonlink.errors import PermanentError, is_permanent
from neonlink.metrics import NoopMetrics
from neonlink.record import HEADER_RETRY_COUNT, Record
from neonlink.tracing import end_span_with_error, extract_trace_context, start_consumer_span

try:
    from opentelemetry import context as otel_context
    _HAS_OTEL = True
except ImportError:
    _HAS_OTEL = False

logger = logging.getLogger(__name__)


class MessageHandler(Protocol):
    """Protocol for message handlers."""

    def handle_message(self, record: Record) -> None:
        """Process a single record.

        Return normally to acknowledge (commit offset).
        Raise an exception to skip commit — Kafka redelivers on next poll.
        """
        ...


class MessageHandlerWithCommit(Protocol):
    """Protocol for handlers that manually control Kafka offset commits."""

    def handle_message_with_commit(
        self,
        record: Record,
        commit: Callable[[], bool],
    ) -> None:
        """Process a record and call commit() only after durable side-effects."""
        ...


class _FailureTracker:
    """In-memory per-record failure counter.

    Fixes the broken retry_count header problem: Kafka redelivers the same
    record with the same headers, so the retry_count header never advances
    unless the producer re-publishes. This tracker maintains a local count
    keyed by topic:partition:offset.
    """

    def __init__(self) -> None:
        self._counts: dict[str, int] = {}
        self._lock = threading.Lock()

    @staticmethod
    def _key(topic: str, partition: int, offset: int) -> str:
        return f"{topic}:{partition}:{offset}"

    def increment(self, topic: str, partition: int, offset: int) -> int:
        with self._lock:
            k = self._key(topic, partition, offset)
            self._counts[k] = self._counts.get(k, 0) + 1
            return self._counts[k]

    def remove(self, topic: str, partition: int, offset: int) -> None:
        with self._lock:
            self._counts.pop(self._key(topic, partition, offset), None)

    def clear_partition(self, topic: str, partition: int) -> None:
        with self._lock:
            prefix = f"{topic}:{partition}:"
            to_delete = [k for k in self._counts if k.startswith(prefix)]
            for k in to_delete:
                del self._counts[k]


def _consumer_poll_backoff(consecutive_errors: int, initial_wait_ms: int, max_wait_ms: int) -> float:
    """Calculate backoff in seconds for consumer poll retries."""
    if consecutive_errors <= 0:
        return 0.0
    base = initial_wait_ms * math.pow(2, consecutive_errors - 1)
    if base > max_wait_ms:
        base = max_wait_ms
    # Jitter: 75-125% of base.
    jittered = base * (0.75 + random.random() * 0.5)
    return jittered / 1000.0  # Convert to seconds


class Consumer:
    """Subscribes to Kafka topics and dispatches records to a MessageHandler.

    On handler success: offset is committed.
    On handler error: offset is NOT committed. Kafka redelivers on next poll.
    If retry_count >= poison pill threshold: offset is committed (skip past).

    Features:
    - In-memory failure tracking (fixes broken retry_count header)
    - Exponential backoff on poll errors (prevents hot-spin)
    - Handler panic recovery (consumer continues after handler exceptions)
    - Rebalance-aware failure tracker cleanup
    """

    def __init__(self, cfg: Config, topics: list[str]) -> None:
        cfg.validate()
        if not cfg.consumer_group:
            raise ValueError("neonlink: consumer group is required")
        if not topics:
            raise ValueError("neonlink: at least one topic is required")

        self._consumer = ConfluentConsumer(cfg.to_consumer_config())
        self._topics = topics
        self._cfg = cfg
        self._tracker = _FailureTracker()
        self._metrics = cfg.metrics if cfg.metrics is not None else NoopMetrics()
        self._cb = CircuitBreaker(
            name="consumer",
            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._stop_event = threading.Event()
        self._cancel_event: Optional[threading.Event] = None
        self._started = False
        self._closed = False
        self._last_poll_error = False
        self._dlq_producer = None
        self._revoked_partitions: set[str] = set()
        self._revoked_lock = threading.Lock()

    def subscribe(
        self,
        handler: MessageHandler,
        *,
        cancel_event: Optional[threading.Event] = None,
    ) -> None:
        """Start the consume loop, blocking until stop() is called.

        Args:
            handler: MessageHandler implementation to process each record.
            cancel_event: Optional threading.Event that, when set, stops the loop.
        """
        if self._started:
            raise RuntimeError("neonlink: consumer already started")
        self._started = True

        if cancel_event is not None:
            self._cancel_event = cancel_event

        def _on_revoke(consumer, partitions):
            with self._revoked_lock:
                for tp in partitions:
                    self._revoked_partitions.add(f"{tp.topic}:{tp.partition}")
                    self._tracker.clear_partition(tp.topic, tp.partition)
                    logger.info(
                        "neonlink: partition revoked topic=%s partition=%d",
                        tp.topic, tp.partition,
                    )
            # NOTE: Do NOT call consumer.commit() here. Per-message commits
            # (self._commit(msg)) are synchronous, so acknowledged offsets are
            # already persisted. A blanket commit() would advance the position
            # past records whose handlers failed, defeating the "don't commit
            # on error" guarantee. The Go SDK uses CommitMarkedOffsets which
            # only flushes explicitly marked offsets — this is the equivalent.

        def _on_assign(consumer, partitions):
            with self._revoked_lock:
                for tp in partitions:
                    self._revoked_partitions.discard(f"{tp.topic}:{tp.partition}")
                    logger.info(
                        "neonlink: partition assigned topic=%s partition=%d",
                        tp.topic, tp.partition,
                    )

        self._consumer.subscribe(
            self._topics,
            on_revoke=_on_revoke,
            on_assign=_on_assign,
        )
        logger.info(
            "neonlink: consumer starting, group=%s topics=%s",
            self._cfg.consumer_group,
            self._topics,
        )

        consecutive_errors = 0

        try:
            while not self._is_stopped():
                msg = self._poll_one()
                if msg is None:
                    # _poll_one returns None on error or timeout.
                    # Check if it was an error (CB or poll error).
                    if self._last_poll_error:
                        consecutive_errors += 1
                        backoff = _consumer_poll_backoff(
                            consecutive_errors,
                            self._cfg.retry_initial_wait_ms,
                            self._cfg.retry_max_wait_ms,
                        )
                        logger.warning(
                            "neonlink: poll error, backing off "
                            "consecutive_errors=%d backoff_sec=%.2f",
                            consecutive_errors, backoff,
                        )
                        if self._wait_for_stop(timeout=backoff):
                            break  # Stop requested during backoff
                        self._last_poll_error = False
                    continue

                consecutive_errors = 0
                record = _from_confluent_message(msg)

                # Rebalance fencing: skip records from revoked partitions.
                with self._revoked_lock:
                    if f"{record.topic}:{record.partition}" in self._revoked_partitions:
                        logger.debug(
                            "neonlink: skipping record from revoked partition "
                            "topic=%s partition=%s offset=%s",
                            record.topic, record.partition, record.offset,
                        )
                        continue

                # Consumer lag metrics.
                self._emit_consumer_lag(record)

                # Extract OTel trace context for distributed tracing continuity.
                otel_token = _attach_trace_context(record)

                # Start consumer span for observability.
                span = start_consumer_span(record)

                # Panic recovery: catch handler exceptions.
                handler_start = time.monotonic()
                handler_err = None
                commit_handler = getattr(handler, "handle_message_with_commit", None)
                has_commit_control = callable(commit_handler)
                try:
                    if has_commit_control:
                        commit_handler(record, lambda: self._commit(msg))
                    else:
                        handler.handle_message(record)
                except Exception as exc:
                    handler_err = exc
                    self._metrics.consume_error(record.topic, record.partition, exc)
                    self._handle_processing_error(record, msg, exc)
                finally:
                    end_span_with_error(span, handler_err)
                    _detach_trace_context(otel_token)

                if handler_err is not None:
                    continue

                # Success: commit and clear failure tracker.
                self._metrics.consume_success(
                    record.topic, record.partition, time.monotonic() - handler_start
                )
                self._tracker.remove(record.topic, record.partition, record.offset)
                if not has_commit_control:
                    self._commit(msg)
        finally:
            self._consumer.close()
            logger.info("neonlink: consumer stopped")

    def subscribe_with_dlq(
        self,
        handler: MessageHandler,
        dlq_producer,
        *,
        cancel_event: Optional[threading.Event] = None,
    ) -> None:
        """Start the consume loop with DLQ routing for poison pills.

        When a poison pill is detected, the record is routed to the DLQ before
        committing the offset. If DLQ routing fails, the offset is NOT committed.
        """
        self._dlq_producer = dlq_producer
        self.subscribe(handler, cancel_event=cancel_event)

    def stop(self) -> None:
        """Request graceful shutdown of the consume loop."""
        self._stop_event.set()
        if self._cancel_event is not None:
            self._cancel_event.set()

    def close(self) -> None:
        """Stop and close the consumer."""
        if self._closed:
            return
        self._closed = True
        self.stop()

    def pause_fetch(self, topics: list[str]) -> None:
        """Pause fetching from the specified topics for custom backpressure."""
        partitions = self._consumer.assignment()
        to_pause = [tp for tp in partitions if tp.topic in topics]
        if to_pause:
            self._consumer.pause(to_pause)
            logger.info("neonlink: fetch paused topics=%s", topics)

    def resume_fetch(self, topics: list[str]) -> None:
        """Resume fetching from the specified topics."""
        partitions = self._consumer.assignment()
        to_resume = [tp for tp in partitions if tp.topic in topics]
        if to_resume:
            self._consumer.resume(to_resume)
            logger.info("neonlink: fetch resumed topics=%s", topics)

    def _emit_consumer_lag(self, record: Record) -> None:
        """Emit consumer lag metric for the record's partition (best-effort)."""
        try:
            tp = TopicPartition(record.topic, record.partition)
            _, high = self._consumer.get_watermark_offsets(tp, cached=True)
            lag = high - record.offset - 1
            if lag < 0:
                lag = 0
            self._metrics.consumer_lag(record.topic, record.partition, lag)
        except Exception:
            pass

    def _poll_one(self) -> Optional[Message]:
        """Poll for a single message with circuit breaker."""
        msg: Optional[Message] = None
        self._last_poll_error = False

        def _do_poll() -> None:
            nonlocal msg
            msg = self._consumer.poll(timeout=1.0)
            if msg is None:
                return
            err = msg.error()
            if err is None:
                return
            if err.code() == KafkaError._PARTITION_EOF:
                msg = None
                return
            raise KafkaException(err)

        try:
            self._cb.execute_with_classifier(
                _do_poll,
                should_trip=lambda err: not is_permanent(err),
            )
        except Exception as exc:
            logger.warning("neonlink: poll error: %s", exc)
            self._last_poll_error = True
            return None

        if msg is None:
            return None

        return msg

    def _commit(self, msg: Message) -> bool:
        """Commit offset for the consumed message with bounded retry.

        Returns True on success, False if all retries exhausted.
        """
        topic = msg.topic() or ""
        partition = msg.partition() or 0
        if self._is_partition_revoked(topic, partition):
            logger.warning(
                "neonlink: commit fenced for revoked partition "
                "topic=%s partition=%s offset=%s",
                topic,
                partition,
                msg.offset(),
            )
            return False

        max_retries = 3
        for attempt in range(max_retries):
            if self._is_partition_revoked(topic, partition):
                logger.warning(
                    "neonlink: commit fenced during retry for revoked partition "
                    "topic=%s partition=%s offset=%s",
                    topic,
                    partition,
                    msg.offset(),
                )
                return False
            try:
                self._consumer.commit(message=msg, asynchronous=False)
                return True
            except KafkaException as exc:
                logger.warning(
                    "neonlink: offset commit attempt %d/%d failed "
                    "topic=%s partition=%s offset=%s: %s",
                    attempt + 1,
                    max_retries,
                    msg.topic(),
                    msg.partition(),
                    msg.offset(),
                    exc,
                )
        logger.error(
            "neonlink: offset commit failed after %d retries "
            "topic=%s partition=%s offset=%s",
            max_retries,
            msg.topic(),
            msg.partition(),
            msg.offset(),
        )
        return False

    def _handle_processing_error(
        self, record: Record, msg: Message, exc: Exception
    ) -> None:
        """Decide whether to skip (commit) or leave for redelivery."""
        header_retry_count = _get_retry_count(record)
        tracker_count = self._tracker.increment(record.topic, record.partition, record.offset)

        effective_retry_count = max(header_retry_count, tracker_count)

        if effective_retry_count >= self._cfg.poison_pill_threshold:
            self._metrics.poison_pill_detected(record.topic, record.partition, record.offset)
            logger.error(
                "neonlink: poison pill detected "
                "topic=%s partition=%s offset=%s effective_retry=%d "
                "header_retry=%d tracker=%d: %s",
                record.topic,
                record.partition,
                record.offset,
                effective_retry_count,
                header_retry_count,
                tracker_count,
                exc,
            )

            if self._dlq_producer is not None:
                from neonlink.dlq import route_to_dlq

                try:
                    route_to_dlq(
                        self._dlq_producer,
                        self._cfg.dlq_topic,
                        record,
                        exc,
                        poison_reason="MAX_RETRIES_EXCEEDED",
                        retry_count=effective_retry_count,
                        consumer_group=self._cfg.consumer_group,
                    )
                    self._metrics.dlq_routed(record.topic, exc)
                    logger.warning(
                        "neonlink: poison pill routed to DLQ, committing "
                        "topic=%s offset=%s",
                        record.topic,
                        record.offset,
                    )
                except Exception as dlq_exc:
                    logger.error(
                        "neonlink: DLQ routing failed, offset NOT committed "
                        "topic=%s offset=%s: %s",
                        record.topic,
                        record.offset,
                        dlq_exc,
                    )
                    # Do NOT commit — leave for redelivery so we don't lose data.
                    return
            else:
                # No DLQ configured — do NOT commit, log CRITICAL.
                self._metrics.poison_pill_blocked(record.topic, record.partition, record.offset)
                logger.error(
                    "neonlink: CRITICAL poison pill with no DLQ, offset NOT committed "
                    "topic=%s partition=%s offset=%s",
                    record.topic,
                    record.partition,
                    record.offset,
                )
                return

            if self._commit(msg):
                self._tracker.remove(record.topic, record.partition, record.offset)
            else:
                logger.error(
                    "neonlink: poison pill commit failed after DLQ route "
                    "topic=%s partition=%s offset=%s",
                    record.topic,
                    record.partition,
                    record.offset,
                )
            return

        logger.warning(
            "neonlink: handler error, will redeliver "
            "topic=%s partition=%s offset=%s effective_retry=%d: %s",
            record.topic,
            record.partition,
            record.offset,
            effective_retry_count,
            exc,
        )

    def _is_stopped(self) -> bool:
        """Check if either the internal stop event or external cancel event is set."""
        if self._stop_event.is_set():
            return True
        if self._cancel_event is not None and self._cancel_event.is_set():
            return True
        return False

    def _wait_for_stop(self, timeout: float) -> bool:
        """Wait up to *timeout* seconds for a stop signal.

        Returns True if a stop was requested (i.e. the loop should break).
        Because threading.Event.wait() can only watch a single event, we poll
        in short intervals so we notice either event promptly.
        """
        deadline = time.monotonic() + timeout
        while True:
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                return self._is_stopped()
            # Poll in 0.1 s slices so we notice cancel_event quickly.
            wait_slice = min(remaining, 0.1)
            if self._stop_event.wait(timeout=wait_slice):
                return True
            if self._cancel_event is not None and self._cancel_event.is_set():
                return True
        return False

    def _is_partition_revoked(self, topic: str, partition: int) -> bool:
        with self._revoked_lock:
            return f"{topic}:{partition}" in self._revoked_partitions


def _attach_trace_context(record: Record):
    """Attach OTel trace context from record headers. Returns token or None."""
    if not _HAS_OTEL:
        return None
    trace_ctx = extract_trace_context(record)
    if trace_ctx is not None:
        return otel_context.attach(trace_ctx)
    return None


def _detach_trace_context(token) -> None:
    """Detach OTel trace context if previously attached."""
    if token is not None and _HAS_OTEL:
        otel_context.detach(token)


def _from_confluent_message(msg: Message) -> Record:
    """Convert a confluent-kafka Message to a Record."""
    headers: dict[str, str | bytes] = {}
    raw_headers = msg.headers()
    if raw_headers:
        for key, val in raw_headers:
            if val is not None:
                headers[key] = val  # store raw bytes, preserving binary metadata

    timestamp = None
    ts_type, ts_value = msg.timestamp()
    if ts_type in (TIMESTAMP_CREATE_TIME, TIMESTAMP_LOG_APPEND_TIME) and ts_value > 0:
        timestamp = datetime.fromtimestamp(ts_value / 1000.0, tz=timezone.utc)

    return Record(
        topic=msg.topic() or "",
        key=msg.key(),
        value=msg.value(),
        headers=headers,
        partition=msg.partition() or 0,
        offset=msg.offset() or 0,
        timestamp=timestamp,
    )


def _get_retry_count(record: Record) -> int:
    """Extract retry_count from record headers."""
    val = record.get_header(HEADER_RETRY_COUNT)
    if not val:
        return 0
    try:
        return int(val)
    except ValueError:
        return 0
