"""
Updates server for streaming graph changes to clients.

Provides gRPC bidirectional streaming for graph updates with mTLS authentication.

RPC methods:
- GetState: Returns the current cached graph state and sequence position
- StreamUpdates: Bidirectional stream for updates and acknowledgments

When all connected clients have acknowledged an update, the cache is refreshed
and old ChangeLog entries are pruned.
"""

import asyncio
import os
import threading
import time
from typing import Any, AsyncIterator, Iterator, Optional
from uuid import uuid4

import dotenv
import grpc
import neo4j
import uvicorn
from fastapi import FastAPI
from fire import Fire  # type: ignore[import-untyped]
from observability_common.certificates import load_tls_credentials
from observability_common.grpc_interceptors import (
    create_service_interceptor,
    disable_rbac,
    require_role,
)
from observability_common.proto import observability_pb2 as pb
from observability_common.proto import observability_pb2_grpc
from pydantic import BaseModel

from .triggers import install_triggers
from .utils import event_from_dict

dotenv.load_dotenv()


class ClientState(BaseModel):
    """Tracks a connected client's position in the update stream."""

    client_id: str
    last_acked_position: int = 0
    last_ack_time: float = 0.0
    connected: bool = True


class GraphCache(BaseModel):
    """Cached graph state with associated metadata."""

    nodes: list[dict[str, Any]] = []
    edges: list[dict[str, Any]] = []
    cached_at_sequence: int = 0


