from __future__ import annotations

from typing import Any, Literal, Optional, TypeAlias, TypedDict

from .agent_stream_state import ToolCallDeltaPayload, ToolRecord


class AssistantDeltaEvent(TypedDict):
    type: Literal["assistant_delta"]
    text: str
    model: Optional[str]


class AssistantReasoningDeltaEvent(TypedDict):
    type: Literal["assistant_reasoning_delta"]
    text: str
    model: Optional[str]


class ToolCallDeltaEvent(TypedDict, total=False):
    type: Literal["tool_call_delta"]
    tool: str
    args: str
    index: Optional[int]
    tool_call_id: str
    delta_id: str
    preview: str


class ToolCallRequestEvent(TypedDict):
    type: Literal["tool_call"]
    tool: str
    delta_id: str
    future: Any


class ToolOutputEvent(TypedDict):
    type: Literal["tool_output"]
    stream: Literal["stdout", "stderr"]
    text: str


class ToolExitEvent(TypedDict):
    type: Literal["tool_exit"]
    code: int


class ApprovalPrecheckPayload(TypedDict, total=False):
    status: str
    message: str | None
    reasons: list[str] | None
    metadata: dict[str, Any] | None


class ApprovalRequestEvent(TypedDict):
    type: Literal["approval_request"]
    id: str
    tool: str
    args: dict[str, Any]
    future: Any
    precheck: ApprovalPrecheckPayload


class ChoiceRequestEvent(TypedDict):
    type: Literal["choice_request"]
    id: str
    tool: str
    prompt: str
    choices: list[str]
    hint: str | None
    future: Any


class ErrorEvent(TypedDict, total=False):
    type: Literal["error"]
    error: str
    exception_type: str
    traceback: str


class AssistantFinalEvent(TypedDict, total=False):
    type: Literal["assistant_final"]
    text: str
    usage_data: Any
    reasoning: str


class MessageSentEvent(TypedDict):
    type: Literal["message_sent"]


class RunAbortedEvent(TypedDict):
    type: Literal["run_aborted"]
    reason: str


class DebugLatencyEvent(TypedDict):
    type: Literal["debug_latency"]
    stage: str
    ms: float


class ToolCallDebugDeltaEvent(TypedDict):
    type: Literal["tool_call_debug"]
    delta: ToolCallDeltaPayload


class ToolCallDebugWarningEvent(TypedDict):
    type: Literal["tool_call_debug"]
    warning: str


class ToolCallDebugAssembledEvent(TypedDict):
    type: Literal["tool_call_debug"]
    assembled: list[ToolRecord]


class ToolCallDebugStageEvent(TypedDict, total=False):
    type: Literal["tool_call_debug"]
    stage: Literal["assembled_final", "invalid_arguments"]
    tool: str
    id: str
    index: int
    args_summary: dict[str, Any]


ToolCallDebugEvent: TypeAlias = (
    ToolCallDebugDeltaEvent
    | ToolCallDebugWarningEvent
    | ToolCallDebugAssembledEvent
    | ToolCallDebugStageEvent
)


StreamOutEvent: TypeAlias = AssistantDeltaEvent | AssistantReasoningDeltaEvent | ToolCallDeltaEvent
CoreEvent: TypeAlias = ErrorEvent | AssistantFinalEvent | MessageSentEvent | RunAbortedEvent | DebugLatencyEvent
ToolEvent: TypeAlias = (
    ToolOutputEvent
    | ToolExitEvent
    | ToolCallRequestEvent
    | ToolCallDeltaEvent
    | ToolCallDebugEvent
    | ApprovalRequestEvent
    | ChoiceRequestEvent
)
EventPayload: TypeAlias = StreamOutEvent | CoreEvent | ToolEvent


def ev_assistant_delta(text: str, model: Optional[str] = None) -> AssistantDeltaEvent:
    return {"type": "assistant_delta", "text": text, "model": model}


