"""
CopilotKit Protocol
"""

import json
from enum import Enum
from typing import Union, Optional
from typing_extensions import TypedDict, Literal, Any, Dict

class RuntimeEventTypes(Enum):
    """CopilotKit Runtime Event Types"""
    TEXT_MESSAGE_START = "TextMessageStart"
    TEXT_MESSAGE_CONTENT = "TextMessageContent"
    TEXT_MESSAGE_END = "TextMessageEnd"
    ACTION_EXECUTION_START = "ActionExecutionStart"
    ACTION_EXECUTION_ARGS = "ActionExecutionArgs"
    ACTION_EXECUTION_END = "ActionExecutionEnd"
    ACTION_EXECUTION_RESULT = "ActionExecutionResult"
    AGENT_STATE_MESSAGE = "AgentStateMessage"
    META_EVENT = "MetaEvent"
    RUN_STARTED = "RunStarted"
    RUN_FINISHED = "RunFinished"
    RUN_ERROR = "RunError"
    NODE_STARTED = "NodeStarted"
    NODE_FINISHED = "NodeFinished"

class RuntimeMetaEventName(Enum):
    """Runtime Meta Event Name"""
    LANG_GRAPH_INTERRUPT_EVENT = "LangGraphInterruptEvent"
    PREDICT_STATE = "PredictState"
    EXIT = "Exit"


class TextMessageStart(TypedDict):
    """Text Message Start Event"""
    type: Literal[RuntimeEventTypes.TEXT_MESSAGE_START]
    messageId: str
    parentMessageId: Optional[str]

class TextMessageContent(TypedDict):
    """Text Message Content Event"""
    type: Literal[RuntimeEventTypes.TEXT_MESSAGE_CONTENT]
    messageId: str
    content: str

class TextMessageEnd(TypedDict):
    """Text Message End Event"""
    type: Literal[RuntimeEventTypes.TEXT_MESSAGE_END]
    messageId: str

class ActionExecutionStart(TypedDict):
    """Action Execution Start Event"""
    type: Literal[RuntimeEventTypes.ACTION_EXECUTION_START]
    actionExecutionId: str
    actionName: str
    parentMessageId: Optional[str]

class ActionExecutionArgs(TypedDict):
    """Action Execution Args Event"""
    type: Literal[RuntimeEventTypes.ACTION_EXECUTION_ARGS]
    actionExecutionId: str
    args: str

class ActionExecutionEnd(TypedDict):
    """Action Execution End Event"""
    type: Literal[RuntimeEventTypes.ACTION_EXECUTION_END]
    actionExecutionId: str

class ActionExecutionResult(TypedDict):
    """Action Execution Result Event"""
    type: Literal[RuntimeEventTypes.ACTION_EXECUTION_RESULT]
    actionName: str
    actionExecutionId: str
    result: str

class AgentStateMessage(TypedDict):
    """Agent State Message Event"""
    type: Literal[RuntimeEventTypes.AGENT_STATE_MESSAGE]
    threadId: str
    agentName: str
    nodeName: str
    runId: str
    active: bool
    role: str
    state: str
    running: bool

class MetaEvent(TypedDict):
    """Meta Event"""
    type: Literal[RuntimeEventTypes.META_EVENT]
    name: RuntimeMetaEventName
    value: Any

class RunStarted(TypedDict):
    """Run Started Event"""
    type: Literal[RuntimeEventTypes.RUN_STARTED]
    state: Dict[str, Any]

class RunFinished(TypedDict):
    """Run Finished Event"""
    type: Literal[RuntimeEventTypes.RUN_FINISHED]
    state: Dict[str, Any]

class RunError(TypedDict):
    """Run Error Event"""
    type: Literal[RuntimeEventTypes.RUN_ERROR]
    error: Any

class NodeStarted(TypedDict):
    """Node Started Event"""
    type: Literal[RuntimeEventTypes.NODE_STARTED]
    node_name: str
    state: Dict[str, Any]

class NodeFinished(TypedDict):
    """Node Finished Event"""
    type: Literal[RuntimeEventTypes.NODE_FINISHED]
    node_name: str
    state: Dict[str, Any]

