"""Module contains typedefs that are used with `Runnable` objects."""

from __future__ import annotations

from typing import TYPE_CHECKING, Any, Literal

from typing_extensions import NotRequired, TypedDict

if TYPE_CHECKING:
    from collections.abc import Sequence


class EventData(TypedDict, total=False):
    """Data associated with a streaming event."""

    input: Any
    """The input passed to the `Runnable` that generated the event.

    Inputs will sometimes be available at the *START* of the `Runnable`, and
    sometimes at the *END* of the `Runnable`.

    If a `Runnable` is able to stream its inputs, then its input by definition
    won't be known until the *END* of the `Runnable` when it has finished streaming
    its inputs.
    """
    error: NotRequired[BaseException]
    """The error that occurred during the execution of the `Runnable`.

    This field is only available if the `Runnable` raised an exception.

    !!! version-added "Added in `langchain-core` 1.0.0"
    """
    output: Any
    """The output of the `Runnable` that generated the event.

    Outputs will only be available at the *END* of the `Runnable`.

    For most `Runnable` objects, this field can be inferred from the `chunk` field,
    though there might be some exceptions for special a cased `Runnable` (e.g., like
    chat models), which may return more information.
    """
    chunk: Any
    """A streaming chunk from the output that generated the event.

    chunks support addition in general, and adding them up should result
    in the output of the `Runnable` that generated the event.
    """


class BaseStreamEvent(TypedDict):
    """Streaming event.

    Schema of a streaming event which is produced from the `astream_events` method.

    Example:
        ```python
        from langchain_core.runnables import RunnableLambda


        async def reverse(s: str) -> str:
            return s[::-1]


        chain = RunnableLambda(func=reverse)

        events = [event async for event in chain.astream_events("hello")]

        # Will produce the following events
        # (where some fields have been omitted for brevity):
        [
            {
                "data": {"input": "hello"},
                "event": "on_chain_start",
                "metadata": {},
                "name": "reverse",
                "tags": [],
            },
            {
                "data": {"chunk": "olleh"},
                "event": "on_chain_stream",
                "metadata": {},
                "name": "reverse",
                "tags": [],
            },
            {
                "data": {"output": "olleh"},
                "event": "on_chain_end",
                "metadata": {},
                "name": "reverse",
                "tags": [],
            },
        ]
        ```
    """

    event: str
    """Event names are of the format: `on_[runnable_type]_(start|stream|end)`.

    Runnable types are one of:

    - **llm** - used by non chat models
    - **chat_model** - used by chat models
    - **prompt** --  e.g., `ChatPromptTemplate`
    - **tool** -- from tools defined via `@tool` decorator or inheriting
        from `Tool`/`BaseTool`
    - **chain** - most `Runnable` objects are of this type

    Further, the events are categorized as one of:

    - **start** - when the `Runnable` starts
    - **stream** - when the `Runnable` is streaming
    - **end* - when the `Runnable` ends

    start, stream and end are associated with slightly different `data` payload.

    Please see the documentation for `EventData` for more details.
    """
    run_id: str
    """An randomly generated ID to keep track of the execution of the given `Runnable`.

    Each child `Runnable` that gets invoked as part of the execution of a parent
    `Runnable` is assigned its own unique ID.
    """
    tags: NotRequired[list[str]]
    """Tags associated with the `Runnable` that generated this event.

    Tags are always inherited from parent `Runnable` objects.

    Tags can either be bound to a `Runnable` using `.with_config({"tags":  ["hello"]})`
    or passed at run time using `.astream_events(..., {"tags": ["hello"]})`.
    """
    metadata: NotRequired[dict[str, Any]]
    """Metadata associated with the `Runnable` that generated this event.

    Metadata can either be bound to a `Runnable` using

        `.with_config({"metadata": { "foo": "bar" }})`

    or passed at run time using

        `.astream_events(..., {"metadata": {"foo": "bar"}})`.
    """

    parent_ids: Sequence[str]
    """A list of the parent IDs associated with this event.

    Root Events will have an empty list.

    For example, if a `Runnable` A calls `Runnable` B, then the event generated by
    `Runnable` B will have `Runnable` A's ID in the `parent_ids` field.

    The order of the parent IDs is from the root parent to the immediate parent.

    Only supported as of v2 of the astream events API. v1 will return an empty list.
    """


class StandardStreamEvent(BaseStreamEvent):
    """A standard stream event that follows LangChain convention for event data."""

    data: EventData
    """Event data.

    The contents of the event data depend on the event type.
    """
    name: str
    """The name of the `Runnable` that generated the event."""


class CustomStreamEvent(BaseStreamEvent):
    """Custom stream event created by the user."""

    # Overwrite the event field to be more specific.
    event: Literal["on_custom_event"]  # type: ignore[misc]
    """The event type."""
    name: str
    """User defined name for the event."""
    data: Any
    """The data associated with the event. Free form and can be anything."""


StreamEvent = StandardStreamEvent | CustomStreamEvent