from __future__ import annotations

import copy
import inspect
import time
from dataclasses import dataclass, field
from typing import (
    Any,
    Awaitable,
    Callable,
    Literal,
    Mapping,
    MutableMapping,
    Protocol,
    Required,
    Sequence,
    TypedDict,
    TypeVar,
    Union,
)

from .event_stream import AssistantMessageEventStream

ThinkingLevel = Literal["off", "minimal", "low", "medium", "high", "xhigh"]
Transport = Literal["sse", "websocket", "auto"]
StopReason = Literal["stop", "length", "toolUse", "error", "aborted"]


class ThinkingBudgets(TypedDict, total=False):
    minimal: int
    low: int
    medium: int
    high: int


class StreamCost(TypedDict):
    input: float
    output: float
    cacheRead: float
    cacheWrite: float
    total: float


class Usage(TypedDict):
    input: int
    output: int
    cacheRead: int
    cacheWrite: int
    totalTokens: int
    cost: StreamCost


class ModelCost(TypedDict):
    input: float
    output: float
    cacheRead: float
    cacheWrite: float


class Model(TypedDict, total=False):
    id: str
    name: str
    api: str
    provider: str
    baseUrl: str
    reasoning: bool
    input: list[Literal["text", "image"]]
    cost: ModelCost
    contextWindow: int
    maxTokens: int
    headers: dict[str, str]


DEFAULT_MODEL: Model = {
    "id": "gemini-2.5-flash-lite-preview-06-17",
    "name": "gemini-2.5-flash-lite-preview-06-17",
    "api": "google-generative-ai",
    "provider": "google",
    "baseUrl": "https://example.invalid",
    "reasoning": False,
    "input": ["text"],
    "cost": {"input": 0.0, "output": 0.0, "cacheRead": 0.0, "cacheWrite": 0.0},
    "contextWindow": 0,
    "maxTokens": 0,
}


class TextContent(TypedDict, total=False):
    type: Literal["text"]
    text: str
    textSignature: str


class ThinkingContent(TypedDict, total=False):
    type: Literal["thinking"]
    thinking: str
    thinkingSignature: str


class ImageContent(TypedDict):
    type: Literal["image"]
    data: str
    mimeType: str


class ToolCall(TypedDict, total=False):
    type: Literal["toolCall"]
    id: str
    name: str
    arguments: dict[str, Any]
    thoughtSignature: str


class UserMessage(TypedDict):
    role: Literal["user"]
    content: str | list[TextContent | ImageContent]
    timestamp: int


class AssistantMessage(TypedDict, total=False):
    role: Literal["assistant"]
    content: list[TextContent | ThinkingContent | ToolCall]
    api: str
    provider: str
    model: str
    usage: Usage
    stopReason: StopReason
    errorMessage: str
    timestamp: int


class ToolResultMessage(TypedDict, total=False):
    role: Literal["toolResult"]
    toolCallId: str
    toolName: str
    content: list[TextContent | ImageContent]
    details: Any
    isError: bool
    timestamp: int


Message = UserMessage | AssistantMessage | ToolResultMessage
# Custom app messages are represented as free-form mappings.
AgentMessage = Message | dict[str, Any]


class Tool(TypedDict, total=False):
    name: str
    description: str
    parameters: dict[str, Any]


class Context(TypedDict, total=False):
    systemPrompt: str
    messages: list[Message]
    tools: list[Tool]


class AssistantStartEvent(TypedDict):
    type: Literal["start"]
    partial: AssistantMessage


class AssistantTextStartEvent(TypedDict):
    type: Literal["text_start"]
    contentIndex: int
    partial: AssistantMessage


class AssistantTextDeltaEvent(TypedDict):
    type: Literal["text_delta"]
    contentIndex: int
    delta: str
    partial: AssistantMessage


class AssistantTextEndEvent(TypedDict):
    type: Literal["text_end"]
    contentIndex: int
    content: str
    partial: AssistantMessage


class AssistantThinkingStartEvent(TypedDict):
    type: Literal["thinking_start"]
    contentIndex: int
    partial: AssistantMessage


class AssistantThinkingDeltaEvent(TypedDict):
    type: Literal["thinking_delta"]
    contentIndex: int
    delta: str
    partial: AssistantMessage


class AssistantThinkingEndEvent(TypedDict):
    type: Literal["thinking_end"]
    contentIndex: int
    content: str
    partial: AssistantMessage


class AssistantToolCallStartEvent(TypedDict):
    type: Literal["toolcall_start"]
    contentIndex: int
    partial: AssistantMessage


class AssistantToolCallDeltaEvent(TypedDict):
    type: Literal["toolcall_delta"]
    contentIndex: int
    delta: str
    partial: AssistantMessage


class AssistantToolCallEndEvent(TypedDict):
    type: Literal["toolcall_end"]
    contentIndex: int
    toolCall: ToolCall
    partial: AssistantMessage


