import json

from dataclasses import dataclass, field
from typing import Optional, List, AsyncIterator, Dict, Any, Mapping, Tuple

from agentica.model.message import Message
from agentica.model.response import ModelResponse
from agentica.model.ollama.chat import Ollama, Metrics
from agentica.utils.log import logger


@dataclass
class MessageData:
    response_role: Optional[str] = None
    response_message: Optional[Dict[str, Any]] = None
    response_content: Any = ""
    response_content_chunk: str = ""
    tool_calls: List[Dict[str, Any]] = field(default_factory=list)
    tool_call_blocks: Any = field(default_factory=list)
    tool_call_chunk: str = ""
    in_tool_call: bool = False
    end_tool_call: bool = False
    response_usage: Optional[Mapping[str, Any]] = None


class Hermes(Ollama):
    """
    A class for interacting with the Hermes model via Ollama. This is a subclass of the Ollama model,
    which customizes tool call streaming for the hermes3 model.
    """

    id: str = "hermes3"
    name: str = "Hermes"
    provider: str = "Ollama"

    def handle_tool_call_chunk(self, content, tool_call_buffer, message_data) -> Tuple[str, bool]:
        """Handle a tool call chunk for response stream."""
        if content != "</tool_call>":
            tool_call_buffer += content

        if message_data.end_tool_call:
            try:
                tool_call_data = json.loads(tool_call_buffer)
                message_data.tool_call_blocks.append(tool_call_data)
                message_data.end_tool_call = False
            except json.JSONDecodeError:
                logger.error("Failed to parse tool call JSON.")
            return "", False

        return tool_call_buffer, True

    async def response_stream(self, messages: List[Message]) -> AsyncIterator[ModelResponse]:
        """Generate a streaming response from Ollama (async-only)."""
        self.sanitize_messages(messages)
        self._log_messages(messages)
        message_data = MessageData()
        metrics: Metrics = Metrics()

        metrics.response_timer.start()
        async for response in self.invoke_stream(messages=messages):
            message_data.response_message = response.get("message", {})
            if message_data.response_message:
                metrics.output_tokens += 1
                if metrics.output_tokens == 1:
                    metrics.time_to_first_token = metrics.response_timer.elapsed

                message_data.response_content_chunk = message_data.response_message.get("content", "").strip("`")
                message_data.response_content_chunk = message_data.response_message.get("content", "").strip(
                    "<|end_of_text|>"
                )
                message_data.response_content_chunk = message_data.response_message.get("content", "").strip(
                    "<|begin_of_text|>"
                )

            if message_data.response_content_chunk:
                if message_data.response_content_chunk.strip().startswith("</tool_call>"):
                    message_data.end_tool_call = True
                if message_data.in_tool_call:
                    message_data.tool_call_chunk, message_data.in_tool_call = self.handle_tool_call_chunk(
                        message_data.response_content_chunk, message_data.tool_call_chunk, message_data
                    )
                elif message_data.response_content_chunk.strip().startswith("<tool_call>"):
                    message_data.in_tool_call = True
                else:
                    yield ModelResponse(content=message_data.response_content_chunk)
                    message_data.response_content += message_data.response_content_chunk

            if response.get("done"):
                message_data.response_usage = response
        metrics.response_timer.stop()

        # Format tool calls
        if message_data.tool_call_blocks is not None:
            for block in message_data.tool_call_blocks:
                tool_name = block.get("name")
                tool_args = block.get("arguments")

                function_def = {
                    "name": tool_name,
                    "arguments": json.dumps(tool_args) if tool_args is not None else None,
                }
                message_data.tool_calls.append({"type": "function", "function": function_def})

        assistant_message = Message(role="assistant", content=message_data.response_content)

        if len(message_data.tool_calls) > 0:
            assistant_message.tool_calls = message_data.tool_calls

        self.update_usage_metrics(
            assistant_message=assistant_message, metrics=metrics, response=message_data.response_usage
        )

        messages.append(assistant_message)
        assistant_message.log()
        metrics.log()

        # Handle tool calls
        if assistant_message.tool_calls is not None and len(assistant_message.tool_calls) > 0 and self.run_tools:
            async for tool_call_response in self.handle_stream_tool_calls(assistant_message, messages):
                yield tool_call_response
            async for post_tool_call_response in self.handle_post_tool_call_messages_stream(messages=messages):
                yield post_tool_call_response