class UpdatesManager:
    """Manages graph state caching, client tracking, and changelog pruning."""

    def __init__(
        self,
        neo4j_uri: str,
        neo4j_auth: tuple[str, str],
        database_name: str,
        poll_interval: float = 5.0,
        client_ack_timeout: float = 60.0,
        notify_debounce: float = 0.01,
    ):
        self.uri = neo4j_uri
        self.auth = neo4j_auth
        self.db_name = database_name
        self.poll_interval = poll_interval
        self.client_ack_timeout = client_ack_timeout

        self._cache = GraphCache()
        self._clients: dict[str, ClientState] = {}
        self._lock = threading.Lock()

        # Async signaling - Condition properly wakes all waiters
        self._condition: asyncio.Condition | None = None
        self._running = False
        self._last_notification_time: float = 0.0
        self._min_available_sequence: int = 0  # Oldest sequence still in changelog

        # Shared updates buffer - fetched once, read by all clients
        self._updates_buffer: list[dict[str, Any]] = []
        self._buffer_start_sequence: int = 0  # Sequence before first buffer item

        # Debounce notifications to batch rapid changes
        self._notify_debounce_delay: float = notify_debounce
        self._notify_pending: bool = False
        self._notify_task: Optional[asyncio.Task[None]] = None

    def _get_condition(self) -> asyncio.Condition:
        """Lazily create condition in the running event loop."""
        if self._condition is None:
            self._condition = asyncio.Condition()
        return self._condition

    def _get_driver(self) -> neo4j.Driver:
        return neo4j.GraphDatabase.driver(self.uri, auth=self.auth)

    def get_current_sequence(self) -> int:
        """Get the current sequence number from Neo4j."""
        with self._get_driver() as driver:
            result = driver.execute_query(
                """
                MATCH (s:ChangeLogSequence {id: 'main'})
                RETURN s.counter AS seq
                """,
                database_=self.db_name,
            )
            if result.records:
                return result.records[0]["seq"] or 0
            return 0

    def refresh_cache(self) -> None:
        """Refresh the cached graph state from Neo4j atomically."""
        with self._get_driver() as driver:
            # Get nodes, edges, and sequence in a single atomic query
            result = driver.execute_query(
                """
                CALL () {
                    OPTIONAL MATCH (s:ChangeLogSequence {id: 'main'})
                    RETURN coalesce(s.counter, 0) AS seq
                }
                CALL () {
                    MATCH (m:Message)
                    RETURN collect({
                        id: m.id,
                        content: m.content,
                        role: m.role,
                        agent: m.agent,
                        tools: coalesce(m.tools, "[]"),
                        derived_from: m.derived_from,
                        created_at: toString(m.created_at)
                    }) AS nodes
                }
                CALL () {
                    MATCH (src:Message)-[r:DEPENDS_ON]->(dst:Message)
                    RETURN collect({
                        source: src.id,
                        destination: dst.id,
                        message_index: r.message_index,
                        proximal: r.proximal
                    }) AS edges
                }
                RETURN seq, nodes, edges
                """,
                database_=self.db_name,
            )

            record = result.records[0] if result.records else None
            current_seq = record["seq"] if record else 0
            nodes = record["nodes"] if record else []
            edges = record["edges"] if record else []

            with self._lock:
                # On first load, initialize buffer/min sequence to current
                # so we don't fetch old changelog entries
                first_load = self._buffer_start_sequence == 0
                if first_load:
                    self._buffer_start_sequence = current_seq
                    self._min_available_sequence = current_seq + 1

                self._cache = GraphCache(
                    nodes=[dict(n) for n in nodes],
                    edges=[dict(e) for e in edges],
                    cached_at_sequence=current_seq,
                )

            # On first load, prune old changelog entries
            if first_load and current_seq > 0:
                driver.execute_query(
                    """
                    MATCH (c:ChangeLog)
                    WHERE c.sequence <= $seq
                    DELETE c
                    """,
                    database_=self.db_name,
                    seq=current_seq,
                )

    def get_cached_state(self) -> dict[str, Any]:
        """Returns the current cached graph state."""
        with self._lock:
            return {
                "nodes": self._cache.nodes,
                "edges": self._cache.edges,
                "sequence": self._cache.cached_at_sequence,
            }

    def needs_full_refresh(self, from_sequence: int) -> bool:
        """
        Check if a client at from_sequence needs a full state refresh.

        Returns True if changelog entries from from_sequence have been pruned.

        Client needs entries > from_sequence. If min_available = 51, then:
        - from_sequence = 50: needs 51+, available → OK
        - from_sequence = 49: needs 50+, but 50 pruned → needs refresh
        """
        with self._lock:
            # Client needs refresh if the next entry they need was pruned
            return from_sequence < self._min_available_sequence - 1

    def _fetch_updates_from_db(self, from_sequence: int) -> list[dict[str, Any]]:
        """Fetch ChangeLog entries from Neo4j."""
        with self._get_driver() as driver:
            result = driver.execute_query(
                """
                MATCH (c:ChangeLog)
                WHERE c.sequence > $from_seq
                RETURN c.sequence AS sequence,
                       c.type AS type,
                       c.id AS id,
                       c.content AS content,
                       c.role AS role,
                       c.agent AS agent,
                       coalesce(c.tools, "[]") AS tools,
                       c.derived_from AS derived_from,
                       c.source_id AS source_id,
                       c.destination_id AS destination_id,
                       c.message_index AS message_index,
                       c.proximal AS proximal,
                       toString(c.timestamp) AS timestamp,
                       // Computation (OTel span) fields
                       c.span_id AS span_id,
                       c.trace_id AS trace_id,
                       c.parent_span_id AS parent_span_id,
                       c.name AS name,
                       c.start_time_ns AS start_time_ns,
                       c.end_time_ns AS end_time_ns,
                       c.duration_ns AS duration_ns,
                       c.status_code AS status_code,
                       c.status_message AS status_message,
                       coalesce(c.attributes, "{}") AS attributes,
                       coalesce(c.events, "[]") AS events,
                       c.service_name AS service_name,
                       // Computation edge fields
                       c.child_span_id AS child_span_id,
                       c.edge_type AS edge_type
                ORDER BY c.sequence ASC
                """,
                database_=self.db_name,
                from_seq=from_sequence,
            )
            return [dict(r) for r in result.records]

    def get_updates_since(self, from_sequence: int) -> list[dict[str, Any]]:
        """Get updates from the shared buffer since the given sequence number."""
        with self._lock:
            if from_sequence < self._buffer_start_sequence:
                # Client is too far behind buffer - they need full refresh
                return []

            # Find updates in buffer after from_sequence
            result = []
            for update in self._updates_buffer:
                if update["sequence"] > from_sequence:
                    result.append(update)
            return result

    def _refresh_buffer(self) -> None:
        """Fetch new updates from DB and append to shared buffer."""
        with self._lock:
            # Determine where to fetch from
            if self._updates_buffer:
                last_seq = self._updates_buffer[-1]["sequence"]
            else:
                last_seq = self._buffer_start_sequence

        # Fetch new updates from DB
        new_updates = self._fetch_updates_from_db(last_seq)

        with self._lock:
            self._updates_buffer.extend(new_updates)

    def register_client(self, client_id: Optional[str] = None) -> str:
        """Register a new client and return its ID."""
        cid = client_id or str(uuid4())
        with self._lock:
            self._clients[cid] = ClientState(
                client_id=cid,
                last_acked_position=self._cache.cached_at_sequence,
                last_ack_time=time.time(),
                connected=True,
            )
        return cid

    def unregister_client(self, client_id: str) -> None:
        """Unregister a client when it disconnects."""
        with self._lock:
            if client_id in self._clients:
                self._clients[client_id].connected = False
                del self._clients[client_id]
        self._maybe_prune_changelog()

    def is_client_registered(self, client_id: str) -> bool:
        """Check if a client is still registered."""
        with self._lock:
            return client_id in self._clients

    def _cleanup_stale_clients(self) -> None:
        """Remove clients that haven't acknowledged within the timeout period."""
        now = time.time()
        stale_clients: list[str] = []

        with self._lock:
            current_seq = (
                self._updates_buffer[-1]["sequence"]
                if self._updates_buffer
                else self._buffer_start_sequence
            )

            for client_id, state in self._clients.items():
                # Client is stale if:
                # 1. They're behind the current sequence (have pending updates)
                # 2. They haven't acknowledged within the timeout
                if state.last_acked_position < current_seq:
                    time_since_ack = now - state.last_ack_time
                    if time_since_ack > self.client_ack_timeout:
                        stale_clients.append(client_id)

        # Unregister stale clients outside the lock
        for client_id in stale_clients:
            self.unregister_client(client_id)

    def acknowledge_position(self, client_id: str, position: int) -> None:
        """Client acknowledges it has received updates up to position."""
        with self._lock:
            if client_id in self._clients:
                self._clients[client_id].last_acked_position = max(
                    self._clients[client_id].last_acked_position, position
                )
                self._clients[client_id].last_ack_time = time.time()
        self._maybe_prune_changelog()

    def _apply_changelog_to_cache(
        self, changes: list[dict[str, Any]], up_to_seq: int
    ) -> None:
        """Apply changelog entries to the cached state."""
        with self._lock:
            nodes_by_id = {n["id"]: n for n in self._cache.nodes}
            edges_set = {
                (e["source"], e["destination"]) for e in self._cache.edges
            }
            edges_by_key = {
                (e["source"], e["destination"]): e for e in self._cache.edges
            }

            for change in changes:
                change_type = change["type"]

                if change_type == "node_created":
                    node_id = change.get("node_id") or change.get("id")
                    if node_id and node_id not in nodes_by_id:
                        nodes_by_id[node_id] = {
                            "id": node_id,
                            "content": change.get("node_content") or change.get("content"),
                            "role": change.get("node_role") or change.get("role"),
                            "agent": change.get("node_agent") or change.get("agent"),
                            "tools": change.get("node_tools") or change.get("tools"),
                            "derived_from": change.get("node_derived_from") or change.get("derived_from"),
                            "created_at": change.get("timestamp"),
                        }

                elif change_type == "node_deleted":
                    node_id = change["node_id"]
                    nodes_by_id.pop(node_id, None)

                elif change_type == "edge_created":
                    key = (change["source_id"], change["destination_id"])
                    if key not in edges_set:
                        edges_set.add(key)
                        edges_by_key[key] = {
                            "source": change["source_id"],
                            "destination": change["destination_id"],
                            "message_index": change.get("message_index"),
                            "proximal": change.get("proximal"),
                        }

                elif change_type == "edge_deleted":
                    key = (change["source_id"], change["destination_id"])
                    edges_set.discard(key)
                    edges_by_key.pop(key, None)

            self._cache = GraphCache(
                nodes=list(nodes_by_id.values()),
                edges=list(edges_by_key.values()),
                cached_at_sequence=up_to_seq,
            )

    def _maybe_prune_changelog(self) -> None:
        """
        If all connected clients have acknowledged past the cached sequence,
        update the cache and prune old changelog entries.

        Uses delta application when changelog is small relative to graph size,
        otherwise does a full refresh.
        """
        with self._lock:
            if not self._clients:
                return

            # Find the minimum acknowledged position among connected clients
            min_acked = min(
                c.last_acked_position for c in self._clients.values() if c.connected
            )

            # Only prune if all clients are past the cached sequence
            if min_acked <= self._cache.cached_at_sequence:
                return

            cached_seq = self._cache.cached_at_sequence
            graph_size = len(self._cache.nodes) + len(self._cache.edges)

        # Get changelog entries to apply
        changes = self.get_updates_since(cached_seq)
        num_changes = len(changes)

        # Heuristic: apply deltas if changes < 25% of graph size, else full refresh
        # Also do full refresh if graph is empty (initial load)
        if graph_size > 0 and num_changes < graph_size * 0.25:
            self._apply_changelog_to_cache(changes, min_acked)
        else:
            self.refresh_cache()

        # Prune old changelog entries and update min available sequence
        pruned_up_to = self._cache.cached_at_sequence
        with self._get_driver() as driver:
            driver.execute_query(
                """
                MATCH (c:ChangeLog)
                WHERE c.sequence <= $min_seq
                DELETE c
                """,
                database_=self.db_name,
                min_seq=pruned_up_to,
            )

        with self._lock:
            self._min_available_sequence = pruned_up_to + 1
            # Also prune the in-memory buffer
            self._updates_buffer = [
                u for u in self._updates_buffer if u["sequence"] > pruned_up_to
            ]
            self._buffer_start_sequence = pruned_up_to

    async def notify_update(self) -> None:
        """
        Called by the /notify endpoint when a trigger fires.
        Debounces rapid notifications to batch changes together.
        """
        self._notify_pending = True

        # If no debounce task is running, start one
        if self._notify_task is None or self._notify_task.done():
            self._notify_task = asyncio.create_task(self._debounced_notify())

    async def _debounced_notify(self) -> None:
        """Wait for debounce delay then process pending notifications."""
        await asyncio.sleep(self._notify_debounce_delay)

        if self._notify_pending:
            self._notify_pending = False
            self._last_notification_time = time.time()

            # Fetch new updates into shared buffer
            self._refresh_buffer()

            # Signal all waiting clients using Condition
            condition = self._get_condition()
            async with condition:
                condition.notify_all()

    async def poll_updates(self) -> None:
        """
        Background task that polls for updates as a fallback mechanism.
        Only polls if no notifications have been received recently.
        Also cleans up stale clients that haven't acknowledged.
        """
        self._running = True
        last_known_seq = self.get_current_sequence()

        while self._running:
            await asyncio.sleep(self.poll_interval)

            # Clean up clients that haven't acknowledged in time
            self._cleanup_stale_clients()

            # Skip polling if we received a notification recently
            time_since_notify = time.time() - self._last_notification_time
            if time_since_notify < self.poll_interval * 2:
                continue

            # Fallback: poll for updates if notifications aren't working
            current_seq = self.get_current_sequence()
            if current_seq > last_known_seq:
                last_known_seq = current_seq
                # Fetch updates and notify
                self._refresh_buffer()
                condition = self._get_condition()
                async with condition:
                    condition.notify_all()

    def stop(self) -> None:
        """Stop the polling task."""
        self._running = False

    async def wait_for_update(self, timeout: float = 30.0) -> bool:
        """Wait for a new update notification or timeout."""
        condition = self._get_condition()
        try:
            async with condition:
                await asyncio.wait_for(condition.wait(), timeout=timeout)
            return True
        except asyncio.TimeoutError:
            return False


