import openai

from ..agents.base import BaseChatAgent
from ..messages import ChatMessage
from ..runtimes import SingleThreadedRuntime


class OpenAIAssistantAgent(BaseChatAgent):
    def __init__(
        self,
        name: str,
        description: str,
        runtime: SingleThreadedRuntime,
        client: openai.AsyncClient,
        assistant_id: str,
        thread_id: str,
    ) -> None:
        super().__init__(name, description, runtime)
        self._client = client
        self._assistant_id = assistant_id
        self._thread_id = thread_id
        self._current_session_window_length = 0

    async def on_chat_message(self, message: ChatMessage) -> ChatMessage:
        print("---------------")
        print(f"{self.name} received message from {message.sender}: {message.body}")
        print("---------------")
        if message.reset:
            # Reset the current session window.
            self._current_session_window_length = 0

        # Save the message to the thread.
        _ = await self._client.beta.threads.messages.create(
            thread_id=self._thread_id,
            content=message.body,
            role="user",
            metadata={"sender": message.sender},
        )
        self._current_session_window_length += 1

        # If the message is a save_message_only message, return early.
        if message.save_message_only:
            return ChatMessage(body="OK", sender=self.name)

        # Create a run and wait until it finishes.
        run = await self._client.beta.threads.runs.create_and_poll(
            thread_id=self._thread_id,
            assistant_id=self._assistant_id,
            truncation_strategy={
                "type": "last_messages",
                "last_messages": self._current_session_window_length,
            },
        )

        if run.status != "completed":
            # TODO: handle other statuses.
            raise ValueError(f"Run did not complete successfully: {run}")

        # Get the last message from the run.
        response = await self._client.beta.threads.messages.list(self._thread_id, run_id=run.id, order="desc", limit=1)
        last_message_content = response.data[0].content

        # TODO: handle array of content.
        text_content = [content for content in last_message_content if content.type == "text"]
        if not text_content:
            raise ValueError(f"Expected text content in the last message: {last_message_content}")

        # TODO: handle multiple text content.
        return ChatMessage(body=text_content[0].text.value, sender=self.name)