RuntimeProtocolEvent = Union[
    TextMessageStart,
    TextMessageContent,
    TextMessageEnd,
    ActionExecutionStart,
    ActionExecutionArgs,
    ActionExecutionEnd,
    ActionExecutionResult,
    AgentStateMessage,
    MetaEvent
]

RuntimeLifecycleEvent = Union[
    RunStarted,
    RunFinished,
    RunError,
    NodeStarted,
    NodeFinished,
]

RuntimeEvent = Union[
    RuntimeProtocolEvent,
    RuntimeLifecycleEvent,
]


class PredictStateConfig(TypedDict):
    """
    Predict State Config
    """
    tool_name: str
    tool_argument: Optional[str]

def text_message_start(
        *,
        message_id: str,
        parent_message_id: Optional[str] = None
  ) -> TextMessageStart:
    """Utility function to create a text message start event"""
    return {
        "type": RuntimeEventTypes.TEXT_MESSAGE_START,
        "messageId": message_id,
        "parentMessageId": parent_message_id
    }

def text_message_content(*, message_id: str, content: str) -> TextMessageContent:
    """Utility function to create a text message content event"""
    return {
        "type": RuntimeEventTypes.TEXT_MESSAGE_CONTENT,
        "messageId": message_id,
        "content": content
    }

def text_message_end(*, message_id: str) -> TextMessageEnd:
    """Utility function to create a text message end event"""
    return {
        "type": RuntimeEventTypes.TEXT_MESSAGE_END,
        "messageId": message_id
    }

def action_execution_start(
        *,
        action_execution_id: str,
        action_name: str,
        parent_message_id: Optional[str] = None
    ) -> ActionExecutionStart:
    """Utility function to create an action execution start event"""
    return {
        "type": RuntimeEventTypes.ACTION_EXECUTION_START,
        "actionExecutionId": action_execution_id,
        "actionName": action_name,
        "parentMessageId": parent_message_id
    }

def action_execution_args(*, action_execution_id: str, args: str) -> ActionExecutionArgs:
    """Utility function to create an action execution args event"""
    return {
        "type": RuntimeEventTypes.ACTION_EXECUTION_ARGS,
        "actionExecutionId": action_execution_id,
        "args": args
    }

def action_execution_end(*, action_execution_id: str) -> ActionExecutionEnd:
    """Utility function to create an action execution end event"""
    return {
        "type": RuntimeEventTypes.ACTION_EXECUTION_END,
        "actionExecutionId": action_execution_id
    }

def action_execution_result(
        *,
        action_name: str,
        action_execution_id: str,
        result: str
    ) -> ActionExecutionResult:
    """Utility function to create an action execution result event"""
    return {
        "type": RuntimeEventTypes.ACTION_EXECUTION_RESULT,
        "actionName": action_name,
        "actionExecutionId": action_execution_id,
        "result": result
    }

def agent_state_message( # pylint: disable=too-many-arguments
        *,
        thread_id: str,
        agent_name: str,
        node_name: str,
        run_id: str,
        active: bool,
        role: str,
        state: str,
        running: bool
  ) -> AgentStateMessage:
    """Utility function to create an agent state message event"""
    return {
        "type": RuntimeEventTypes.AGENT_STATE_MESSAGE,
        "threadId": thread_id,
        "agentName": agent_name,
        "nodeName": node_name,
        "runId": run_id,
        "active": active,
        "role": role,
        "state": state,
        "running": running
    }

def meta_event(*, name: RuntimeMetaEventName, value: Any) -> MetaEvent:
    """Utility function to create a meta event"""
    return {
        "type": RuntimeEventTypes.META_EVENT,
        "name": name,
        "value": value
    }

def emit_runtime_events(*events: RuntimeProtocolEvent) -> str:
    """Emit a list of runtime events"""
    def serialize_event(event):
        # Convert enum values to their string representation
        if isinstance(event, dict):
            return {k: (v.value if isinstance(v, Enum) else v) for k, v in event.items()}
        return event

    return "\n".join(json.dumps(serialize_event(event)) for event in events) + "\n"

def emit_runtime_event(event: RuntimeProtocolEvent) -> str:
    """Emit a single runtime event"""
    return emit_runtime_events(event)