from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Literal, Optional, TypedDict, TypeAlias, cast


class ReasoningDetail(TypedDict, total=False):
    text: str


class UsageData(TypedDict, total=False):
    prompt_tokens: int
    completion_tokens: int
    total_tokens: int
    cache_read_input_tokens: int
    cache_creation_input_tokens: int


class ToolFunctionPayload(TypedDict):
    name: str
    arguments: str


class ToolRecord(TypedDict):
    index: int
    id: str
    function: ToolFunctionPayload


class ToolCallDeltaPayload(TypedDict, total=False):
    model: str
    output_index: int
    tool_call_id: str
    name: str
    arguments: str


class AssistantDeltaPayload(TypedDict, total=False):
    text: str
    model: str


class AssistantReasoningDeltaPayload(TypedDict, total=False):
    text: str
    model: str
    extra: "ReasoningExtra | None"


class AssistantDonePayload(TypedDict, total=False):
    usage: UsageData
    model: str


class EmptyPayload(TypedDict):
    pass


class ReasoningExtra(TypedDict, total=False):
    reasoning_details: list[ReasoningDetail]


@dataclass
class ToolCallAccumulator:
    index: int
    id: str
    name: str = ""
    arguments: str = ""


@dataclass
class TurnStreamState:
    final: str = ""
    reasoning: str = ""
    reasoning_details: list[ReasoningDetail] = field(default_factory=list)
    tool_phase: bool = False
    usage_data: Optional[UsageData] = None
    model_name: Optional[str] = None
    tool_calls_by_idx: dict[int, ToolCallAccumulator] = field(default_factory=dict)
    tool_calls_by_id: dict[str, ToolCallAccumulator] = field(default_factory=dict)
    next_tool_idx: int = 0


SignalType = Literal[
    "assistant_delta",
    "assistant_reasoning_delta",
    "assistant_done",
    "tool_call_delta",
    "tool_calls_done",
]


@dataclass
class StreamSignal:
    type: Literal["assistant_delta"]
    payload: AssistantDeltaPayload = field(default_factory=dict)


@dataclass
class StreamTransition:
    kind: Literal["assistant_delta"]
    payload: AssistantDeltaPayload = field(default_factory=dict)


@dataclass
class AssistantReasoningDeltaSignal:
    type: Literal["assistant_reasoning_delta"]
    payload: AssistantReasoningDeltaPayload = field(default_factory=dict)


@dataclass
class AssistantDoneSignal:
    type: Literal["assistant_done"]
    payload: AssistantDonePayload = field(default_factory=dict)


@dataclass
class ToolCallDeltaSignal:
    type: Literal["tool_call_delta"]
    payload: ToolCallDeltaPayload = field(default_factory=dict)


@dataclass
class ToolCallsDoneSignal:
    type: Literal["tool_calls_done"]
    payload: EmptyPayload = field(default_factory=dict)


StreamSignalUnion: TypeAlias = (
    StreamSignal
    | AssistantReasoningDeltaSignal
    | AssistantDoneSignal
    | ToolCallDeltaSignal
    | ToolCallsDoneSignal
)


class AssistantDeltaTransitionPayload(TypedDict):
    text: str
    model: Optional[str]


class AssistantDoneTransitionPayload(TypedDict, total=False):
    usage: UsageData


class ToolCallTransitionPayload(TypedDict):
    index: int
    id: str
    name: str
    arguments: str


class EmptyTransitionPayload(TypedDict):
    pass


@dataclass
class AssistantDeltaTransition:
    kind: Literal["assistant_delta"]
    payload: AssistantDeltaTransitionPayload


@dataclass
class AssistantReasoningDeltaTransition:
    kind: Literal["assistant_reasoning_delta"]
    payload: AssistantDeltaTransitionPayload


@dataclass
class AssistantDoneTransition:
    kind: Literal["assistant_done"]
    payload: AssistantDoneTransitionPayload = field(default_factory=dict)


@dataclass
class ToolCallDeltaTransition:
    kind: Literal["tool_call_delta"]
    payload: ToolCallTransitionPayload


@dataclass
class ToolCallDeltaDroppedTransition:
    kind: Literal["tool_call_delta_dropped"]
    payload: EmptyTransitionPayload = field(default_factory=dict)


@dataclass
class ToolCallsDoneTransition:
    kind: Literal["tool_calls_done"]
    payload: EmptyTransitionPayload = field(default_factory=dict)


@dataclass
class NoopTransition:
    kind: Literal["noop"]
    payload: EmptyTransitionPayload = field(default_factory=dict)


StreamTransitionUnion: TypeAlias = (
    AssistantDeltaTransition
    | AssistantReasoningDeltaTransition
    | AssistantDoneTransition
    | ToolCallDeltaTransition
    | ToolCallDeltaDroppedTransition
    | ToolCallsDoneTransition
    | NoopTransition
)


