from typing import Dict, Tuple, List
from collections import defaultdict
import asyncio, time, json, uuid
from fastapi import BackgroundTasks
from animuz_core.utils import init_logger

import anthropic

from animuz_core.genai.base import BaseLLM, BaseAgent
from animuz_core.genai.prompts import CLAUDE_RAG_DEFAULT_PROMPT, CLAUDE_TOOL_DEFAULT_PROMPT, SUMMARIZE_PROMPT, COMBINE_SUMMARIES_PROMPT
from animuz_core.genai.tools import AnthropicTool

LOGGER = init_logger(__name__)

class AnthropicClient(BaseLLM):
    def __init__(self, system_prompt: str = None) -> None:
        self.client = anthropic.AsyncClient()
        self.system_prompt = CLAUDE_RAG_DEFAULT_PROMPT if system_prompt is None else system_prompt
        self.message_history = []
        self.temperature = 0.0
        self.model = "claude-3-haiku-20240307"
        self.max_token = 1000
        self.semaphore = asyncio.Semaphore(2)

    async def __aenter__(self):
        self.client_context = await self.client.__aenter__()
        return self

    async def __aexit__(self, exc_type, exc_value, traceback):
        await self.client_context.__aexit__(exc_type, exc_value, traceback)

    async def get_reply(self, text: str) -> str:
        """
        Appends the user's input text to the message history, creates a message using the client with specified parameters, 
        and appends the assistant's response to the message history. Returns the assistant's response text.
        
        Parameters:
            text (str): The user's input text.
        
        Returns:
            str: The assistant's response text.
        """
        self.message_history.append({"role": "user", "content": text})
      
        message = await self.client.messages.create(
            model=self.model,
            max_tokens=self.max_token,
            temperature=self.temperature,
            system=self.system_prompt,
            messages=self.message_history
        )
        
        self.message_history.append({"role": "assistant", "content": message.content[-1].text})
        return message.content[-1].text
    

