
import os
import traceback

from playwright.async_api import Frame
from playwright.async_api import Page

from Agent_E.ae.config import PROJECT_SOURCE_ROOT
from Agent_E.ae.utils.js_helper import escape_js_message
from Agent_E.ae.utils.logger import logger
from Agent_E.ae.utils.ui_messagetype import MessageType


class UIManager:
    """
    Manages the UI overlay for this application. The application uses playwright for the browser driver.
    This includes handling navigation events, showing or hiding overlays, and maintaining
    a conversation history within the UI overlay.

    Attributes:
        overlay_is_collapsed (bool): Indicates if the overlay is currently collapsed.
        conversation_history (list[dict[str, str]]): The chat history between user and system. Each entry contains 'from' and 'message' keys.
        __update_overlay_chat_history_running (bool): A flag to prevent concurrent updates to the chat history.
    """

    overlay_is_collapsed: bool = True

    overlay_processing_state: str = "init"  #init: initialised, processing: processing is ongoing, done: processing is done
    overlay_show_details:bool = True

    conversation_history:list[dict[str, str]] = []
    __update_overlay_chat_history_running: bool = False


    def __init__(self):
        """
        Initializes the UIManager instance by adding default system messages to the conversation history.
        """
        self.add_default_system_messages()


    async def handle_navigation(self, frame: Frame):
        """
        Handles navigation events by injecting JavaScript code into the frame to manage the overlay state
        and updating the overlay chat history.

        Args:
            frame (Frame): The Playwright Frame object to inject JavaScript into and manage.
        """
        try:
            await frame.wait_for_load_state("load")
            overlay_injection_file = os.path.join(PROJECT_SOURCE_ROOT, "ui", "injectOverlay.js")
            with open(overlay_injection_file, 'r') as file:  # noqa: UP015
                js_code = file.read()

            # Inject the JavaScript code into the page
            await frame.evaluate(js_code)
            js_bool = str(self.overlay_show_details).lower()
            if self.overlay_is_collapsed:
                await frame.evaluate(f"showCollapsedOverlay('{self.overlay_processing_state}', {js_bool});")
            else:
                await frame.evaluate(f"showExpandedOverlay('{self.overlay_processing_state}', {js_bool});")

            #update chat history in the overlay
            await self.update_overlay_chat_history(frame)

        except Exception as e:
            if "Frame was detached" not in str(e):
                raise e


    async def show_overlay(self, page: Page):
        """
        Displays the overlay in an expanded state on the given page if it's currently collapsed.

        Args:
            page (Page): The Playwright Page object on which to show the overlay.
        """
        if not self.overlay_is_collapsed:
            logger.debug("Overlay is already expanded, ignoring show_overlay call")
            return
        await page.evaluate("showExpandedOverlay();")
        self.overlay_is_collapsed = True


    def update_overlay_state(self, is_collapsed: bool):
        """
        Updates the state of the overlay to either collapsed or expanded.

        Args:
            is_collapsed (bool): True to collapse the overlay, False to expand it.
        """
        self.overlay_is_collapsed = is_collapsed



    async def update_overlay_show_details(self, show_details: bool, page: Page):
        """
        Updates the state of the overlay to either show steps or not.

        Args:
            show_steps (bool): True to show steps, False to hide them.
        """
        self.overlay_show_details = show_details
        await self.update_overlay_chat_history(page)


    async def update_processing_state(self, state: str, page: Page):
        """
        Updates the processing state of the overlay.

        Args:
            state (str): The processing state to update.
        """
        self.overlay_processing_state = state
        try:
            js_bool = str(self.overlay_is_collapsed).lower()
            await page.evaluate(f"updateOverlayState('{self.overlay_processing_state}', {js_bool});")
        except Exception as e:
            logger.debug(f"JavaScript error: {e}")

    async def update_overlay_chat_history(self, frame_or_page: Frame | Page):
        """
        Updates the chat history in the overlay. If the overlay is expanded and not currently being updated,
        it clears existing messages and adds them fresh from the conversation history.

        Args:
            frame_or_page (Frame | Page): The Playwright Frame or Page object to update the chat history in.
        """
        logger.debug("Updating overlay chat history")

        if self.overlay_is_collapsed:
            logger.debug("Overlay is collapsed, not updating chat history")
            return
        if self.__update_overlay_chat_history_running:
            logger.debug("update_overlay_chat_history is already running, returning" + frame_or_page.url)
            return

        self.__update_overlay_chat_history_running = True
        #update chat history in the overlay by removing all messages and adding them again fresh
        try:
            await frame_or_page.evaluate("clearOverlayMessages();")
            for message in self.conversation_history:
                safe_message = escape_js_message(message["message"])
                safe_message_type = escape_js_message(message.get("message_type", MessageType.STEP.value))
                if message["from"] == "user":
                    await frame_or_page.evaluate(f"addUserMessage({safe_message});")
                else:
                   #choose chich message types to be shown depending on UI setting
                    if self.overlay_show_details == False:  # noqa: E712
                        if message["message_type"] not in (MessageType.PLAN.value, MessageType.QUESTION.value, MessageType.ANSWER.value, MessageType.INFO.value):
                            continue
                    else:
                        if message["message_type"] not in (MessageType.PLAN.value,  MessageType.QUESTION.value , MessageType.ANSWER.value,  MessageType.INFO, MessageType.STEP.value):
                            continue

                    js_code = f"addSystemMessage({safe_message}, is_awaiting_user_response=false, message_type={safe_message_type});"
                    await frame_or_page.evaluate(js_code)
            logger.debug("Chat history updated in overlay, removing update lock flag")
        except Exception:
            traceback.print_exc()
        finally:
            self.__update_overlay_chat_history_running = False

    def clear_conversation_history(self):
        """
        Clears the conversation history.
        """
        self.conversation_history = []
        self.add_default_system_messages()

    def get_conversation_history(self):
        """
        Returns the current conversation history.

        Returns:
            list[dict[str, str]]: The conversation history.
        """
        return self.conversation_history


    def new_user_message(self, message: str):
        """
        Adds a new user message to the conversation history.

        Args:
            message (str): The message text to add.
        """

        self.conversation_history.append({"from":"user", "message":message})


    def new_system_message(self, message: str, type:MessageType=MessageType.STEP):
        """
        Adds a new system message to the conversation history.

        Args:
            message (str): The message text to add.
        """

        self.conversation_history.append({"from":"system", "message":message, "message_type":type.value})
        print(f"Adding system message: {message}")

    def add_default_system_messages(self):
        """
        Adds default system messages to the conversation history to greet the user or provide initial instructions.
        """
        pass

    async def command_completed(self, page: Page, command: str, elapsed_time: float|None = None):
        """
        Handles the completion of a command, focusing on the overlay input and indicating that the command has finished.

        Args:
            page (Page): The Playwright Page object where the command was executed.
            command (str): The command that was completed.
            elapsed_time (float | None, optional): The time taken to complete the command, if relevant.
        """
        if not self.overlay_is_collapsed:
            await page.evaluate("focusOnOverlayInput();")
            await page.evaluate("commandExecutionCompleted();")
