"""This module contains the classes for the context and context message"""
import datetime
import uuid
import asyncio
from typing import List, Union, TYPE_CHECKING

from fastapi import Request
from litellm.types.completion import ChatCompletionMessageParam

from npiai.core import callback


# if TYPE_CHECKING:
#     # from npiai.core import BaseTool


class Task:
    """the message wrapper of a message in the context"""

    def __init__(self, goal: str) -> None:
        self.task_id = str(uuid.uuid4())
        self.goal = goal
        self.born_at = datetime.datetime.now()
        self.dialogues: List[Union[ChatCompletionMessageParam]] = []
        self.response: str
        self._session: 'Context' | None = None

    async def step(self, msgs: List[Union[ChatCompletionMessageParam]]) -> None:
        """add a message to the context"""
        self.dialogues.extend(msgs)

        for msg in msgs:
            if msg.get('content'):
                await self._session.send(callback.Callable(msg.get('content')))

    def set_session(self, session: 'Context') -> None:
        self._session = session

    def conversations(self) -> List[ChatCompletionMessageParam]:
        """return the raw message"""
        return self.dialogues.copy()


class Context:
    __last_screenshot: str | None = None

    def __init__(self, req: Request | None = None) -> None:
        # self._parent_session = parent
        # self._task = task
        self.id = str(uuid.uuid4())
        self.q = asyncio.Queue()
        self.cb_dict: dict[str, callback.Callable] = {}
        self.__is_finished = False
        self.__result: str = ''
        self.__is_failed = False
        self.__failed_msg: str = ''
        self.__active_tool: Union[BaseTool, None] = None
        self.__request = req

    def authorization(self) -> str:
        return self.__request.headers.get('Authorization')

    def entry(self):
        pass

    def exit(self):
        pass

    def with_task(self, task: Task) -> 'Context':
        task.set_session(self)
        return self

    def fork(self, task: Task) -> 'Context':
        """fork a child message, typically when call a new tools"""
        return self.with_task(task)

    async def refresh_screenshot(self) -> str | None:
        """
        Refresh and return the latest screenshot of the running tools.

        Returns:
            None if no screenshot is available or the screenshot stays unchanged.
            Otherwise, return the latest screenshot.
        """
        if not self.__active_tool or self.is_finished():
            return None

        screenshot = await self.__active_tool.get_screenshot()

        if screenshot == self.__last_screenshot:
            return None

        self.__last_screenshot = screenshot

        return screenshot

    async def send(self, cb: callback.Callable) -> None:
        """send a message to the context"""
        self.cb_dict[cb.id()] = cb
        await self.q.put(cb)

    async def fetch(self) -> callback.Callable:
        """receive a message"""
        while not self.is_failed() and not self.is_finished():
            try:
                item = self.q.get_nowait()
                if item:
                    return item
            except asyncio.QueueEmpty:
                await asyncio.sleep(0.01)

    def get_callback(self, cb_id: str) -> callback.Callable:
        return self.cb_dict[cb_id]

    def finish(self, msg: str):
        self.__result = msg
        self.__is_finished = True
        self.__last_screenshot = None
        self.q.task_done()

    def failed(self, msg: str):
        self.__failed_msg = msg
        self.__is_failed = True
        self.q.task_done()

    def is_finished(self) -> bool:
        return self.__is_finished

    def get_result(self) -> str:
        return self.__result

    def is_failed(self) -> bool:
        return self.__is_failed

    def get_failed_msg(self) -> str:
        return self.__failed_msg

    def retrieve(self, msg: str) -> str:
        pass

    @classmethod
    def from_request(cls, req: Request) -> 'Context':
        return Context(req=req)

    def bind(self, tool):
        self.__active_tool = tool

    def get_tool(self):
        return self.__active_tool