class AnthropicAgentClient(BaseAgent):
    def __init__(self, system_prompt: str = None, tools: Dict[str, Tuple[AnthropicTool, callable]] = None) -> None:
        self.client = anthropic.AsyncClient()
        self.system_prompt = CLAUDE_TOOL_DEFAULT_PROMPT if system_prompt is None else system_prompt
        self.temperature = 0.0
        self.model = "claude-3-haiku-20240307"
        self.max_token = 1000
        if tools:
            self.add_tools(tools)

        self.result_store = {}

    def add_tools(self, tools: Dict[str, Tuple[AnthropicTool, callable]]):
        self.tools = tools
        self.tool_dump = [tool[0].model_dump() for tool in self.tools.values()]

    def _filter_for_anthropic(self, messages: List[Dict]) -> List[Dict]:
        """Remove non-Anthropic fields (metadata) before sending to API."""
        clean = []
        for msg in messages:
            m = msg.copy()
            m.pop("metadata", None)
            clean.append(m)
        return clean

    async def get_reply_frontend(self, messages: List[Dict[str, str]], user_chat_id: str, system_prompt: str = None) -> str:
        """
        Asynchronously retrieves a reply from the Anthropic Agent using the provided messages and system prompt.

        Args:
            messages (List[Dict[str, str]]): A list of dictionaries representing the messages in the conversation.
            system_prompt (str): The system prompt to use when creating the message.
            collection_id (str): The ID of the collection to use when using the retrieval tool.

        Returns:
            str: The content of the last message in the conversation.

        This function creates a message using the provided messages and system prompt, and appends the assistant's response to the messages list.
        If the message stop reason is not "tool_use", the assistant's response is appended to the output list.
        If the message stop reason is "tool_use", the function enters a loop where it processes the tool use response,
        appends the tool result to the messages and output lists, and creates a new message.
        The assistant's response for each new message is appended to the messages and output lists.
        The function returns the content of the last message in the conversation.
        """

        message = await self.create_message(messages, system_prompt)
        output = []
        if message.stop_reason != "tool_use":
           messages.append({"role": "assistant", "content": message.content[-1].text})
           output.append({"role": "assistant", "content": message.content[-1].text})
        else:
            messages.append({"role": "assistant", "content": message.content})
            output.append({"role": "assistant", "content": message.content})

        while message.stop_reason == "tool_use":
            tool_use_response = message.content[-1]
            result, metadata = await self.process_tool_use(tool_use_response.input, tool_use_response.name, user_chat_id)

            tool_result = {
                "role": "user",
                "content": [{"type": "tool_result", "tool_use_id": tool_use_response.id, "content": result, "is_error": result == 'ERROR'}],
                "metadata": metadata
            }
            messages.append(tool_result)
            output.append(tool_result)

            message = await self.create_message(messages, system_prompt)
            messages.append({"role": "assistant", "content": message.content[-1].text})
            output.append({"role": "assistant", "content": message.content[-1].text})

        return output

    async def get_reply_frontend_tool_use_ack(self, messages: List[Dict[str, str]], user_chat_id: str, background_tasks: BackgroundTasks, system_prompt: str = None) -> str:
        message = await self.create_message(messages, system_prompt)
        output = []

        if message.stop_reason != "tool_use":
            messages.append({"role": "assistant", "content": message.content[-1].text})
            output.append({"role": "assistant", "content": message.content[-1].text})
            return output, None
        else:
            messages.append({"role": "assistant", "content": message.content})
            output.append({"role": "assistant", "content": message.content})
            chat_id = str(uuid.uuid4()) 
            background_tasks.add_task(self.process_full_response, message, messages, user_chat_id, chat_id)

            return output, chat_id

    async def process_full_response(self, message, messages: List[Dict[str, str]], user_chat_id: str, chat_id: str, system_prompt: str = None) -> str:
        output = []

        while message.stop_reason == "tool_use":
            tool_use_response = message.content[-1]
            LOGGER.debug("waiting for tool")
            result, metadata = await self.process_tool_use(tool_use_response.input, tool_use_response.name, user_chat_id)

            tool_result = {
                "role": "user",
                "content": [{"type": "tool_result", "tool_use_id": tool_use_response.id, "content": result, "is_error": result == 'ERROR'}],
                "metadata": metadata
            }
            messages.append(tool_result)
            output.append(tool_result)
            
            message = await self.create_message(messages, system_prompt)
            messages.append({"role": "assistant", "content": message.content[-1].text})
            output.append({"role": "assistant", "content": message.content[-1].text})

        self.result_store[chat_id] = output

    async def get_chat_result(self, chat_id: str):
        if chat_id in self.result_store:
            result = self.result_store[chat_id]
            del self.result_store[chat_id]  # Clean up the stored result
            return result
        else:
            return "Result not ready yet"

    async def create_message(self, history: List[Dict] = None, system_prompt: str = None):
        start = time.time()
        res = await self.client.messages.create(
            model=self.model,
            max_tokens=self.max_token,
            temperature=self.temperature,
            system=self.system_prompt if system_prompt is None else system_prompt,
            messages=self._filter_for_anthropic(self.message_history if history is None else history),
            tools=self.tool_dump
        )

        LOGGER.debug(json.dumps({
            "timestamp": time.strftime('[%Y-%b-%d %H:%M:%S %z]'),
            "message": "Anthropic API latency",
            "latency": round(time.time() - start, 4),
            "exception": None
        }))
        return res

    async def stream_reply_frontend(self, messages: List[Dict[str, str]], system_prompt: str = None):
        current_tool_call = None
        tool_input_buffer = ""

        async with self.client.messages.stream(
            model=self.model,
            max_tokens=self.max_token,
            temperature=self.temperature,
            system=self.system_prompt if system_prompt is None else system_prompt,
            messages=self._filter_for_anthropic(messages),
            tools=self.tool_dump
        ) as stream:
            async for chunk in stream:
                if chunk.type == "content_block_start":
                    if chunk.content_block.type == "tool_use":
                        current_tool_call = {"id": chunk.content_block.id, "name": chunk.content_block.name} 
                        tool_input_buffer = ""
                        yield {"type": "tool_call_start", "content": {"tool_name": current_tool_call["name"]}}
                    else:
                        # start of text
                        yield {"type": "content_start", "content": ""}

                elif chunk.type == "content_block_delta":
                    # Streaming text content
                    if current_tool_call:
                        tool_input_buffer += chunk.delta.partial_json
                        yield {"type": "content", "content": chunk.delta.partial_json}
                    else:
                        yield {"type": "content", "content": chunk.delta.text}
                 
                elif chunk.type == "content_block_stop":
                    # End of content block
                    if current_tool_call:
                        yield {
                            "type": "tool_call_end",
                            "content": chunk.content_block,
                        }
                        current_tool_call = None
                        tool_input_buffer = ""

                    else:
                        yield {"type": "content_end", "content": ""}
                elif chunk.type == "message_stop":
                    # End of message
                    yield {"type": "message_end", "content": ""}
            
    async def get_reply_frontend_streaming(self, messages: List[Dict[str, str]], user_chat_id: str, system_prompt: str = None):
        output = []
        current_message = {"role": "assistant", "content": ""}
        tool_use = False
        async for chunk in self.stream_reply_frontend(messages, system_prompt):
            if chunk["type"] == "content":
                current_message["content"] += chunk["content"]
                yield chunk
            elif chunk["type"] == "content_start":
                yield chunk
            elif chunk["type"] == "content_end":
                yield chunk
            elif chunk["type"] == "tool_call_start":
                # current_message["content"] += chunk["content"]
                yield chunk
            elif chunk["type"] == "tool_call_end":
                chunk["content"] = chunk["content"].to_dict()
                yield chunk # yield tool call message

                # use tool
                tool_use = True
                result, metadata = await self.process_tool_use(chunk["content"]["input"], chunk["content"]["name"], user_chat_id)
                tool_result = {
                    "role": "user",
                    "content": [{
                        "type": "tool_result",
                        "tool_use_id": chunk["content"]["id"],
                        "content": result,
                        "is_error": result == "ERROR"
                    }],
                    "metadata": metadata
                }
                output.append(tool_result)
                messages.append({"role": "assistant", "content": [chunk["content"]]}) # add tool call
                messages.append(tool_result) # add tool result

                yield {"type": "tool_result", "content": tool_result} # yield tool result message
            elif chunk["type"] == "message_end":
                if current_message["content"]:
                    output.append(current_message)
                yield {"type": "end", "content": output}

        while tool_use:
            current_message = {"role": "assistant", "content": ""}
            tool_use = False
            async for chunk in self.stream_reply_frontend(messages, user_chat_id, system_prompt):
                if chunk["type"] == "content":
                    current_message["content"] += chunk["content"]
                    yield chunk
                elif chunk["type"] == "content_start":
                    yield chunk
                elif chunk["type"] == "content_end":
                    yield chunk
                elif chunk["type"] == "tool_call_start":
                    current_message["content"] += chunk["content"]
                    yield chunk
                elif chunk["type"] == "tool_call_end":
                    chunk["content"] = chunk["content"].to_dict()
                    yield chunk # yield tool call message

                    tool_use = True
                    result, metadata = await self.process_tool_use(chunk["content"]["input"], chunk["content"]["name"], user_chat_id)
                    tool_result = {
                        "role": "user",
                        "content": [{
                            "type": "tool_result",
                            "tool_use_id": chunk["content"]["id"],
                            "content": result,
                            "is_error": result == "ERROR"
                        }],
                        "metadata": metadata
                    }
                    output.append(tool_result)
                    messages.append({"role": "assistant", "content": [chunk["content"].to_dict()]}) # add tool call
                    messages.append(tool_result) # add tool result

                    yield {"type": "tool_result", "content": tool_result} # yield tool result message
                elif chunk["type"] == "message_end":
                    if current_message["content"]:
                        output.append(current_message)
                    yield {"type": "end", "content": output}

        if output:
            yield {"type": "final_output", "content": output}