# Global manager instance
_manager: Optional[UpdatesManager] = None


def get_manager() -> UpdatesManager:
    """Get the global UpdatesManager instance."""
    if _manager is None:
        raise RuntimeError("UpdatesManager not initialized")
    return _manager




def _dict_to_edge(d: dict) -> pb.Edge:
    """Convert a dict to an Edge protobuf message."""
    msg_idx = d.get("message_index")
    proximal = d.get("proximal")
    proximal_bool = proximal == "true" if isinstance(proximal, str) else bool(proximal)
    return pb.Edge(
        source=d.get("source", ""),
        destination=d.get("destination", ""),
        message_index=int(msg_idx) if msg_idx is not None else None,
        proximal=proximal_bool,
    )


def _computation_from_dict(d: dict) -> pb.Computation:
    """Convert a dict to a Computation protobuf message."""
    # Map status code string to enum
    status_str = d.get("status_code", "STATUS_UNSET")
    if status_str == "OK" or status_str == "STATUS_OK":
        status_code = pb.SpanStatusCode.STATUS_OK
    elif status_str == "ERROR" or status_str == "STATUS_ERROR":
        status_code = pb.SpanStatusCode.STATUS_ERROR
    else:
        status_code = pb.SpanStatusCode.STATUS_UNSET

    return pb.Computation(
        trace_id=d.get("trace_id", ""),
        span_id=d.get("span_id", ""),
        parent_span_id=d.get("parent_span_id"),
        name=d.get("name", ""),
        start_time_ns=d.get("start_time_ns", 0) or 0,
        end_time_ns=d.get("end_time_ns", 0) or 0,
        duration_ns=d.get("duration_ns", 0) or 0,
        status_code=status_code,
        status_message=d.get("status_message"),
        attributes_json=d.get("attributes", "{}"),
        events_json=d.get("events", "[]"),
        service_name=d.get("service_name"),
    )


