from __future__ import annotations

from dataclasses import dataclass, field
from threading import Lock
import time
from typing import Callable

from flownet.v1.runtime.topics.normalization import _normalize_non_negative_int, _normalize_topic_uri


@dataclass
class EventGroupLedger:
    event_group_id: str
    event_chain_pending: int = 0
    producer_done: bool = False
    emitted_event_count: int = 0
    final_seq: int | None = None
    expected_total_events_hint: int | None = None
    request_done_emitted: bool = False
    updated_at: float = field(default_factory=time.time)


@dataclass
class CoordinatorTopicState:
    topic_uri: str
    epoch: int
    consuming_flow_process_uris: set[str] = field(default_factory=set)
    producing_flow_process_uris: set[str] = field(default_factory=set)
    subscriber_ids: set[str] = field(default_factory=set)
    subscriber_addresses: dict[str, str] = field(default_factory=dict)
    event_group_ledgers: dict[str, EventGroupLedger] = field(default_factory=dict)
    updated_at: float = field(default_factory=time.time)


class TopicCoordinatorRegistry:
    """
    Local coordinator-state holder.

    State is lazy and keyed by (topic_uri, epoch).
    """

    def __init__(self, *, time_fn: Callable[[], float] | None = None):
        self._states: dict[tuple[str, int], CoordinatorTopicState] = {}
        self._lock = Lock()
        self._time_fn = time_fn or time.time

    def get_or_create(self, topic_uri: str, epoch: int) -> CoordinatorTopicState:
        normalized_topic_uri = _normalize_topic_uri(topic_uri)
        normalized_epoch = _normalize_non_negative_int(epoch, field_name="epoch")
        key = (normalized_topic_uri, normalized_epoch)
        now = self._time_fn()
        with self._lock:
            state = self._states.get(key)
            if state is None:
                state = CoordinatorTopicState(
                    topic_uri=normalized_topic_uri,
                    epoch=normalized_epoch,
                    updated_at=now,
                )
                self._states[key] = state
            else:
                state.updated_at = now
            return state

    def get(self, topic_uri: str, epoch: int) -> CoordinatorTopicState | None:
        normalized_topic_uri = _normalize_topic_uri(topic_uri)
        normalized_epoch = _normalize_non_negative_int(epoch, field_name="epoch")
        with self._lock:
            return self._states.get((normalized_topic_uri, normalized_epoch))

    def gc_idle(self, *, max_idle_seconds: float) -> int:
        threshold = max(0.0, float(max_idle_seconds))
        now = self._time_fn()
        removed = 0
        with self._lock:
            for key, state in list(self._states.items()):
                self._prune_finished_event_group_ledgers(
                    state=state,
                    now=now,
                    threshold=threshold,
                )
                if self._has_unfinished_event_group(state):
                    continue
                if state.event_group_ledgers:
                    # keep recently-finished ledgers for observability until TTL.
                    continue
                if (now - state.updated_at) <= threshold:
                    continue
                if state.consuming_flow_process_uris:
                    continue
                if state.producing_flow_process_uris:
                    continue
                if state.subscriber_ids:
                    continue
                if state.subscriber_addresses:
                    continue
                self._states.pop(key, None)
                removed += 1
        return removed

    @staticmethod
    def _has_unfinished_event_group(state: CoordinatorTopicState) -> bool:
        for ledger in state.event_group_ledgers.values():
            if not TopicCoordinatorRegistry._is_finished_event_group_ledger(ledger):
                return True
        return False

    @staticmethod
    def _is_finished_event_group_ledger(ledger: EventGroupLedger) -> bool:
        if not ledger.request_done_emitted:
            return False
        if not ledger.producer_done:
            return False
        if int(ledger.event_chain_pending) != 0:
            return False
        return True

    @staticmethod
    def _prune_finished_event_group_ledgers(
        *,
        state: CoordinatorTopicState,
        now: float,
        threshold: float,
    ) -> None:
        for event_group_id, ledger in list(state.event_group_ledgers.items()):
            if not TopicCoordinatorRegistry._is_finished_event_group_ledger(ledger):
                continue
            if (now - ledger.updated_at) <= threshold:
                continue
            state.event_group_ledgers.pop(event_group_id, None)


__all__ = [
    "EventGroupLedger",
    "CoordinatorTopicState",
    "TopicCoordinatorRegistry",
]
