"""
The ``chat`` module provides a programmatic chat interface for interacting with
deployed Cortex agents.
"""
import dataclasses
from typing import List, Optional

import httpx

from relationalai.util.docutils import include_in_docs
from relationalai.util.error import exc
from .api.agent import CortexAgentApi
from .api.thread import CortexThreadApi

__include_in_docs__ = True


@dataclasses.dataclass
class InputTokens:
    total: int = 0
    cache_read: int = 0
    cache_write: int = 0
    uncached: int = 0

    @classmethod
    def from_dict(cls, data: dict) -> "InputTokens":
        return cls(
            total=data.get("total", 0),
            cache_read=data.get("cache_read", 0),
            cache_write=data.get("cache_write", 0),
            uncached=data.get("uncached", 0),
        )


@dataclasses.dataclass
class OutputTokens:
    total: int = 0

    @classmethod
    def from_dict(cls, data: dict) -> "OutputTokens":
        return cls(total=data.get("total", 0))


@dataclasses.dataclass
class ModelTokenUsage:
    model_name: str = ""
    context_window: int = 0
    input_tokens: InputTokens = dataclasses.field(default_factory=InputTokens)
    output_tokens: OutputTokens = dataclasses.field(default_factory=OutputTokens)

    @classmethod
    def from_dict(cls, data: dict) -> "ModelTokenUsage":
        return cls(
            model_name=data.get("model_name", ""),
            context_window=data.get("context_window", 0),
            input_tokens=InputTokens.from_dict(data.get("input_tokens", {})),
            output_tokens=OutputTokens.from_dict(data.get("output_tokens", {})),
        )


@dataclasses.dataclass
class Usage:
    tokens_consumed: list[ModelTokenUsage] = dataclasses.field(default_factory=list)

    @classmethod
    def from_dict(cls, data: dict) -> "Usage":
        return cls(
            tokens_consumed=[
                ModelTokenUsage.from_dict(item)
                for item in data.get("tokens_consumed", [])
            ]
        )

    @property
    def total_input_tokens(self) -> int:
        return sum(m.input_tokens.total for m in self.tokens_consumed)

    @property
    def total_output_tokens(self) -> int:
        return sum(m.output_tokens.total for m in self.tokens_consumed)

    @property
    def total_tokens(self) -> int:
        return self.total_input_tokens + self.total_output_tokens

    @property
    def total_cache_read_tokens(self) -> int:
        return sum(m.input_tokens.cache_read for m in self.tokens_consumed)

    @property
    def model_names(self) -> list[str]:
        return list({m.model_name for m in self.tokens_consumed if m.model_name})


@dataclasses.dataclass
class ToolCall:
    name: str
    arguments: dict

    @classmethod
    def from_dict(cls, data: dict):
        args = data.get("arguments") or data.get("input", {})
        return cls(
            name=data.get("name", ""),
            arguments=args
        )


@dataclasses.dataclass
class ToolResult:
    name: str
    status: str
    content: list

    def json_result(self) -> Optional[dict]:
        for item in self.content:
            if item.get("type") == "json":
                return item.get("json", {}).get("result")
        return None

    @classmethod
    def from_dict(cls, data: dict):
        return cls(
            name=data.get("name", ""),
            status=data.get("status", ""),
            content=data.get("content", [])
        )


@dataclasses.dataclass
class ContentItem:
    type: str
    role: Optional[str] = None
    content: Optional[str] = None
    tool_call: Optional[ToolCall] = None
    tool_result: Optional[ToolResult] = None
    raw: dict = dataclasses.field(default_factory=dict)

    @classmethod
    def from_dict(cls, data: dict):
        tool_call = None
        if "tool_call" in data:
            tool_call = ToolCall.from_dict(data["tool_call"])
        elif "tool_use" in data:
            tool_call = ToolCall.from_dict(data["tool_use"])

        tool_result = None
        if "tool_result" in data:
            tool_result = ToolResult.from_dict(data["tool_result"])

        # The SSE API uses "text" for text message content,
        # while the dataclass field is named "content".
        content = data.get("content") or data.get("text")

        return cls(
            type=data.get("type", ""),
            role=data.get("role"),
            content=content,
            tool_call=tool_call,
            tool_result=tool_result,
            raw=data
        )