class AssistantDoneEvent(TypedDict):
    type: Literal["done"]
    reason: Literal["stop", "length", "toolUse"]
    message: AssistantMessage


class AssistantErrorEvent(TypedDict):
    type: Literal["error"]
    reason: Literal["aborted", "error"]
    error: AssistantMessage


AssistantMessageEvent = Union[
    AssistantStartEvent,
    AssistantTextStartEvent,
    AssistantTextDeltaEvent,
    AssistantTextEndEvent,
    AssistantThinkingStartEvent,
    AssistantThinkingDeltaEvent,
    AssistantThinkingEndEvent,
    AssistantToolCallStartEvent,
    AssistantToolCallDeltaEvent,
    AssistantToolCallEndEvent,
    AssistantDoneEvent,
    AssistantErrorEvent,
]


class SupportsAbort(Protocol):
    @property
    def aborted(self) -> bool: ...


class StreamResponse(Protocol):
    def __aiter__(self): ...

    async def result(self) -> AssistantMessage: ...


class AgentToolResult(TypedDict, total=False):
    content: Required[list[TextContent | ImageContent]]
    details: Any


AgentToolUpdateCallback = Callable[[AgentToolResult], None]


class AgentToolProtocol(Protocol):
    name: str
    label: str
    description: str
    parameters: Mapping[str, Any]

    async def execute(
        self,
        tool_call_id: str,
        params: Any,
        signal: SupportsAbort | None = None,
        on_update: AgentToolUpdateCallback | None = None,
    ) -> AgentToolResult: ...


AgentTool = AgentToolProtocol | Mapping[str, Any]


class AgentContext(TypedDict, total=False):
    systemPrompt: str
    messages: Required[list[AgentMessage]]
    tools: list[AgentTool]


ConvertToLlmFn = Callable[[list[AgentMessage]], Sequence[Message] | Awaitable[Sequence[Message]]]
TransformContextFn = Callable[[list[AgentMessage], SupportsAbort | None], Awaitable[list[AgentMessage]]]
GetApiKeyFn = Callable[[str], str | None | Awaitable[str | None]]
GetQueuedMessagesFn = Callable[[], Awaitable[list[AgentMessage]]]


class AgentLoopConfig(TypedDict, total=False):
    model: Required[Model]
    convertToLlm: Required[ConvertToLlmFn]
    transformContext: TransformContextFn
    getApiKey: GetApiKeyFn
    getSteeringMessages: GetQueuedMessagesFn
    getFollowUpMessages: GetQueuedMessagesFn
    reasoning: ThinkingLevel
    apiKey: str
    sessionId: str
    transport: Transport
    thinkingBudgets: ThinkingBudgets
    maxRetryDelayMs: int
    temperature: float
    maxTokens: int


class AgentStartEvent(TypedDict):
    type: Literal["agent_start"]


class AgentEndEvent(TypedDict):
    type: Literal["agent_end"]
    messages: list[AgentMessage]


class TurnStartEvent(TypedDict):
    type: Literal["turn_start"]


class TurnEndEvent(TypedDict):
    type: Literal["turn_end"]
    message: AgentMessage
    toolResults: list[ToolResultMessage]


class MessageStartEvent(TypedDict):
    type: Literal["message_start"]
    message: AgentMessage


class MessageUpdateEvent(TypedDict):
    type: Literal["message_update"]
    message: AgentMessage
    assistantMessageEvent: AssistantMessageEvent


class MessageEndEvent(TypedDict):
    type: Literal["message_end"]
    message: AgentMessage


class ToolExecutionStartEvent(TypedDict):
    type: Literal["tool_execution_start"]
    toolCallId: str
    toolName: str
    args: Any


class ToolExecutionUpdateEvent(TypedDict):
    type: Literal["tool_execution_update"]
    toolCallId: str
    toolName: str
    args: Any
    partialResult: Any


class ToolExecutionEndEvent(TypedDict):
    type: Literal["tool_execution_end"]
    toolCallId: str
    toolName: str
    result: Any
    isError: bool


AgentEvent = Union[
    AgentStartEvent,
    AgentEndEvent,
    TurnStartEvent,
    TurnEndEvent,
    MessageStartEvent,
    MessageUpdateEvent,
    MessageEndEvent,
    ToolExecutionStartEvent,
    ToolExecutionUpdateEvent,
    ToolExecutionEndEvent,
]


class SimpleStreamOptions(TypedDict, total=False):
    reasoning: ThinkingLevel
    thinkingBudgets: ThinkingBudgets
    temperature: float
    maxTokens: int
    signal: SupportsAbort
    apiKey: str
    transport: Transport
    sessionId: str
    maxRetryDelayMs: int


StreamFn = Callable[
    [Model, Context, SimpleStreamOptions | AgentLoopConfig | dict[str, Any] | None],
    AssistantMessageEventStream | StreamResponse | Awaitable[AssistantMessageEventStream | StreamResponse],
]


