from typing import Annotated, ClassVar, Final, Literal

from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError


type JupyterChannel = Literal["iopub", "shell", "stdin"]
type MessageStatus = Literal["ok", "error", "abort"]

STREAM: Final[Literal["stream"]] = "stream"
DISPLAY_DATA: Final[Literal["display_data"]] = "display_data"
UPDATE_DISPLAY_DATA: Final[Literal["update_display_data"]] = "update_display_data"
EXECUTE_RESULT: Final[Literal["execute_result"]] = "execute_result"
ERROR: Final[Literal["error"]] = "error"
STATUS: Final[Literal["status"]] = "status"
EXECUTE_INPUT: Final[Literal["execute_input"]] = "execute_input"
CLEAR_OUTPUT: Final[Literal["clear_output"]] = "clear_output"
EXECUTE_REPLY: Final[Literal["execute_reply"]] = "execute_reply"
INPUT_REQUEST: Final[Literal["input_request"]] = "input_request"


class MimeContent(BaseModel):
    """The shared payload for rich display messages."""

    data: dict[str, object]
    metadata: dict[str, object] = Field(default_factory=dict)
    transient: dict[str, object] | None = None


class ExecuteResultContent(MimeContent):
    """Rich display data plus the kernel execution counter."""

    execution_count: int


class StreamContent(BaseModel):
    """A stdout or stderr chunk emitted by the kernel."""

    name: Literal["stdout", "stderr"]
    text: str


class ErrorContent(BaseModel):
    """An exception reported on the IOPub channel."""

    ename: str
    evalue: str
    traceback: list[str]


class ExecuteReplyContent(BaseModel):
    """The terminal shell-channel reply for an execution request."""

    status: MessageStatus
    execution_count: int | None = None
    ename: str | None = None
    evalue: str | None = None
    traceback: list[str] | None = None


class StatusContent(BaseModel):
    """Kernel busy/idle state from an IOPub status message."""

    execution_state: Literal["busy", "idle", "starting"]


class InputRequestContent(BaseModel):
    """The prompt and echo policy for a kernel input request."""

    prompt: str
    password: bool = False


class ExecuteInputContent(BaseModel):
    """The code and execution counter echoed by the kernel."""

    code: str
    execution_count: int


class ClearOutputContent(BaseModel):
    """Whether a clear-output request should wait for new output."""

    wait: bool = False


class JupyterMessage(BaseModel):
    """Common Jupyter envelope, also capable of holding unknown messages.

    The protocol's message vocabulary is extensible, so ``msg_type`` stays a
    string on this raw-compatible model. Known messages are narrowed into the
    typed subclasses below by :func:`parse_jupyter_message`.
    """

    model_config: ClassVar[ConfigDict] = ConfigDict(extra="allow")

    channel: JupyterChannel = "iopub"
    msg_type: str
    message_id: str = ""
    parent_id: str | None = None
    metadata: dict[str, object] = Field(default_factory=dict)
    content: dict[str, object] = Field(default_factory=dict)
    buffers: list[bytes] = Field(default_factory=list)


class _TypedMessage(BaseModel):
    """Common fields for typed messages without narrowing mutable fields."""

    model_config: ClassVar[ConfigDict] = ConfigDict(extra="allow")

    message_id: str = ""
    parent_id: str | None = None
    metadata: dict[str, object] = Field(default_factory=dict)
    buffers: list[bytes] = Field(default_factory=list)


class StreamMessage(_TypedMessage):
    msg_type: Literal["stream"] = STREAM
    channel: Literal["iopub"] = "iopub"
    content: StreamContent


class DisplayDataMessage(_TypedMessage):
    msg_type: Literal["display_data"] = DISPLAY_DATA
    channel: Literal["iopub"] = "iopub"
    content: MimeContent


class UpdateDisplayDataMessage(_TypedMessage):
    msg_type: Literal["update_display_data"] = UPDATE_DISPLAY_DATA
    channel: Literal["iopub"] = "iopub"
    content: MimeContent


class ExecuteResultMessage(_TypedMessage):
    msg_type: Literal["execute_result"] = EXECUTE_RESULT
    channel: Literal["iopub"] = "iopub"
    content: ExecuteResultContent


class ErrorMessage(_TypedMessage):
    msg_type: Literal["error"] = ERROR
    channel: Literal["iopub"] = "iopub"
    content: ErrorContent


