from __future__ import annotations

"""Shared types used across providers, tools, memory, and agents."""

from dataclasses import dataclass, field
from enum import Enum
from typing import Any
from uuid import uuid4


class Role(str, Enum):
    """Message role."""

    SYSTEM = "system"
    USER = "user"
    ASSISTANT = "assistant"
    TOOL = "tool"


@dataclass(slots=True)
class TextContentPart:
    """Plain text content."""

    text: str


@dataclass(slots=True)
class ImageContentPart:
    """Inline image reference."""

    data: str
    mime_type: str | None = None


@dataclass(slots=True)
class JsonContentPart:
    """Structured JSON payload."""

    value: Any


@dataclass(slots=True)
class ResourceContentPart:
    """External resource reference."""

    uri: str
    text: str | None = None


ContentPartType = (
    TextContentPart | ImageContentPart | JsonContentPart | ResourceContentPart
)


class ContentPart:
    """Namespace containing content part classes."""

    Text = TextContentPart
    Image = ImageContentPart
    Json = JsonContentPart
    Resource = ResourceContentPart


@dataclass(slots=True)
class MessageContent:
    """Simple text or rich multipart message content."""

    text_value: str | None = None
    parts: list[ContentPartType] | None = None

    @classmethod
    def from_text(cls, text: str) -> "MessageContent":
        return cls(text_value=text)

    @classmethod
    def parts_content(cls, parts: list[ContentPartType]) -> "MessageContent":
        return cls(parts=parts)

    def as_text(self) -> str:
        if self.text_value is not None:
            return self.text_value
        if not self.parts:
            return ""
        lines: list[str] = []
        for part in self.parts:
            if isinstance(part, TextContentPart):
                lines.append(part.text)
            elif isinstance(part, JsonContentPart):
                lines.append(str(part.value))
            elif isinstance(part, ResourceContentPart) and part.text is not None:
                lines.append(part.text)
        return "\n".join(lines)


@dataclass(slots=True)
class ToolCall:
    """Tool call emitted by a provider."""

    id: str
    name: str
    arguments: Any


@dataclass(slots=True)
class ToolCallDelta:
    """Partial tool-call delta used in streaming."""

    index: int
    id: str | None = None
    name: str | None = None
    arguments_delta: str | None = None


@dataclass(slots=True)
class ToolDefinition:
    """Tool-call definition sent to a provider."""

    name: str
    description: str
    schema: dict[str, Any]


@dataclass(slots=True)
class Message:
    """A chat message exchanged with an LLM."""

    role: Role
    content: MessageContent
    tool_calls: list[ToolCall] | None = None
    tool_call_id: str | None = None
    metadata: dict[str, Any] = field(default_factory=dict)

    @classmethod
    def system(cls, content: str) -> "Message":
        return cls.new(Role.SYSTEM, MessageContent.from_text(content))

    @classmethod
    def user(cls, content: str) -> "Message":
        return cls.new(Role.USER, MessageContent.from_text(content))

    @classmethod
    def assistant(cls, content: str) -> "Message":
        return cls.new(Role.ASSISTANT, MessageContent.from_text(content))

    @classmethod
    def new(cls, role: Role, content: MessageContent) -> "Message":
        return cls(role=role, content=content)

    @classmethod
    def tool_result(
        cls, tool_call_id: str, tool_name: str, output: "ToolOutput"
    ) -> "Message":
        return cls(
            role=Role.TOOL,
            content=MessageContent.from_text(output.text_content()),
            tool_call_id=tool_call_id,
            metadata={"tool_name": tool_name},
        )

    def text_content(self) -> str:
        return self.content.as_text()


class StopReason(str, Enum):
    """Reason generation stopped."""

    STOP = "stop"
    END_TURN = "end_turn"
    TOOL_USE = "tool_use"
    MAX_TOKENS = "max_tokens"
    OTHER = "other"