def ingest_stream_event(ev: Any) -> list[StreamSignalUnion]:
    if not isinstance(ev, dict):
        return []
    et = ev.get("type")
    if et == "assistant_delta":
        payload: AssistantDeltaPayload = {"text": ev.get("text", ""), "model": ev.get("model")}
        return [StreamSignal(type="assistant_delta", payload=payload)]
    if et == "assistant_reasoning_delta":
        raw_extra = ev.get("extra")
        extra: ReasoningExtra | None = cast(ReasoningExtra, raw_extra) if isinstance(raw_extra, dict) else None
        payload: AssistantReasoningDeltaPayload = {
            "text": ev.get("text", ""),
            "model": ev.get("model"),
            "extra": extra,
        }
        return [
            AssistantReasoningDeltaSignal(
                type="assistant_reasoning_delta",
                payload=payload,
            )
        ]
    if et == "assistant_done":
        payload: AssistantDonePayload = {"usage": ev.get("usage"), "model": ev.get("model")}
        return [AssistantDoneSignal(type="assistant_done", payload=payload)]
    if et == "tool_call_delta":
        payload: ToolCallDeltaPayload = {
            "model": ev.get("model"),
            "output_index": ev.get("output_index"),
            "tool_call_id": ev.get("tool_call_id"),
            "name": ev.get("name"),
            "arguments": ev.get("arguments"),
        }
        return [
            ToolCallDeltaSignal(
                type="tool_call_delta",
                payload=payload,
            )
        ]
    if et == "tool_calls_done":
        return [ToolCallsDoneSignal(type="tool_calls_done", payload={})]
    return []


def reduce_stream_signal(state: TurnStreamState, signal: StreamSignalUnion) -> StreamTransitionUnion:
    if signal.type == "assistant_delta":
        text = str(signal.payload.get("text", ""))
        if text:
            state.final += text
        current_model_name = signal.payload.get("model")
        if current_model_name is not None:
            state.model_name = current_model_name
        return AssistantDeltaTransition(kind="assistant_delta", payload={"text": text, "model": state.model_name})

    if signal.type == "assistant_reasoning_delta":
        text = str(signal.payload.get("text", ""))
        if text:
            state.reasoning += text
        extra = signal.payload.get("extra")
        if isinstance(extra, dict) and extra.get("reasoning_details") is not None:
            details = extra.get("reasoning_details")
            if isinstance(details, list) and details:
                if not state.reasoning_details:
                    first = details[0]
                    if isinstance(first, dict):
                        state.reasoning_details.append(dict(first))
                else:
                    if isinstance(details[0], dict):
                        state.reasoning_details[0]["text"] = str(state.reasoning_details[0].get("text", "")) + str(
                            details[0].get("text", "")
                        )
        current_model_name = signal.payload.get("model")
        if current_model_name is not None:
            state.model_name = current_model_name
        return AssistantReasoningDeltaTransition(
            kind="assistant_reasoning_delta",
            payload={"text": text, "model": state.model_name},
        )

    if signal.type == "assistant_done":
        usage_data = signal.payload.get("usage")
        if usage_data is not None:
            state.usage_data = usage_data
        current_model_name = signal.payload.get("model")
        if current_model_name is not None:
            state.model_name = current_model_name
        payload: AssistantDoneTransitionPayload = {}
        if usage_data is not None:
            payload["usage"] = usage_data
        return AssistantDoneTransition(kind="assistant_done", payload=payload)

    if signal.type == "tool_call_delta":
        state.tool_phase = True
        current_model_name = signal.payload.get("model")
        if current_model_name is not None:
            state.model_name = current_model_name
        idx = signal.payload.get("output_index")
        tcid = signal.payload.get("tool_call_id")
        rec = None
        if isinstance(idx, int):
            rec = state.tool_calls_by_idx.get(idx)
            if rec is None:
                rec = ToolCallAccumulator(index=idx, id=str(tcid) if tcid else "")
                state.tool_calls_by_idx[idx] = rec
            elif (not rec.id) and tcid:
                rec.id = str(tcid)
            if tcid:
                state.tool_calls_by_id[str(tcid)] = rec
            if idx >= state.next_tool_idx:
                state.next_tool_idx = idx + 1
        elif tcid:
            rec = state.tool_calls_by_id.get(str(tcid))
            if rec is None:
                rec = ToolCallAccumulator(index=state.next_tool_idx, id=str(tcid))
                state.tool_calls_by_idx[state.next_tool_idx] = rec
                state.tool_calls_by_id[str(tcid)] = rec
                state.next_tool_idx += 1

        if rec is None:
            return ToolCallDeltaDroppedTransition(kind="tool_call_delta_dropped")

        name = signal.payload.get("name")
        if name:
            rec.name = str(name)
        arguments = signal.payload.get("arguments")
        if arguments:
            rec.arguments += str(arguments)
        return ToolCallDeltaTransition(
            kind="tool_call_delta",
            payload={
                "index": rec.index,
                "id": rec.id,
                "name": rec.name,
                "arguments": rec.arguments,
            },
        )

    if signal.type == "tool_calls_done":
        state.tool_phase = True
        return ToolCallsDoneTransition(kind="tool_calls_done")

    return NoopTransition(kind="noop")


def build_assembled_debug(state: TurnStreamState) -> list[ToolRecord]:
    assembled = []
    for k, rec in sorted(state.tool_calls_by_idx.items(), key=lambda x: x[0]):
        assembled.append(
            {
                "index": k,
                "id": rec.id,
                "function": {"name": rec.name, "arguments": rec.arguments},
            }
        )
    return assembled


def build_tool_records(state: TurnStreamState) -> list[ToolRecord]:
    out: list[ToolRecord] = []
    for _, rec in sorted(state.tool_calls_by_idx.items(), key=lambda x: x[0]):
        out.append(
            {
                "index": rec.index,
                "id": rec.id,
                "function": {"name": rec.name, "arguments": rec.arguments},
            }
        )
    return out