@include_in_docs
@dataclasses.dataclass
class ChatResponse:
    """Complete response from an agent conversation turn.

    Contains all content items generated during a single agent interaction,
    including text responses, tool calls, and tool results. Provides utility
    methods for filtering, asserting, and debugging response content.

    Attributes
    ----------
    messages : list of ContentItem
        Content items in the response.
    raw : dict
        Original API response data for debugging and advanced use.
    """
    messages: list[ContentItem]
    raw: dict

    @classmethod
    def from_dict(cls, data: dict):
        """Parse a ``ChatResponse`` from API response data.

        Parameters
        ----------
        data : dict
            Dictionary containing response data with ``"content"`` field.

        Returns
        -------
        ChatResponse
            Parsed instance with messages.
        """
        content = data.get("content", [])
        messages = [ContentItem.from_dict(msg) for msg in content]
        return cls(messages=messages, raw=data)

    @include_in_docs
    def usage(self) -> Optional[Usage]:
        """Token usage data from the Cortex API for this turn, if available.

        Returns
        -------
        Usage or None
            Token usage for this turn, or ``None`` if not available.
        """
        raw_usage = self.raw.get("metadata", {}).get("usage")
        if raw_usage is None:
            return None
        return Usage.from_dict(raw_usage)

    @include_in_docs
    def tool_results(self, name: Optional[str] = None, status: Optional[str] = None) -> list[ToolResult]:
        """Filter tool results by name and/or status.

        Parameters
        ----------
        name : str, optional
            Tool name to filter by.
        status : str, optional
            Status to filter by (e.g., ``"success"``, ``"error"``).

        Returns
        -------
        list of ToolResult
            Matching tool result objects.
        """
        results = []
        for msg in self.messages:
            if msg.tool_result:
                if name and msg.tool_result.name != name:
                    continue
                if status and msg.tool_result.status != status:
                    continue
                results.append(msg.tool_result)
        return results

    @include_in_docs
    def tool_calls(self, name: Optional[str] = None) -> list[ToolCall]:
        """Filter tool calls by name.

        Parameters
        ----------
        name : str, optional
            Tool name to filter by.

        Returns
        -------
        list of ToolCall
            Matching tool call objects.
        """
        calls = []
        for msg in self.messages:
            if msg.tool_call:
                if name and msg.tool_call.name != name:
                    continue
                calls.append(msg.tool_call)
        return calls

    @include_in_docs
    def text_content(self) -> list[str]:
        """Extract all text content from messages.

        Returns
        -------
        list of str
            Text strings from all messages with content.
        """
        return [msg.content for msg in self.messages if msg.content]

    @include_in_docs
    def full_text(self) -> str:
        """Get all text content concatenated.

        Returns
        -------
        str
            All message text joined by newlines.
        """
        return "\n".join(self.text_content())

    @include_in_docs
    def tool_usage_summary(self) -> dict[str, int]:
        """Get tool call counts by tool name.

        Returns
        -------
        dict of str to int
            Dictionary mapping tool names to call counts.
        """
        from collections import Counter
        calls = [call.name for msg in self.messages if msg.tool_call for call in [msg.tool_call]]
        return dict(Counter(calls))

    def debug_print(self, include_content: bool = True):
        """Print a readable summary of the response for debugging.

        Outputs a formatted view of all messages showing message types,
        tool calls, tool results, and content previews.

        Parameters
        ----------
        include_content : bool
            Whether to include content previews. Default: ``True``.
        """
        print(f"Messages: {len(self.messages)}")
        for i, msg in enumerate(self.messages):
            print(f"  [{i}] {msg.type}", end="")
            if msg.tool_call:
                print(f" - Tool Call: {msg.tool_call.name}")
            elif msg.tool_result:
                print(f" - Tool Result: {msg.tool_result.name} ({msg.tool_result.status})")
            elif msg.content and include_content:
                preview = msg.content[:50] + "..." if len(msg.content) > 50 else msg.content
                print(f" - Content: {preview}")
            else:
                print()