def _computation_edge_from_dict(d: dict) -> pb.ComputationEdge:
    """Convert a dict to a ComputationEdge protobuf message."""
    return pb.ComputationEdge(
        parent_span_id=d.get("parent_span_id", ""),
        child_span_id=d.get("child_span_id", ""),
    )


def _computation_message_edge_from_dict(d: dict) -> pb.ComputationMessageEdge:
    """Convert a dict to a ComputationMessageEdge protobuf message."""
    edge_type_str = d.get("edge_type", "PRODUCES")
    if edge_type_str == "CONSUMES":
        edge_type = pb.ComputationMessageEdgeType.CONSUMES
    else:
        edge_type = pb.ComputationMessageEdgeType.PRODUCES

    msg_idx = d.get("message_index")
    return pb.ComputationMessageEdge(
        span_id=d.get("span_id", ""),
        message_id=d.get("message_id") or d.get("id", ""),
        edge_type=edge_type,
        message_index=int(msg_idx) if msg_idx is not None else None,
    )


def _update_type_from_str(type_str: str) -> pb.UpdateType:
    """Convert update type string to enum."""
    mapping = {
        "node_created": pb.UpdateType.NODE_CREATED,
        "node_deleted": pb.UpdateType.NODE_DELETED,
        "edge_created": pb.UpdateType.EDGE_CREATED,
        "edge_deleted": pb.UpdateType.EDGE_DELETED,
        "full_state": pb.UpdateType.FULL_STATE,
        "heartbeat": pb.UpdateType.HEARTBEAT,
        # Computation (OTel span) update types
        "computation_created": pb.UpdateType.COMPUTATION_CREATED,
        "computation_deleted": pb.UpdateType.COMPUTATION_DELETED,
        "computation_edge_created": pb.UpdateType.COMPUTATION_EDGE_CREATED,
        "computation_message_link_created": pb.UpdateType.COMPUTATION_MESSAGE_LINK_CREATED,
    }
    return mapping.get(type_str, pb.UpdateType.HEARTBEAT)


