import datetime as _dt
from dataclasses import dataclass, field
from typing import Any, Literal, Optional

from pydantic import BaseModel


class ContentBlock(BaseModel):
    """A single piece of content within a message."""

    type: Literal["text", "tool_call", "tool_result"]
    text: Optional[str] = None
    tool_name: Optional[str] = None
    tool_args: Optional[dict[str, Any]] = None
    tool_output: Optional[str] = None

    def display_text(self) -> str:
        if self.type == "text":
            return self.text or ""
        if self.type == "tool_call":
            args = ", ".join(f"{k}={v!r}" for k, v in (self.tool_args or {}).items())
            return f"{self.tool_name or 'unknown'}({args})"
        if self.type == "tool_result":
            return self.tool_output or ""
        return ""


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

    role: str
    content: str | list[ContentBlock]
    timestamp: Optional[_dt.datetime] = None

    @classmethod
    def from_openai(cls, msg: dict[str, Any]) -> "Message":
        """Parse an OpenAI chat-format message dict.

        Handles both simple ``{"role": "...", "content": "..."}`` and
        multi-part content arrays with tool calls.
        """
        role = msg.get("role", "user")
        raw = msg.get("content", "")

        if msg.get("tool_call_id"):
            return cls(
                role=role,
                content=[ContentBlock(type="tool_result", tool_output=str(raw))],
            )

        if isinstance(raw, str) and not msg.get("tool_calls"):
            return cls(role=role, content=raw)

        blocks: list[ContentBlock] = []
        if isinstance(raw, list):
            for part in raw:
                if isinstance(part, str):
                    blocks.append(ContentBlock(type="text", text=part))
                elif isinstance(part, dict):
                    ptype = part.get("type", "text")
                    if ptype == "text":
                        blocks.append(ContentBlock(type="text", text=part.get("text", "")))

        for tc in msg.get("tool_calls", []):
            fn = tc.get("function", {})
            import json

            try:
                args = json.loads(fn.get("arguments", "{}"))
            except (json.JSONDecodeError, TypeError):
                args = {}
            blocks.append(ContentBlock(type="tool_call", tool_name=fn.get("name"), tool_args=args))

        return cls(role=role, content=blocks if blocks else str(raw or ""))

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "Message":
        """Parse a plain ``{role, content}`` dict."""
        return cls(
            role=data.get("role", "user"),
            content=data.get("content", ""),
            timestamp=data.get("timestamp"),
        )

    def to_dict(self) -> dict[str, Any]:
        d: dict[str, Any] = {"role": self.role}
        if isinstance(self.content, str):
            d["content"] = self.content
        else:
            d["content"] = [block.model_dump(exclude_none=True) for block in self.content]
        if self.timestamp:
            d["timestamp"] = self.timestamp.isoformat()
        return d

    def plain_text(self) -> str:
        """Return a plain-text rendering of the message content."""
        if isinstance(self.content, str):
            return self.content
        return "\n".join(block.display_text() for block in self.content if block.display_text())


class SummaryRecord(BaseModel):
    """A summary produced by the hierarchical summarization pipeline."""

    summary: str
    message_count: int
    date: Optional[_dt.date] = None
    hour: Optional[int] = None
    minute: Optional[int] = None
    week_start_date: Optional[_dt.date] = None
    year: Optional[int] = None
    month: Optional[int] = None


@dataclass
class LLMResponse:
    """Response from an LLM completion call."""

    text: str
    input_tokens: int = 0
    output_tokens: int = 0


@dataclass
class MemoryData:
    """Raw data fetched from the store for building memory context.

    A structured bundle of summaries and messages that the caller can
    inspect or modify before rendering into a context string.
    """

    monthly_summaries: list[SummaryRecord] = field(default_factory=list)
    weekly_summaries: list[SummaryRecord] = field(default_factory=list)
    daily_summaries: list[SummaryRecord] = field(default_factory=list)
    hourly_summaries: list[SummaryRecord] = field(default_factory=list)
    five_minute_summaries: list[SummaryRecord] = field(default_factory=list)
    messages: list[Message] = field(default_factory=list)
    last_summarized_time: Optional[_dt.datetime] = None