@dataclass(slots=True)
class AgentState:
    system_prompt: str = ""
    model: Model = field(default_factory=lambda: copy.deepcopy(DEFAULT_MODEL))
    thinking_level: ThinkingLevel = "off"
    tools: list[AgentTool] = field(default_factory=list)
    messages: list[AgentMessage] = field(default_factory=list)
    is_streaming: bool = False
    stream_message: AgentMessage | None = None
    pending_tool_calls: set[str] = field(default_factory=set)
    error: str | None = None

    @property
    def systemPrompt(self) -> str:
        return self.system_prompt

    @systemPrompt.setter
    def systemPrompt(self, value: str) -> None:
        self.system_prompt = value

    @property
    def thinkingLevel(self) -> ThinkingLevel:
        return self.thinking_level

    @thinkingLevel.setter
    def thinkingLevel(self, value: ThinkingLevel) -> None:
        self.thinking_level = value

    @property
    def isStreaming(self) -> bool:
        return self.is_streaming

    @isStreaming.setter
    def isStreaming(self, value: bool) -> None:
        self.is_streaming = value

    @property
    def streamMessage(self) -> AgentMessage | None:
        return self.stream_message

    @streamMessage.setter
    def streamMessage(self, value: AgentMessage | None) -> None:
        self.stream_message = value

    @property
    def pendingToolCalls(self) -> set[str]:
        return self.pending_tool_calls

    @pendingToolCalls.setter
    def pendingToolCalls(self, value: set[str]) -> None:
        self.pending_tool_calls = value


class AgentOptions(TypedDict, total=False):
    initialState: Mapping[str, Any] | AgentState
    convertToLlm: ConvertToLlmFn
    transformContext: TransformContextFn
    steeringMode: Literal["all", "one-at-a-time"]
    followUpMode: Literal["all", "one-at-a-time"]
    streamFn: StreamFn
    sessionId: str
    getApiKey: GetApiKeyFn
    thinkingBudgets: ThinkingBudgets
    transport: Transport
    maxRetryDelayMs: int


def now_ms() -> int:
    return int(time.time() * 1000)


def create_default_model() -> Model:
    return copy.deepcopy(DEFAULT_MODEL)


T = TypeVar("T")


def get_value(source: Mapping[str, Any] | object, *names: str, default: T | None = None) -> Any | T | None:
    for name in names:
        if isinstance(source, Mapping) and name in source:
            return source[name]
        if hasattr(source, name):
            return getattr(source, name)
    return default


def set_value(target: MutableMapping[str, Any] | object, name: str, value: Any) -> None:
    if isinstance(target, MutableMapping):
        target[name] = value
        return
    setattr(target, name, value)


def clone_message(message: AgentMessage) -> AgentMessage:
    return copy.deepcopy(message)


def clone_messages(messages: Sequence[AgentMessage]) -> list[AgentMessage]:
    return [clone_message(message) for message in messages]


def message_role(message: AgentMessage | Mapping[str, Any] | object) -> str | None:
    role = get_value(message, "role", default=None)
    return role if isinstance(role, str) else None


def message_content(message: AgentMessage | Mapping[str, Any] | object) -> Any:
    return get_value(message, "content", default=None)


def stop_reason(message: AgentMessage | Mapping[str, Any] | object) -> str | None:
    value = get_value(message, "stopReason", "stop_reason", default=None)
    return value if isinstance(value, str) else None


def model_provider(model: Model | Mapping[str, Any]) -> str:
    provider = get_value(model, "provider", default="")
    return provider if isinstance(provider, str) else ""


def ensure_usage(value: Mapping[str, Any] | None = None) -> Usage:
    if value is None:
        return {
            "input": 0,
            "output": 0,
            "cacheRead": 0,
            "cacheWrite": 0,
            "totalTokens": 0,
            "cost": {"input": 0.0, "output": 0.0, "cacheRead": 0.0, "cacheWrite": 0.0, "total": 0.0},
        }
    # Keep any caller-provided usage fields, but guarantee mandatory keys.
    usage = dict(value)
    usage.setdefault("input", 0)
    usage.setdefault("output", 0)
    usage.setdefault("cacheRead", 0)
    usage.setdefault("cacheWrite", 0)
    usage.setdefault("totalTokens", 0)
    cost = dict(usage.get("cost", {}))
    cost.setdefault("input", 0.0)
    cost.setdefault("output", 0.0)
    cost.setdefault("cacheRead", 0.0)
    cost.setdefault("cacheWrite", 0.0)
    cost.setdefault("total", 0.0)
    usage["cost"] = cost
    return usage  # type: ignore[return-value]


def is_awaitable(value: Any) -> bool:
    return inspect.isawaitable(value)


async def maybe_await(value: Any) -> Any:
    if inspect.isawaitable(value):
        return await value
    return value


def default_convert_to_llm(messages: list[AgentMessage]) -> list[Message]:
    filtered: list[Message] = []
    for message in messages:
        if message_role(message) in {"user", "assistant", "toolResult"}:
            filtered.append(message)  # type: ignore[arg-type]
    return filtered
