"""LLM node I/O models.

Uses JsonObject and typed fields instead of dict[str, Any].
"""

from pydantic import BaseModel, Field

from core.common import JsonObject
from core.types.llm import LLMToolCall


class LLMUsage(BaseModel):
    """Token usage statistics from LLM calls."""

    input_tokens: int = 0
    output_tokens: int = 0
    tokens_used: int = 0


class LLMInput(BaseModel):
    """Input for LLM nodes."""

    model_config = {"extra": "forbid"}

    prompt: str
    model: str = "gpt-4"
    temperature: float = 0.7
    max_output_tokens: int | None = None
    system_prompt: str | None = None
    context: JsonObject = Field(default_factory=dict)


class LLMOutput(BaseModel):
    """Output from LLM nodes.

    Attributes
    ----------
        response: The generated text (also available as 'content' for convenience)
        model: The model that was used
        provider: Provider used (openrouter, ollama, etc.)
        usage: Token usage statistics (input_tokens, output_tokens, tokens_used)
        finish_reason: Provider/model finish reason when available
        tool_calls: Tool calls requested by the model (if tool calling was enabled)
        cost_usd: Estimated cost in USD when available
        metadata: Additional metadata
    """

    model_config = {"extra": "forbid"}

    response: str
    model: str
    provider: str | None = None
    usage: LLMUsage | None = None
    finish_reason: str | None = None
    tool_calls: list[LLMToolCall] | None = None
    cost_usd: float | None = None
    metadata: JsonObject = Field(default_factory=dict)

    @property
    def content(self) -> str:
        """Alias for response - for template convenience ({{node.content}})."""
        return self.response

    def model_dump(self, **kwargs: object) -> dict[str, object]:
        """Include 'content' as alias for 'response' in output."""
        data = super().model_dump(**kwargs)
        # Add content as alias for response so both {{node.response}} and {{node.content}} work
        data["content"] = data["response"]
        return data


__all__ = ["LLMInput", "LLMOutput", "LLMUsage"]