@include_in_docs
class CortexAgentChat:
    """High-level client for conversing with deployed Cortex agents.

    Provides a simple interface for interacting with agents defined in
    Snowflake. The agent handles tool execution server-side, manages
    conversation state through threads, and you simply send messages and
    receive structured responses.

    Parameters
    ----------
    client : httpx.Client
        Configured ``httpx.Client`` with authentication headers.
    name : str
        Agent name to interact with.
    db : str
        Database name. Default: ``"SNOWFLAKE_INTELLIGENCE"``.
    schema : str
        Schema name. Default: ``"AGENTS"``.
    app_name : str
        Optional application identifier for thread grouping.

    Examples
    --------

    >>> chat = CortexAgentChat(client=client, name="my_agent")
    >>> response = chat.send("What can I ask about?")
    >>> print(response.full_text())
    """
    def __init__(
            self,
            client: httpx.Client,
            name: str,
            db="SNOWFLAKE_INTELLIGENCE",
            schema="AGENTS",
            app_name="",
    ):
        self._agent_api = CortexAgentApi(client)
        self._thread_api = CortexThreadApi(client)
        self._db = db
        self._schema = schema
        self._name = name
        self._app_name = app_name
        self._thread_id: Optional[int] = None
        self._parent_message_id = 0

    @include_in_docs
    def config(self) -> dict:
        """Retrieve the agent's configuration.

        Fetches the complete agent specification from the API and returns it
        as a dictionary.

        Returns
        -------
        dict
            Dictionary containing agent name and ``cortex_agent`` specification.

        Raises
        ------
        AssertionError
            If ``agent_spec`` is not present in the response.
        httpx.HTTPStatusError
            If the API request fails.
        """
        response = self._agent_api.describe(
            database=self._db,
            schema=self._schema,
            name=self._name,
        )
        assert response.agent_spec
        return {
            "name": self._name,
            "cortex_agent": dataclasses.asdict(response.agent_spec)
        }

    def _ensure_thread_id(self):
        """Ensure a thread ID exists, creating one if necessary."""
        if self._thread_id is None:
            response = self._thread_api.create(self._app_name)
            self._thread_id = response.thread_id

    @include_in_docs
    def send(self, prompt: str) -> ChatResponse:
        """Send a message to the agent and get the response.

        Creates a thread if this is the first message, then sends the prompt
        to the agent and returns the complete response including any tool calls
        and results. Automatically maintains conversation context across
        multiple ``send()`` calls.

        Parameters
        ----------
        prompt : str
            User message to send to the agent.

        Returns
        -------
        ChatResponse
            The agent's reply including tool interactions.

        Raises
        ------
        httpx.HTTPStatusError
            If the API request fails.
        AssertionError
            If the response doesn't contain expected content.
        """
        self._ensure_thread_id()
        assert self._thread_id is not None, "Thread ID should be set after _ensure_thread_id"
        data, self._parent_message_id = self._agent_api.run(
            database=self._db,
            schema=self._schema,
            name=self._name,
            thread_id=self._thread_id,
            parent_message_id=self._parent_message_id,
            prompt=prompt,
        )
        return ChatResponse.from_dict(data)

    @include_in_docs
    def thread_id(self) -> int:
        """Get the current thread ID, creating one if necessary.

        Returns
        -------
        int
            Thread ID for this conversation.
        """
        self._ensure_thread_id()
        assert self._thread_id is not None, "Thread ID should be set after _ensure_thread_id"
        return self._thread_id

    @include_in_docs
    def messages(self) -> List:
        """Retrieve all messages from the current conversation thread.

        Fetches the complete message history from the server, including both
        user and assistant messages. Messages are returned in descending order
        (newest first) by default.

        Returns
        -------
        list of CortexThreadApi.Message
            Conversation history.

        Raises
        ------
        RAIException
            If no thread exists yet (call ``send()`` first).
        httpx.HTTPStatusError
            If the API request fails.
        """
        if self._thread_id is None:
            exc(
                "NoThread",
                "No thread exists yet.",
                ["Call send() first to create a thread and send a message."],
            )

        response = self._thread_api.describe(self._thread_id)
        return response.messages