class StatusMessage(_TypedMessage):
    msg_type: Literal["status"] = STATUS
    channel: Literal["iopub"] = "iopub"
    content: StatusContent


class ExecuteInputMessage(_TypedMessage):
    msg_type: Literal["execute_input"] = EXECUTE_INPUT
    channel: Literal["iopub"] = "iopub"
    content: ExecuteInputContent


class ClearOutputMessage(_TypedMessage):
    msg_type: Literal["clear_output"] = CLEAR_OUTPUT
    channel: Literal["iopub"] = "iopub"
    content: ClearOutputContent


class ExecuteReplyMessage(_TypedMessage):
    msg_type: Literal["execute_reply"] = EXECUTE_REPLY
    channel: Literal["shell"] = "shell"
    content: ExecuteReplyContent


class InputRequestMessage(_TypedMessage):
    msg_type: Literal["input_request"] = INPUT_REQUEST
    channel: Literal["stdin"] = "stdin"
    content: InputRequestContent


class UnknownJupyterMessage(JupyterMessage):
    """An extension or future message preserved without lossy parsing."""

    # Allow any channel value so known-but-malformed fallbacks never fail
    # on a strict Literal. The parser preserves the transport channel verbatim.
    channel: str = "iopub"  # type: ignore[assignment]
    raw: dict[str, object] = Field(default_factory=dict)


type KnownJupyterMessage = Annotated[
    StreamMessage
    | DisplayDataMessage
    | UpdateDisplayDataMessage
    | ExecuteResultMessage
    | ErrorMessage
    | StatusMessage
    | ExecuteInputMessage
    | ClearOutputMessage
    | ExecuteReplyMessage
    | InputRequestMessage,
    Field(discriminator="msg_type"),
]
type ParsedJupyterMessage = KnownJupyterMessage | UnknownJupyterMessage

_KNOWN_MESSAGE_TYPES: Final[frozenset[str]] = frozenset(
    {
        STREAM,
        DISPLAY_DATA,
        UPDATE_DISPLAY_DATA,
        EXECUTE_RESULT,
        ERROR,
        STATUS,
        EXECUTE_INPUT,
        CLEAR_OUTPUT,
        EXECUTE_REPLY,
        INPUT_REQUEST,
    }
)
_KNOWN_MESSAGE_ADAPTER = TypeAdapter(KnownJupyterMessage)


def parse_jupyter_message(
    raw: dict[str, object],
    *,
    channel: JupyterChannel,
) -> ParsedJupyterMessage:
    """Parse a decoded Jupyter message without dropping extensions."""
    parent_header = raw.get("parent_header")
    parent_id = (
        parent_header.get("msg_id")
        if isinstance(parent_header, dict)
        and isinstance(parent_header.get("msg_id"), str)
        else None
    )
    raw_metadata = raw.get("metadata")
    metadata = raw_metadata if isinstance(raw_metadata, dict) else {}
    raw_content = raw.get("content")
    content = raw_content if isinstance(raw_content, dict) else {}
    raw_buffers = raw.get("buffers")
    buffers = (
        [
            bytes(buffer)
            for buffer in raw_buffers
            if isinstance(buffer, (bytes, bytearray, memoryview))
        ]
        if isinstance(raw_buffers, list)
        else []
    )
    raw_msg_type = raw.get("msg_type")
    raw_message_id = raw.get("msg_id")
    values: dict[str, object] = {
        "channel": channel,
        "msg_type": raw_msg_type if isinstance(raw_msg_type, str) else "",
        "message_id": raw_message_id if isinstance(raw_message_id, str) else "",
        "parent_id": parent_id,
        "metadata": metadata,
        "content": content,
        "buffers": buffers,
    }

    if values["msg_type"] not in _KNOWN_MESSAGE_TYPES:
        return UnknownJupyterMessage.model_validate({**values, "raw": raw})
    try:
        return _KNOWN_MESSAGE_ADAPTER.validate_python(values)
    except ValidationError:
        # Known msg_type but content or channel doesn't match the strict
        # model (e.g., extra execution_state, wrong stdout/stderr name,
        # or msg delivered on an unexpected channel). Degrade losslessly
        # instead of aborting the whole call.
        return UnknownJupyterMessage.model_validate({**values, "raw": raw})


# Existing callers use this name for the transport-level message. Keep it as
# an alias while the more expressive JupyterMessage name becomes canonical.
KernelMessage = JupyterMessage