@dataclass(slots=True)
class Usage:
    """Token usage metadata."""

    input_tokens: int = 0
    output_tokens: int = 0
    total_tokens: int = 0


@dataclass(slots=True)
class CompletionRequest:
    """Unified completion request."""

    messages: list[Message]
    tools: list[ToolDefinition] | None = None
    model: str = ""
    temperature: float | None = None
    max_tokens: int | None = None
    stream: bool = False
    system: str | None = None
    extra: dict[str, Any] = field(default_factory=dict)


@dataclass(slots=True)
class CompletionResponse:
    """Unified completion response."""

    message: Message
    stop_reason: StopReason = StopReason.STOP
    usage: Usage = field(default_factory=Usage)
    model: str = ""
    raw: Any | None = None

    @classmethod
    def text(cls, text: str) -> "CompletionResponse":
        return cls(message=Message.assistant(text))

    @classmethod
    def tool_call(cls, name: str, arguments: Any) -> "CompletionResponse":
        message = Message.assistant("")
        message.tool_calls = [ToolCall(id=str(uuid4()), name=name, arguments=arguments)]
        return cls(message=message, stop_reason=StopReason.TOOL_USE)


@dataclass(slots=True)
class StreamChunk:
    """Streaming chunk emitted by providers."""

    delta: str = ""
    tool_call_delta: list[ToolCallDelta] | None = None
    finish_reason: StopReason | None = None


@dataclass(slots=True)
class TextToolContent:
    """Text tool output."""

    text: str


@dataclass(slots=True)
class ImageToolContent:
    """Image tool output."""

    data: str
    mime_type: str


@dataclass(slots=True)
class ResourceToolContent:
    """Resource tool output."""

    uri: str
    text: str | None = None


ToolContentType = TextToolContent | ImageToolContent | ResourceToolContent


class ToolContent:
    """Namespace containing tool content classes."""

    Text = TextToolContent
    Image = ImageToolContent
    Resource = ResourceToolContent


@dataclass(slots=True)
class ToolOutput:
    """Structured tool output payload."""

    content: list[ToolContentType]
    is_error: bool = False

    @classmethod
    def text(cls, text: str) -> "ToolOutput":
        return cls(content=[TextToolContent(text=text)], is_error=False)

    @classmethod
    def error(cls, text: str) -> "ToolOutput":
        return cls(content=[TextToolContent(text=text)], is_error=True)

    def text_content(self) -> str:
        lines: list[str] = []
        for item in self.content:
            if isinstance(item, TextToolContent):
                lines.append(item.text)
            elif isinstance(item, ResourceToolContent) and item.text is not None:
                lines.append(item.text)
        return "\n".join(lines)


@dataclass(slots=True)
class AgentOutput:
    """Final output of an agent run."""

    text: str = ""
    steps: int = 0
    usage: Usage = field(default_factory=Usage)
    messages: list[Message] = field(default_factory=list)
    metadata: dict[str, Any] = field(default_factory=dict)


@dataclass(slots=True)
class ThinkingEvent:
    """Event marking a reasoning step."""

    text: str


@dataclass(slots=True)
class ToolCallEvent:
    """Event marking a tool invocation."""

    name: str
    input: Any


@dataclass(slots=True)
class ToolResultEvent:
    """Event marking a tool result."""

    name: str
    output: str


@dataclass(slots=True)
class TokenEvent:
    """Incremental token event."""

    token: str


@dataclass(slots=True)
class DoneEvent:
    """Terminal event with final output."""

    output: AgentOutput


AgentEventType = (
    ThinkingEvent | ToolCallEvent | ToolResultEvent | TokenEvent | DoneEvent
)


class AgentEvent:
    """Namespace containing streaming event classes."""

    Thinking = ThinkingEvent
    ToolCall = ToolCallEvent
    ToolResult = ToolResultEvent
    Token = TokenEvent
    Done = DoneEvent
