from __future__ import annotations

from typing import Any, Callable

from flownet.v1.runtime.topics.coordinator_registry import (
    CoordinatorTopicState,
    EventGroupLedger,
)
from flownet.v1.runtime.topics.normalization import (
    _normalize_non_empty,
    _normalize_non_negative_int,
)


class EventGroupLedgerManager:
    """Coordinator-side convergence ledger manager keyed by event_group_id."""

    def __init__(
        self,
        *,
        time_fn: Callable[[], float],
        on_request_done: Callable[[dict[str, Any]], None] | None = None,
    ) -> None:
        self._time_fn = time_fn
        self._on_request_done = on_request_done

    def normalize_event_group_id(self, event_group_id: str) -> str:
        return _normalize_non_empty(event_group_id, field_name="event_group_id")

    def ledger(
        self,
        *,
        state: CoordinatorTopicState,
        event_group_id: str,
    ) -> EventGroupLedger:
        normalized_event_group_id = self.normalize_event_group_id(event_group_id)
        ledger = state.event_group_ledgers.get(normalized_event_group_id)
        if ledger is None:
            ledger = EventGroupLedger(event_group_id=normalized_event_group_id)
            state.event_group_ledgers[normalized_event_group_id] = ledger
        return ledger

    def snapshot(
        self,
        *,
        state: CoordinatorTopicState,
        event_group_id: str,
    ) -> dict[str, Any] | None:
        normalized_event_group_id = self.normalize_event_group_id(event_group_id)
        ledger = state.event_group_ledgers.get(normalized_event_group_id)
        if ledger is None:
            return None
        return {
            "event_group_id": ledger.event_group_id,
            "event_chain_pending": ledger.event_chain_pending,
            "producer_done": ledger.producer_done,
            "emitted_event_count": ledger.emitted_event_count,
            "final_seq": ledger.final_seq,
            "expected_total_events_hint": ledger.expected_total_events_hint,
            "request_done_emitted": ledger.request_done_emitted,
        }

    def record_local_event_publish(
        self,
        *,
        state: CoordinatorTopicState,
        event_group_id: str,
        payload: Any,
        tags: dict[str, str] | None,
        seq: int | None,
    ) -> dict[str, Any]:
        ledger = self.ledger(state=state, event_group_id=event_group_id)
        normalized_seq = self._resolve_event_seq(tags=tags, seq=seq, ledger=ledger)
        ledger.emitted_event_count += 1
        if normalized_seq is not None:
            if ledger.final_seq is None or normalized_seq > ledger.final_seq:
                ledger.final_seq = normalized_seq
        ledger.updated_at = self._time_fn()
        state.updated_at = ledger.updated_at
        return {
            "topic_uri": state.topic_uri,
            "epoch": state.epoch,
            "event_group_id": ledger.event_group_id,
            "payload": payload,
            "tags": dict(tags or {}),
            "seq": normalized_seq,
            "emitted_event_count": ledger.emitted_event_count,
        }

    def apply_event_chain_delta(
        self,
        *,
        state: CoordinatorTopicState,
        event_group_id: str,
        delta: int,
    ) -> dict[str, Any]:
        ledger = self.ledger(state=state, event_group_id=event_group_id)
        next_pending = ledger.event_chain_pending + int(delta)
        if next_pending < 0:
            raise ValueError("event_chain_pending_underflow")
        ledger.event_chain_pending = next_pending
        ledger.updated_at = self._time_fn()
        state.updated_at = ledger.updated_at
        request_done = self._maybe_request_done(state=state, ledger=ledger)
        return {
            "topic_uri": state.topic_uri,
            "epoch": state.epoch,
            "event_group_id": ledger.event_group_id,
            "event_chain_pending": ledger.event_chain_pending,
            "producer_done": ledger.producer_done,
            "request_done": request_done,
        }

    def apply_producer_done(
        self,
        *,
        state: CoordinatorTopicState,
        event_group_id: str,
        final_seq: int | None,
        expected_total_events: int | None,
    ) -> dict[str, Any]:
        ledger = self.ledger(state=state, event_group_id=event_group_id)
        ledger.producer_done = True
        if final_seq is not None:
            normalized_final_seq = _normalize_non_negative_int(final_seq, field_name="final_seq")
            if ledger.final_seq is None or normalized_final_seq > ledger.final_seq:
                ledger.final_seq = normalized_final_seq
        if expected_total_events is not None:
            ledger.expected_total_events_hint = _normalize_non_negative_int(
                expected_total_events,
                field_name="expected_total_events",
            )
        ledger.updated_at = self._time_fn()
        state.updated_at = ledger.updated_at
        request_done = self._maybe_request_done(state=state, ledger=ledger)
        return {
            "topic_uri": state.topic_uri,
            "epoch": state.epoch,
            "event_group_id": ledger.event_group_id,
            "event_chain_pending": ledger.event_chain_pending,
            "producer_done": ledger.producer_done,
            "request_done": request_done,
        }

    def _resolve_event_seq(
        self,
        *,
        tags: dict[str, str] | None,
        seq: int | None,
        ledger: EventGroupLedger,
    ) -> int | None:
        if seq is None and isinstance(tags, dict):
            raw_seq = tags.get("__seq__")
            if raw_seq is not None:
                seq = int(raw_seq)
        if seq is not None:
            return _normalize_non_negative_int(seq, field_name="seq")
        return ledger.emitted_event_count + 1

    def _maybe_request_done(
        self,
        *,
        state: CoordinatorTopicState,
        ledger: EventGroupLedger,
    ) -> dict[str, Any] | None:
        if ledger.request_done_emitted:
            return None
        if not ledger.producer_done:
            return None
        if ledger.event_chain_pending != 0:
            return None
        ledger.request_done_emitted = True
        expected_total_events = ledger.expected_total_events_hint
        if expected_total_events is None:
            if ledger.final_seq is not None:
                expected_total_events = ledger.final_seq
            else:
                expected_total_events = ledger.emitted_event_count
        request_done = {
            "topic_uri": state.topic_uri,
            "epoch": state.epoch,
            "event_group_id": ledger.event_group_id,
            "final_seq": ledger.final_seq,
            "expected_total_events": expected_total_events,
            "status": "completed",
        }
        if callable(self._on_request_done):
            self._on_request_done(dict(request_done))
        return request_done


__all__ = ["EventGroupLedgerManager"]