def ev_assistant_reasoning_delta(text: str, model: Optional[str] = None) -> AssistantReasoningDeltaEvent:
    return {"type": "assistant_reasoning_delta", "text": text, "model": model}


def ev_error(error: str, *, exception_type: Optional[str] = None, traceback_text: Optional[str] = None) -> ErrorEvent:
    payload: ErrorEvent = {"type": "error", "error": error}
    if exception_type:
        payload["exception_type"] = exception_type
    if traceback_text:
        payload["traceback"] = traceback_text
    return payload


def ev_assistant_final(text: str, *, usage_data: Optional[Any] = None, reasoning: Optional[str] = None) -> AssistantFinalEvent:
    payload: AssistantFinalEvent = {"type": "assistant_final", "text": text}
    if usage_data is not None:
        payload["usage_data"] = usage_data
    if reasoning is not None:
        payload["reasoning"] = reasoning
    return payload


def ev_message_sent() -> MessageSentEvent:
    return {"type": "message_sent"}


def ev_run_aborted(reason: str = "cancelled") -> RunAbortedEvent:
    return {"type": "run_aborted", "reason": reason}


def ev_debug_latency(*, stage: str, ms: float) -> DebugLatencyEvent:
    return {"type": "debug_latency", "stage": stage, "ms": float(ms)}


def ev_tool_output(*, stream: Literal["stdout", "stderr"], text: str) -> ToolOutputEvent:
    return {"type": "tool_output", "stream": stream, "text": text}


def ev_tool_exit(code: int) -> ToolExitEvent:
    return {"type": "tool_exit", "code": int(code)}


def ev_tool_call_delta(
    *,
    tool: str,
    args: str,
    delta_id: str,
    index: Optional[int] = None,
    tool_call_id: Optional[str] = None,
    preview: Optional[str] = None,
) -> ToolCallDeltaEvent:
    payload: ToolCallDeltaEvent = {
        "type": "tool_call_delta",
        "tool": tool,
        "args": args,
        "delta_id": delta_id,
    }
    if index is not None:
        payload["index"] = index
    if tool_call_id:
        payload["tool_call_id"] = tool_call_id
    if preview:
        payload["preview"] = preview
    return payload


def ev_tool_call_request(*, tool: str, delta_id: str, future: Any) -> ToolCallRequestEvent:
    return {"type": "tool_call", "tool": tool, "delta_id": delta_id, "future": future}


def ev_approval_request(
    *,
    id: str,
    tool: str,
    args: dict[str, Any],
    future: Any,
    precheck: ApprovalPrecheckPayload,
) -> ApprovalRequestEvent:
    return {
        "type": "approval_request",
        "id": id,
        "tool": tool,
        "args": args,
        "future": future,
        "precheck": precheck,
    }


def ev_choice_request(
    *,
    id: str,
    tool: str,
    prompt: str,
    choices: list[str],
    hint: str | None,
    future: Any,
) -> ChoiceRequestEvent:
    return {
        "type": "choice_request",
        "id": id,
        "tool": tool,
        "prompt": prompt,
        "choices": choices,
        "hint": hint,
        "future": future,
    }


def ev_tool_call_debug_delta(delta: ToolCallDeltaPayload) -> ToolCallDebugDeltaEvent:
    return {"type": "tool_call_debug", "delta": delta}


def ev_tool_call_debug_warning(warning: str) -> ToolCallDebugWarningEvent:
    return {"type": "tool_call_debug", "warning": warning}


def ev_tool_call_debug_assembled(assembled: list[ToolRecord]) -> ToolCallDebugAssembledEvent:
    return {"type": "tool_call_debug", "assembled": assembled}


def ev_tool_call_debug_stage(
    *,
    stage: Literal["assembled_final", "invalid_arguments"],
    tool: str,
    id: str,
    index: int,
    args_summary: dict[str, Any],
) -> ToolCallDebugStageEvent:
    return {
        "type": "tool_call_debug",
        "stage": stage,
        "tool": tool,
        "id": id,
        "index": index,
        "args_summary": args_summary,
    }
