"""Core type definitions for the pi-agents AI layer.

Port of pi-mono/packages/pi-ai/src/types.ts to Python dataclasses.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Literal, Union


# --- String literal types (represented as type aliases) ---

KnownApi = Literal[
    "openai-completions",
    "mistral-conversations",
    "openai-responses",
    "azure-openai-responses",
    "openai-codex-responses",
    "anthropic-messages",
    "bedrock-converse-stream",
    "google-generative-ai",
    "google-gemini-cli",
    "google-vertex",
]

Api = str  # KnownApi | arbitrary string

KnownProvider = Literal[
    "amazon-bedrock",
    "anthropic",
    "google",
    "google-gemini-cli",
    "google-antigravity",
    "google-vertex",
    "openai",
    "azure-openai-responses",
    "openai-codex",
    "github-copilot",
    "xai",
    "groq",
    "cerebras",
    "openrouter",
    "vercel-ai-gateway",
    "zai",
    "mistral",
    "minimax",
    "minimax-cn",
    "huggingface",
    "opencode",
    "opencode-go",
    "kimi-coding",
]

KNOWN_PROVIDERS: list[str] = [
    "amazon-bedrock",
    "anthropic",
    "google",
    "google-gemini-cli",
    "google-antigravity",
    "google-vertex",
    "openai",
    "azure-openai-responses",
    "openai-codex",
    "github-copilot",
    "xai",
    "groq",
    "cerebras",
    "openrouter",
    "vercel-ai-gateway",
    "zai",
    "mistral",
    "minimax",
    "minimax-cn",
    "huggingface",
    "opencode",
    "opencode-go",
    "kimi-coding",
]

Provider = str  # KnownProvider | arbitrary string

ThinkingLevel = Literal["minimal", "low", "medium", "high", "xhigh"]

CacheRetention = Literal["none", "short", "long"]

Transport = Literal["sse", "websocket", "auto"]

StopReason = Literal["stop", "length", "toolUse", "error", "aborted"]

InputModality = Literal["text", "image"]


# --- Data structures ---


@dataclass
class ThinkingBudgets:
    minimal: int | None = None
    low: int | None = None
    medium: int | None = None
    high: int | None = None


@dataclass
class TextContent:
    type: Literal["text"] = "text"
    text: str = ""
    text_signature: str | None = None  # textSignature in JSON


@dataclass
class ThinkingContent:
    type: Literal["thinking"] = "thinking"
    thinking: str = ""
    thinking_signature: str | None = None  # thinkingSignature in JSON
    redacted: bool | None = None


@dataclass
class ImageContent:
    type: Literal["image"] = "image"
    data: str = ""  # base64 encoded
    mime_type: str = ""  # mimeType in JSON


@dataclass
class ToolCall:
    type: Literal["toolCall"] = "toolCall"
    id: str = ""
    name: str = ""
    arguments: dict[str, Any] = field(default_factory=dict)
    thought_signature: str | None = None  # thoughtSignature in JSON


@dataclass
class Cost:
    input: float = 0.0
    output: float = 0.0
    cache_read: float = 0.0  # cacheRead in JSON
    cache_write: float = 0.0  # cacheWrite in JSON
    total: float = 0.0


@dataclass
class ModelCost:
    input: float = 0.0  # $/million tokens
    output: float = 0.0
    cache_read: float = 0.0  # cacheRead in JSON
    cache_write: float = 0.0  # cacheWrite in JSON


@dataclass
class Usage:
    input: int = 0
    output: int = 0
    cache_read: int = 0  # cacheRead in JSON
    cache_write: int = 0  # cacheWrite in JSON
    total_tokens: int = 0  # totalTokens in JSON
    cost: Cost = field(default_factory=Cost)


@dataclass
class OpenAICompletionsCompat:
    supports_store: bool | None = None
    supports_developer_role: bool | None = None
    supports_reasoning_effort: bool | None = None
    reasoning_effort_map: dict[str, str] | None = None
    supports_usage_in_streaming: bool | None = None
    max_tokens_field: Literal["max_completion_tokens", "max_tokens"] | None = None
    requires_tool_result_name: bool | None = None
    requires_assistant_after_tool_result: bool | None = None
    requires_thinking_as_text: bool | None = None
    thinking_format: Literal["openai", "zai", "qwen", "qwen-chat-template"] | None = None
    open_router_routing: dict[str, Any] | None = None
    vercel_gateway_routing: dict[str, Any] | None = None
    supports_strict_mode: bool | None = None


@dataclass
class OpenAIResponsesCompat:
    pass


@dataclass
class Model:
    id: str = ""
    name: str = ""
    api: str = ""  # Api
    provider: str = ""  # Provider
    base_url: str = ""  # baseUrl in JSON
    reasoning: bool = False
    input: list[str] = field(default_factory=lambda: ["text", "image"])
    cost: ModelCost = field(default_factory=ModelCost)
    context_window: int = 0  # contextWindow in JSON
    max_tokens: int = 0  # maxTokens in JSON
    headers: dict[str, str] | None = None
    compat: OpenAICompletionsCompat | OpenAIResponsesCompat | None = None


@dataclass
class Tool:
    name: str = ""
    description: str = ""
    parameters: dict[str, Any] = field(default_factory=dict)  # JSON Schema


@dataclass
class UserMessage:
    role: Literal["user"] = "user"
    content: str | list[TextContent | ImageContent] = ""
    timestamp: int = 0


@dataclass
class AssistantMessage:
    role: Literal["assistant"] = "assistant"
    content: list[TextContent | ThinkingContent | ToolCall] = field(default_factory=list)
    api: str = ""  # Api
    provider: str = ""  # Provider
    model: str = ""
    usage: Usage = field(default_factory=Usage)
    stop_reason: str = "stop"  # StopReason, stopReason in JSON
    error_message: str | None = None  # errorMessage in JSON
    timestamp: int = 0


@dataclass
class ToolResultMessage:
    role: Literal["toolResult"] = "toolResult"
    tool_call_id: str = ""  # toolCallId in JSON
    tool_name: str = ""  # toolName in JSON
    content: list[TextContent | ImageContent] = field(default_factory=list)
    details: Any = None
    is_error: bool = False  # isError in JSON
    timestamp: int = 0


Message = Union[UserMessage, AssistantMessage, ToolResultMessage]


@dataclass
class Context:
    system_prompt: str | None = None  # systemPrompt in JSON
    messages: list[Message] = field(default_factory=list)
    tools: list[Tool] | None = None


@dataclass
class StreamOptions:
    temperature: float | None = None
    max_tokens: int | None = None  # maxTokens in JSON
    api_key: str | None = None  # apiKey in JSON
    transport: str | None = None
    cache_retention: str | None = None  # cacheRetention in JSON
    session_id: str | None = None  # sessionId in JSON
    headers: dict[str, str] | None = None
    max_retry_delay_ms: int | None = None  # maxRetryDelayMs in JSON
    metadata: dict[str, Any] | None = None


@dataclass
class SimpleStreamOptions(StreamOptions):
    reasoning: str | None = None  # ThinkingLevel
    thinking_budgets: ThinkingBudgets | None = None  # thinkingBudgets in JSON


# --- AssistantMessageEvent types ---

AssistantMessageEvent = dict[str, Any]
# Events are represented as dicts with 'type' key for simplicity in JSON serialization.
# Types: "start", "text_start", "text_delta", "text_end",
#        "thinking_start", "thinking_delta", "thinking_end",
#        "toolcall_start", "toolcall_delta", "toolcall_end",
#        "done", "error"