# Maximum chunk size for streaming large states (3MB, leaving headroom below 4MB gRPC limit)
MAX_CHUNK_SIZE = 3 * 1024 * 1024


def _generate_chunked_state(
    nodes: list[dict[str, Any]], edges: list[dict[str, Any]], sequence: int
) -> Iterator[pb.GraphUpdates]:
    """Generate chunked full state updates.

    Protocol uses repeated updates field to minimize messages:
    - Small graph: single message with [START, CHUNK(data), COMMIT]
    - Large graph: [START, CHUNK] → [CHUNK]... → [CHUNK, COMMIT]

    Nodes are sent first, then edges. Chunks may contain both if they fit.
    Client only acks after receiving COMMIT.
    """
    pending_updates: list[pb.GraphUpdate] = []
    started = False

    def make_chunk_update(chunk: pb.GraphStateChunk) -> pb.GraphUpdate:
        return pb.GraphUpdate(type=pb.UpdateType.FULL_STATE_CHUNK, state_chunk=chunk)

    def flush_message() -> pb.GraphUpdates:
        nonlocal pending_updates
        msg = pb.GraphUpdates(sequence=sequence, updates=pending_updates)
        pending_updates = []
        return msg

    # Always start with START
    pending_updates.append(pb.GraphUpdate(type=pb.UpdateType.FULL_STATE_START))
    started = True

    chunk = pb.GraphStateChunk()

    # 1. Add nodes to chunks
    for node_dict in nodes:
        event = event_from_dict(node_dict)
        chunk.events.append(event)
        if chunk.ByteSize() > MAX_CHUNK_SIZE:
            chunk.events.pop()  # Remove oversized item
            pending_updates.append(make_chunk_update(chunk))
            yield flush_message()
            chunk = pb.GraphStateChunk()
            chunk.events.append(event)

    # 2. Add edges to chunks (continue using existing chunk if it has room)
    for edge_dict in edges:
        edge = _dict_to_edge(edge_dict)
        chunk.edges.append(edge)
        if chunk.ByteSize() > MAX_CHUNK_SIZE:
            chunk.edges.pop()
            pending_updates.append(make_chunk_update(chunk))
            yield flush_message()
            chunk = pb.GraphStateChunk()
            chunk.edges.append(edge)

    # 3. Final message: remaining chunk (if any) + COMMIT
    if chunk.events or chunk.edges:
        pending_updates.append(make_chunk_update(chunk))
    pending_updates.append(pb.GraphUpdate(type=pb.UpdateType.FULL_STATE_COMMIT))
    yield flush_message()


