"""Shared types used across the KOAN framework."""

from __future__ import annotations

from typing import Any, Literal

from pydantic import BaseModel, Field


class Message(BaseModel):
    """A single message in a conversation."""

    role: Literal["system", "user", "assistant", "tool"]
    content: str
    name: str | None = None
    tool_call_id: str | None = None
    tool_calls: list[dict[str, Any]] | None = None


class Usage(BaseModel):
    """Token usage from a single LLM call."""

    input_tokens: int = 0
    output_tokens: int = 0

    @property
    def total_tokens(self) -> int:
        return self.input_tokens + self.output_tokens


class LLMResponse(BaseModel):
    """Normalized response from any LLM provider."""

    text: str
    model: str
    provider: str
    usage: Usage = Field(default_factory=Usage)
    stop_reason: str | None = None
    raw: dict | None = Field(default=None, exclude=True)


class StreamDelta(BaseModel):
    """A single chunk from a streaming LLM response."""

    text: str = ""
    model: str | None = None
    provider: str | None = None
    finish_reason: str | None = None
    usage: Usage | None = None

    # Tool call accumulation fields
    tool_call_id: str | None = None
    tool_name: str | None = None
    tool_args_delta: str | None = None
