"""Event sink contracts used by the runner and storage adapters.

The runner only needs append-only event writes and status transitions, so tests
can use ``InMemoryRunEventSink`` without Redis. Production storage implements the
same protocol with Redis streams in ``dify_agent.storage.redis_run_store``. The
terminal success helper writes either the final JSON-safe output or one deferred
tool request together with the resumable session snapshot in a single event so
consumers can stop at ``run_succeeded`` without correlating separate payload
events.
"""

from collections import defaultdict
from typing import Protocol, cast

from pydantic import JsonValue
from pydantic_ai.messages import AgentStreamEvent

from agenton.compositor import CompositorSessionSnapshot
from dify_agent.protocol.schemas import (
    AgentRunUsage,
    DeferredToolCallPayload,
    EmptyRunEventData,
    PydanticAIStreamRunEvent,
    RunEvent,
    RunCancelledEvent,
    RunCancelledEventData,
    RunFailedEvent,
    RunFailedEventData,
    RunStartedEvent,
    RunStatus,
    RunSucceededEvent,
    RunSucceededEventData,
    utc_now,
)


_UNSET = object()


class RunEventSink(Protocol):
    """Boundary used by runtime code to publish observable run progress."""

    async def append_event(self, event: RunEvent) -> str:
        """Persist ``event`` and return its cursor id."""
        ...

    async def update_status(self, run_id: str, status: RunStatus, error: str | None = None) -> None:
        """Persist the current run status."""
        ...


class InMemoryRunEventSink:
    """Small async-compatible sink for local unit tests and examples."""

    events: dict[str, list[RunEvent]]
    statuses: dict[str, RunStatus]
    errors: dict[str, str | None]

    def __init__(self) -> None:
        self.events = defaultdict(list)
        self.statuses = {}
        self.errors = {}

    async def append_event(self, event: RunEvent) -> str:
        """Store an event and assign a monotonic per-run cursor."""
        event_id = str(len(self.events[event.run_id]) + 1)
        stored = event.model_copy(update={"id": event_id})
        self.events[event.run_id].append(stored)
        return event_id

    async def update_status(self, run_id: str, status: RunStatus, error: str | None = None) -> None:
        """Record the latest status; timestamps are owned by run stores."""
        self.statuses[run_id] = status
        self.errors[run_id] = error


async def emit_run_event(
    sink: RunEventSink,
    *,
    event: RunEvent,
) -> str:
    """Append an already typed public run event."""
    return await sink.append_event(event)


async def emit_run_started(sink: RunEventSink, *, run_id: str) -> str:
    """Emit the first lifecycle event for one run."""
    return await emit_run_event(
        sink,
        event=RunStartedEvent(run_id=run_id, data=EmptyRunEventData(), created_at=utc_now()),
    )


async def emit_pydantic_ai_event(
    sink: RunEventSink,
    *,
    run_id: str,
    data: AgentStreamEvent,
    agent_message_delta: str | None = None,
) -> str:
    """Emit one typed Pydantic AI stream event."""
    return await emit_run_event(
        sink,
        event=PydanticAIStreamRunEvent(
            run_id=run_id,
            data=data,
            agent_message_delta=agent_message_delta,
            created_at=utc_now(),
        ),
    )


async def emit_run_succeeded(
    sink: RunEventSink,
    *,
    run_id: str,
    output: JsonValue | None | object = _UNSET,
    deferred_tool_call: DeferredToolCallPayload | object = _UNSET,
    session_snapshot: CompositorSessionSnapshot,
    usage: AgentRunUsage | None = None,
) -> str:
    """Emit the terminal success event with output or deferred continuation.

    Callers must activate exactly one result branch. ``_UNSET`` is used instead
    of ``None`` to preserve the distinction between an omitted inactive branch
    and an active ``output`` branch whose JSON value is explicitly ``null``.
    Without that sentinel, ``output=None`` would be indistinguishable from
    “output field absent”, which would break nullable-success payloads.
    """
    data: dict[str, JsonValue | DeferredToolCallPayload | CompositorSessionSnapshot | AgentRunUsage | None] = {
        "session_snapshot": session_snapshot,
    }
    if output is not _UNSET:
        data["output"] = cast(JsonValue | None, output)
    if deferred_tool_call is not _UNSET:
        data["deferred_tool_call"] = cast(DeferredToolCallPayload, deferred_tool_call)
    if usage is not None:
        data["usage"] = usage

    return await emit_run_event(
        sink,
        event=RunSucceededEvent(
            run_id=run_id,
            data=RunSucceededEventData.model_validate(data),
            created_at=utc_now(),
        ),
    )


async def emit_run_failed(
    sink: RunEventSink,
    *,
    run_id: str,
    error: str,
    reason: str | None = None,
) -> str:
    """Emit the terminal failure lifecycle event."""
    return await emit_run_event(
        sink,
        event=RunFailedEvent(run_id=run_id, data=RunFailedEventData(error=error, reason=reason), created_at=utc_now()),
    )


async def emit_run_cancelled(
    sink: RunEventSink,
    *,
    run_id: str,
    reason: str | None = None,
    message: str | None = None,
) -> str:
    """Emit the terminal cancellation lifecycle event."""
    return await emit_run_event(
        sink,
        event=RunCancelledEvent(
            run_id=run_id,
            data=RunCancelledEventData(reason=reason, message=message),
            created_at=utc_now(),
        ),
    )


__all__ = [
    "InMemoryRunEventSink",
    "RunEventSink",
    "emit_pydantic_ai_event",
    "emit_run_cancelled",
    "emit_run_event",
    "emit_run_failed",
    "emit_run_started",
    "emit_run_succeeded",
]