class UpdatesServicer(observability_pb2_grpc.ObservabilityUpdatesServicer):
    """gRPC servicer for updates streaming with mTLS authentication."""

    def __init__(self, updates_manager: UpdatesManager):
        self.updates_manager = updates_manager

    @require_role("observability-reader")
    def GetState(
        self,
        request: pb.StateRequest,
        context: grpc.ServicerContext,
    ) -> pb.GraphState:
        """Get the current cached graph state."""
        state = self.updates_manager.get_cached_state()
        return pb.GraphState(
            events=[event_from_dict(n) for n in state["nodes"]],
            edges=[_dict_to_edge(e) for e in state["edges"]],
            sequence=state["sequence"],
        )

    @require_role("observability-reader")
    async def StreamUpdates(
        self,
        request_iterator: AsyncIterator[pb.StreamMessage],
        context: grpc.aio.ServicerContext,
    ) -> AsyncIterator[pb.GraphUpdates]:
        """
        Bidirectional streaming for graph updates.

        Client sends:
        - StreamMessage.subscribe: Initial subscription with from_sequence
        - StreamMessage.ack: Acknowledgment of received updates

        Server sends:
        - GraphUpdates messages with batched updates and sequence
        """
        client_id = str(uuid4())
        current_position = 0
        subscribed = False

        # Register client
        self.updates_manager.register_client(client_id)

        # Task to handle incoming client messages (acks)
        async def handle_client_messages() -> None:
            nonlocal current_position, subscribed
            try:
                async for message in request_iterator:
                    if message.HasField("subscribe"):
                        current_position = message.subscribe.from_sequence
                        subscribed = True
                    elif message.HasField("ack"):
                        self.updates_manager.acknowledge_position(
                            client_id, message.ack.position
                        )
            except Exception:
                pass  # Stream closed

        # Start client message handler in background
        client_task = asyncio.create_task(handle_client_messages())

        try:
            # Wait for subscription message
            for i in range(100):  # ~10 second timeout
                if subscribed:
                    break
                await asyncio.sleep(0.001 if i < 100 else 0.1)

            if not subscribed:
                await context.abort(
                    grpc.StatusCode.INVALID_ARGUMENT,
                    "Client must send subscribe message first",
                )
                return

            # Send initial heartbeat
            yield pb.GraphUpdates(
                sequence=current_position,
                updates=[pb.GraphUpdate(type=pb.UpdateType.HEARTBEAT)],
            )

            while self.updates_manager.is_client_registered(client_id):
                # Check if client needs full refresh
                if self.updates_manager.needs_full_refresh(current_position):
                    state = self.updates_manager.get_cached_state()
                    # Use chunked streaming for full state transfer
                    for chunk_msg in _generate_chunked_state(
                        state["nodes"], state["edges"], state["sequence"]
                    ):
                        yield chunk_msg
                    current_position = state["sequence"]
                    continue

                # Get pending updates and batch them
                updates = self.updates_manager.get_updates_since(current_position)
                batch: list[pb.GraphUpdate] = []

                for update in updates:
                    update_type = _update_type_from_str(update.get("type", ""))

                    msg = pb.GraphUpdate(
                        type=update_type,
                        timestamp=update.get("timestamp"),
                    )

                    if update_type == pb.UpdateType.NODE_CREATED:
                        event = event_from_dict(update)
                        msg.event.CopyFrom(event)
                    elif update_type == pb.UpdateType.NODE_DELETED:
                        msg.node_id = update.get("id", "")
                    elif update_type == pb.UpdateType.EDGE_CREATED:
                        msg_idx = update.get("message_index")
                        proximal = update.get("proximal")
                        proximal_bool = proximal == "true" if isinstance(proximal, str) else bool(proximal)
                        msg.edge.CopyFrom(pb.Edge(
                            source=update.get("source_id", ""),
                            destination=update.get("destination_id", ""),
                            message_index=int(msg_idx) if msg_idx is not None else None,
                            proximal=proximal_bool,
                        ))
                    elif update_type == pb.UpdateType.EDGE_DELETED:
                        msg.source_id = update.get("source_id", "")
                        msg.destination_id = update.get("destination_id", "")
                    # Computation (OTel span) update types
                    elif update_type == pb.UpdateType.COMPUTATION_CREATED:
                        msg.computation.CopyFrom(_computation_from_dict(update))
                    elif update_type == pb.UpdateType.COMPUTATION_DELETED:
                        msg.node_id = update.get("span_id", "")
                    elif update_type == pb.UpdateType.COMPUTATION_EDGE_CREATED:
                        msg.computation_edge.CopyFrom(
                            _computation_edge_from_dict(update)
                        )
                    elif update_type == pb.UpdateType.COMPUTATION_MESSAGE_LINK_CREATED:
                        msg.computation_message_edge.CopyFrom(
                            _computation_message_edge_from_dict(update)
                        )

                    batch.append(msg)
                    current_position = update["sequence"]

                # Send batch if we have updates, otherwise send heartbeat
                if batch:
                    yield pb.GraphUpdates(sequence=current_position, updates=batch)
                else:
                    yield pb.GraphUpdates(
                        sequence=current_position,
                        updates=[pb.GraphUpdate(type=pb.UpdateType.HEARTBEAT)],
                    )

                # Wait for new updates
                await self.updates_manager.wait_for_update(timeout=15.0)

        finally:
            client_task.cancel()
            try:
                await client_task
            except asyncio.CancelledError:
                pass
            if self.updates_manager.is_client_registered(client_id):
                self.updates_manager.unregister_client(client_id)


