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

from openai import AsyncOpenAI

from animuz_core.genai.base import BaseAgent
from animuz_core.genai.prompts import CLAUDE_RAG_DEFAULT_PROMPT, CLAUDE_TOOL_DEFAULT_PROMPT
from animuz_core.genai.tools import OpenAITool

LOGGER = init_logger(__name__)

class OpenAIAgentClient(BaseAgent):
    def __init__(self, tools: Dict[str, Tuple[OpenAITool, callable]] = None) -> None:
        self.client = AsyncOpenAI()
        self.temperature = 0.0
        self.model = "gpt-4o"
        self.max_token = 1000
        if tools:
            self.add_tools(tools)

        self.result_store = {}

    def add_tools(self, tools: Dict[str, Tuple[OpenAITool, callable]]):
        self.tools = tools
        self.tool_dump = [{"type":"function", "function": tool[0].model_dump()} for tool in self.tools.values()]

    def _filter_for_openai(self, messages: List[Dict]) -> List[Dict]:
        """Remove non-OpenAI 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) -> str:
        """
        Asynchronously retrieves a reply from the OpenAI 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.
            user_chat_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.
        """

        res = await self.create_message(messages)
        response_message = res.choices[0].message 
        tool_calls = response_message.tool_calls

        output = []

        messages.append(response_message)
        output.append(response_message)

        while tool_calls is not None:
            for tool_call in tool_calls:
                function_name = tool_call.function.name
                function_args = json.loads(tool_call.function.arguments)
                result, metadata = await self.process_tool_use(function_args, function_name, user_chat_id)

                tool_result = {
                    "tool_call_id": tool_call.id,
                    "role": "tool",
                    "name": function_name,
                    "content": result,
                    "metadata": metadata
                }
                messages.append(tool_result)
                output.append(tool_result)

            res = await self.create_message(messages)
            response_message = res.choices[0].message
            tool_calls = response_message.tool_calls

            messages.append(response_message)
            output.append(response_message)

        return output

    async def get_reply_frontend_tool_use_ack(self, messages: List[Dict[str, str]], user_chat_id: str, background_tasks: BackgroundTasks) -> str:

        res = await self.create_message(messages)
        response_message = res.choices[0].message 
        tool_calls = response_message.tool_calls

        output = []

        messages.append(response_message)
        output.append(response_message)

        if tool_calls is None:
            return output, None
        else:
            chat_id = str(uuid.uuid4()) 
            background_tasks.add_task(self.process_full_response, tool_calls, messages, user_chat_id, chat_id)
            return output, chat_id
        # message = await self.create_message(messages)
        # 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, tool_calls, messages: List[Dict[str, str]], user_chat_id: str, chat_id: str) -> str:
        output = []

        while tool_calls is not None:
            for tool_call in tool_calls:
                function_name = tool_call.function.name
                function_args = json.loads(tool_call.function.arguments)
                result, metadata = await self.process_tool_use(function_args, function_name, user_chat_id)

                tool_result = {
                    "tool_call_id": tool_call.id,
                    "role": "tool",
                    "name": function_name,
                    "content": result,
                    "metadata": metadata
                }
                messages.append(tool_result)
                output.append(tool_result)

            res = await self.create_message(messages)
            response_message = res.choices[0].message
            tool_calls = response_message.tool_calls

            messages.append(response_message)
            output.append(response_message)

        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, stream: bool= False):
        start = time.time()
        res = await self.client.chat.completions.create(
            model=self.model,
            max_tokens=self.max_token,
            temperature=self.temperature,
            messages=self._filter_for_openai(self.message_history if history is None else history),
            tools=self.tool_dump,
            stream=stream
        )

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

    async def stream_reply_frontend(self, messages: List[Dict[str, str]]):
        """Gets the response from OpenAI, updates the messages array, yields
        content, and calls functions as needed."""

        # Define variables to hold the streaming content and function call
        streaming_content = ""
        function_call = None
        content_started = False
        stream = await self.create_message(messages, stream = True)

        async for chunk in stream:
            msg = chunk.choices[0]
            # print(msg)
            if hasattr(msg, "delta"):
                if msg.delta.tool_calls:
                    if not function_call:
                        function_call = {"name": msg.delta.tool_calls[0].function.name, "arguments": "", "id": msg.delta.tool_calls[0].id}
                        yield {"type": "tool_call_start", "content": {"tool_name": msg.delta.tool_calls[0].function.name}}

                    if msg.delta.tool_calls[0].function.arguments:
                        function_call["arguments"] += msg.delta.tool_calls[0].function.arguments
                        yield {"type": "content", "content": msg.delta.tool_calls[0].function.arguments}                            

                # If it's content, add it to the streaming content and yield it
                elif hasattr(msg.delta, "content"):
                    if msg.delta.content:
                        streaming_content += msg.delta.content
                        if not content_started:
                            yield {"type":"content_start", "content": ""}
                            content_started = True
                        yield {"type":"content", "content":msg.delta.content}

            # If it's the end of the response and it's a text response, update the messages array with it
            if msg.finish_reason == "stop":
                yield {"type": "content_end", "content": streaming_content}

            elif msg.finish_reason == "tool_calls":
                yield {"type": "tool_call_end", "content": function_call}

        yield {"type": "message_end", "content": ""}

    async def get_reply_frontend_streaming(self, messages: List[Dict[str, str]], user_chat_id: str):
        output = []
        current_message = {"role": "assistant", "content": ""}
        tool_use = False
        async for chunk in self.stream_reply_frontend(messages):
            print(chunk)
            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":
                yield chunk
            elif chunk["type"] == "tool_call_end":
                yield chunk # yield tool call message

                # use tool
                tool_use = True
                result, metadata = await self.process_tool_use(json.loads(chunk["content"]["arguments"]), chunk["content"]["name"], user_chat_id)
                tool_result = {
                    "role": "function",
                    "tool_call_id": chunk["content"]["id"],
                    "content": result,
                    "name": chunk["content"]["name"],
                    "metadata": metadata
                }
                output.append(tool_result)
                messages.append({"role": "assistant", "content": None, "function_call":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):
                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(json.loads(chunk["content"]["arguments"]), chunk["content"]["name"], user_chat_id)
                    tool_result = {
                        "role": "function",
                        "tool_call_id": chunk["content"]["id"],
                        "content": result,
                        "name": chunk["content"]["name"],
                        "metadata": metadata
                    }
                    output.append(tool_result)
                    messages.append({"role": "assistant", "content": None, "function_call":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}

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