async def serve(
    host: str = "localhost",
    port: int = 10086,
    database_name: str = "neo4j",
    neo4j_uri: Optional[str] = None,
    neo4j_username: Optional[str] = None,
    neo4j_password: Optional[str] = None,
    poll_interval: float = 5.0,
    client_ack_timeout: float = 60.0,
    notify_debounce: float = 0.01,
    tls: bool = True,
    mtls: bool = True,
    cert_path: Optional[str] = None,
    key_path: Optional[str] = None,
    ca_path: Optional[str] = None,
    use_mkcert: bool = False,
    cert_dir: str = "certs",
    auto_install_triggers: bool = True,
    install_notify_trigger: bool = False,
    notify_url: Optional[str] = None,
    auth_config_path: Optional[str] = None,
) -> None:
    """
    Start the gRPC updates server with mTLS authentication.

    Args:
        host: Host to bind to
        port: Port to listen on (gRPC with mTLS)
        database_name: Neo4j database name
        neo4j_uri: Neo4j connection URI (or NEO4J_URI env var)
        neo4j_username: Neo4j username (or NEO4J_USERNAME env var)
        neo4j_password: Neo4j password (or NEO4J_PASSWORD env var)
        poll_interval: How often to poll for updates (seconds)
        client_ack_timeout: Seconds before unregistering unresponsive clients
        notify_debounce: Debounce delay for batching rapid notifications (seconds)
        tls: Enable TLS
        mtls: Enable mutual TLS (requires client certificates)
        cert_path: Path to server certificate
        key_path: Path to server private key
        ca_path: Path to CA certificate for client verification
        use_mkcert: Use mkcert for certificate generation
        cert_dir: Directory for generated certificates
        auto_install_triggers: Automatically install APOC triggers if not present
        install_notify_trigger: Install APOC trigger for push notifications (requires
            Neo4j to have network access to notify_url). If False, uses polling only.
        notify_url: URL for Neo4j to notify this server (defaults to http://host:port+1)
        auth_config_path: Path to auth config YAML for role lookups
    """
    global _manager

    uri = neo4j_uri or os.environ["NEO4J_URI"]
    auth = (
        neo4j_username or os.environ["NEO4J_USERNAME"],
        neo4j_password or os.environ["NEO4J_PASSWORD"],
    )

    _manager = UpdatesManager(
        neo4j_uri=uri,
        neo4j_auth=auth,
        database_name=database_name,
        poll_interval=poll_interval,
        client_ack_timeout=client_ack_timeout,
        notify_debounce=notify_debounce,
    )

    # Auto-install APOC triggers if enabled
    internal_port = port + 1  # e.g., 10087 for internal HTTP
    if auto_install_triggers:
        # Notify URL should point to the internal HTTP server
        updates_url = notify_url or f"http://{host}:{internal_port}"
        mode = "push+poll" if install_notify_trigger else "poll-only"
        print(f"Installing APOC triggers ({mode} mode)...")
        try:
            install_triggers(
                neo4j_uri=uri,
                neo4j_username=auth[0],
                neo4j_password=auth[1],
                database_name=database_name,
                updates_server_url=updates_url,
                install_notify_trigger=install_notify_trigger,
            )
        except Exception as e:
            print(f"Warning: Failed to install triggers: {e}")
            print("  Triggers may already be installed or APOC is not available")

    # Configure interceptors
    interceptors = []
    if mtls:
        interceptors.append(
            create_service_interceptor(
                "updates-server",
                auth_config_path=auth_config_path,
            )
        )
        print("mTLS auth enabled with role-based access control")
    else:
        disable_rbac()
        print("WARNING: Auth disabled - RBAC checks will be skipped")

    # Create gRPC server
    server = grpc.aio.server(interceptors=interceptors)  # type: ignore[arg-type]
    observability_pb2_grpc.add_ObservabilityUpdatesServicer_to_server(
        UpdatesServicer(_manager), server
    )

    url = f"{host}:{port}"

    # Configure TLS/mTLS for gRPC
    if tls or mtls:
        server_cert, server_key, ca_cert = load_tls_credentials(
            cert_path=cert_path,
            key_path=key_path,
            ca_path=ca_path,
            use_mkcert=use_mkcert,
            cert_name=host,
            cert_dir=cert_dir,
        )

        # Always pass CA cert so client certs can be optionally accepted
        # When mtls=True, client certs are required; when False, they're optional
        credentials = grpc.ssl_server_credentials(
            [(server_key, server_cert)],
            root_certificates=ca_cert,
            require_client_auth=mtls,
        )
        tls_mode = "mTLS (required)" if mtls else "TLS (mTLS optional)"
        print(f"{tls_mode} enabled")

        bound_port = server.add_secure_port(url, credentials)
    else:
        bound_port = server.add_insecure_port(url)
        print("Warning: Running without TLS (insecure)")

    if bound_port == 0:
        print(f"ERROR: Failed to bind port {url}")
        return

    # Initialize manager and start server
    _manager.refresh_cache()
    await server.start()

    print(f"Updates gRPC server started on {url}")
    print("gRPC methods:")
    print("  GetState      - Get current graph state")
    print("  StreamUpdates - Bidirectional stream for updates and acks")

    # Start background tasks
    poll_task = asyncio.create_task(_manager.poll_updates())

    # Start internal HTTP server for health checks and notifications
    print(f"Internal HTTP server on http://{host}:{internal_port} (for health/notify)")

    internal_app = FastAPI()

    @internal_app.get("/health")
    async def internal_health():
        return {"status": "healthy"}

    if install_notify_trigger:
        # Accept both GET and POST - apoc.load.jsonParams may use GET despite config
        @internal_app.api_route("/notify", methods=["GET", "POST"])
        async def internal_notify():
            await _manager.notify_update()
            return {"status": "notified"}

    internal_config = uvicorn.Config(
        internal_app, host=host, port=internal_port, log_level="warning"
    )
    internal_server = uvicorn.Server(internal_config)
    http_task = asyncio.create_task(internal_server.serve())

    try:
        await server.wait_for_termination()
    finally:
        _manager.stop()
        poll_task.cancel()
        http_task.cancel()
        try:
            await poll_task
        except asyncio.CancelledError:
            pass
        try:
            await http_task
        except asyncio.CancelledError:
            pass


if __name__ == "__main__":
    Fire(serve)
