import json
import queue

from typing import Dict, Optional

from PySide6.QtCore import Slot, QTimer, QPropertyAnimation
from PySide6.QtWidgets import QSpacerItem, QSizePolicy
from PySide6.QtCore import Qt

from airunner.enums import SignalCode, LLMActionType, ModelType, ModelStatus
from airunner.widgets.base_widget import BaseWidget
from airunner.widgets.llm.templates.chat_prompt_ui import Ui_chat_prompt
from airunner.widgets.llm.message_widget import MessageWidget
from airunner.data.models import Conversation
from airunner.utils.strip_names_from_message import strip_names_from_message
from airunner.utils import create_worker
from airunner.handlers.llm.llm_request import LLMRequest
from airunner.handlers.llm.llm_response import LLMResponse
from airunner.workers import LLMResponseWorker
from airunner.settings import AIRUNNER_ART_ENABLED


class ChatPromptWidget(BaseWidget):
    widget_class_ = Ui_chat_prompt
    
    def __init__(self, *args, **kwargs):
        self.signal_handlers = {
            SignalCode.AUDIO_PROCESSOR_RESPONSE_SIGNAL: self.on_hear_signal,
            SignalCode.SET_CONVERSATION: self.on_set_conversation,
            SignalCode.MODEL_STATUS_CHANGED_SIGNAL: self.on_model_status_changed_signal,
            SignalCode.CHATBOT_CHANGED: self.on_chatbot_changed,
            SignalCode.CONVERSATION_DELETED: self.on_delete_conversation,
            SignalCode.LLM_CLEAR_HISTORY_SIGNAL: self.on_clear_conversation,
            SignalCode.LOAD_CONVERSATION: self.on_load_conversation,
            SignalCode.LLM_TOKEN_SIGNAL: self.on_token_signal,
            SignalCode.LLM_TEXT_STREAMED_SIGNAL: self.on_add_bot_message_to_conversation,
        }
        self.splitters = [
            "chat_prompt_splitter"
        ]
        super().__init__()
        self.token_buffer = []
        self.ui_update_timer = QTimer(self)
        self.ui_update_timer.setInterval(50)
        self.ui_update_timer.timeout.connect(self.flush_token_buffer)
        self.ui_update_timer.start()
        self.registered: bool = False
        self.scroll_bar = None
        self.is_modal = True
        self.generating = False
        self.prefix = ""
        self.prompt = ""
        self.suffix = ""
        self.spacer = None
        self.promptKeyPressEvent = None
        self.originalKeyPressEvent = None
        self.action_menu_displayed = None
        self.action_menu_displayed = None
        self.messages_spacer = None
        self.chat_loaded = False
        self._conversation: Optional[Conversation] = None
        self._conversation_id: Optional[int] = None
        self.conversation_history = []

        self.ui.action.blockSignals(True)
        self.ui.action.addItem("Auto")
        self.ui.action.addItem("Chat")
        self.ui.action.addItem("RAG")
        
        if AIRUNNER_ART_ENABLED:
            self.ui.action.addItem("Image")
        
        action = self.action
        if action is LLMActionType.APPLICATION_COMMAND:
            self.ui.action.setCurrentIndex(0)
        elif action is LLMActionType.CHAT:
            self.ui.action.setCurrentIndex(1)
        elif action is LLMActionType.GENERATE_IMAGE:
            self.ui.action.setCurrentIndex(2)
        elif action is LLMActionType.PERFORM_RAG_SEARCH:
            self.ui.action.setCurrentIndex(3)
        elif action is LLMActionType.STORE_DATA:
            self.ui.action.setCurrentIndex(4)
        self.ui.action.blockSignals(False)
        self.originalKeyPressEvent = None
        self.originalKeyPressEvent = self.ui.prompt.keyPressEvent
        self.held_message = None
        self._disabled = False
        self.scroll_animation = None
        self._llm_response_worker = create_worker(
            LLMResponseWorker,
            sleep_time_in_ms=1
        )
        self.load_conversation()

    @property
    def conversation(self) -> Optional[Conversation]:
        return self._conversation
    
    @conversation.setter
    def conversation(self, val: Optional[Conversation]):
        self._conversation = val
        self._conversation_id = val.id if val else None

    @property
    def conversation_id(self) -> Optional[int]:
        return self._conversation_id
    
    @conversation_id.setter
    def conversation_id(self, val: Optional[int]):
        self._conversation_id = val
        if val:
            self._conversation = Conversation.objects.filter_by_first(
                id=val
            )
        else:
            self._conversation = None

    def on_load_conversation(self, data):
        self.load_conversation(data["conversation_id"])
    
    def on_delete_conversation(self, data):
        if self.conversation_id == data["conversation_id"] or self.conversation_id is None:
            self._clear_conversation_widgets()
        self.conversation = None

    def load_conversation(self, index: Optional[int] = None):
        conversation = None
        if index is not None:
            conversation = Conversation.objects.get(index)
        
        if conversation is None:
            conversation = Conversation.objects.order_by(
                Conversation.id.desc()
            ).first()
        if conversation is not None:
            self.conversation = conversation
            self.emit_signal(SignalCode.LLM_CLEAR_HISTORY_SIGNAL, {
                "conversation_id": self.conversation_id
            })
            self.on_clear_conversation()
            self._set_conversation_widgets([
                {
                    "name": (
                        self.user.username
                        if message["role"] == "user"
                        else self.chatbot.name
                    ),
                    "content": message["blocks"][0]["text"],
                    "is_bot": message["role"] == "assistant",
                    "id": message_id
                } for message_id, message in enumerate(self.conversation.value or [])
            ])

    @Slot(str)
    def handle_token_signal(self, val: str):
        if val != "[END]":
            text = self.ui.conversation.toPlainText()
            text += val
            self.ui.conversation.setPlainText(text)
        else:
            self.stop_progress_bar()
            self.generating = False
            self.enable_send_button()

    def on_model_status_changed_signal(self, data):
        if data["model"] == ModelType.LLM:
            self.chat_loaded = data["status"] is ModelStatus.LOADED

        if not self.chat_loaded:
            self.disable_send_button()
        else:
            self.enable_send_button()

    def on_chatbot_changed(self):
        self.emit_signal(SignalCode.LLM_CLEAR_HISTORY_SIGNAL)
        self.on_clear_conversation()

    def on_set_conversation(self, message):
        self._clear_conversation_widgets()
        if len(message["messages"]) > 0:
            conversation_id = message["messages"][0]["conversation_id"]
            self.conversation = Conversation.objects.filter_by_first(
                id=conversation_id
            )
        self._set_conversation_widgets([{
                "name": message["additional_kwargs"]["name"],
                "content": message["text"],
                "is_bot": message["role"] == "assistant",
                "id": index
            } for index, message in enumerate(json.loads(self.conversation.value))
        ])

    def _set_conversation_widgets(self, messages):
        for message in messages:
            self.add_message_to_conversation(
                name=message["name"],
                message=message["content"],
                is_bot=message["is_bot"],
                first_message=True,
            )
        QTimer.singleShot(100, self.scroll_to_bottom)

    def on_hear_signal(self, data: Dict):
        transcription = data["transcription"]
        self.prompt = transcription
        self.do_generate()

    def on_add_to_conversation_signal(self, name, text, is_bot):
        self.add_message_to_conversation(name=name, message=text, is_bot=is_bot)

    def on_add_bot_message_to_conversation(self, data: Dict):
        llm_response: LLMResponse = data.get("response", None)
        if not llm_response:
            raise ValueError("No LLMResponse object found in data")
    
        self.token_buffer.append(llm_response.message)

        if llm_response.is_first_message:
            self.stop_progress_bar()

        if llm_response.is_end_of_message:
            self.enable_generate()
    
    def flush_token_buffer(self):
        """
        Flush the token buffer and update the UI.
        """

        combined_message = "".join(self.token_buffer)
        self.token_buffer.clear()

        if combined_message != "":
            self.add_message_to_conversation(
                name=self.chatbot.name,
                message=combined_message,
                is_bot=True,
                first_message=False
            )

    def enable_generate(self):
        self.generating = False
        if self.held_message is not None:
            self.do_generate(prompt_override=self.held_message)
            self.held_message = None
        self.enable_send_button()

    @Slot()
    def action_button_clicked_clear_conversation(self):
        self.emit_signal(SignalCode.LLM_CLEAR_HISTORY_SIGNAL)
    
    def on_clear_conversation(self):
        self._clear_conversation()

    def _clear_conversation(self):
        self.conversation_history = []
        self._clear_conversation_widgets()

    def _clear_conversation_widgets(self):
        # Iterate through the layout items and delete the widgets
        while self.ui.scrollAreaWidgetContents.layout().count():
            item = self.ui.scrollAreaWidgetContents.layout().takeAt(0)
            widget = item.widget()
            if widget is not None:
                widget.deleteLater()
        # Ensure the layout is updated
        self.ui.scrollAreaWidgetContents.layout().update()
    
    @Slot(bool)
    def action_button_clicked_send(self):
        self.do_generate()

    def interrupt_button_clicked(self):
        self.emit_signal(SignalCode.INTERRUPT_PROCESS_SIGNAL)
        self.stop_progress_bar()
        self.generating = False
        self.enable_send_button()

    @property
    def action(self) -> LLMActionType:
        return LLMActionType[self.llm_generator_settings.action]

    def do_generate(self, prompt_override=None):
        prompt = self.prompt if (prompt_override is None or prompt_override == "") else prompt_override
        if prompt is None or prompt == "":
            self.logger.warning("Prompt is empty")
            return

        if self.generating:
            if self.held_message is None:
                self.held_message = prompt
                self.disable_send_button()
                self.interrupt_button_clicked()
            return
        self.generating = True

        widget = self.add_message_to_conversation(
            name=self.user.username,
            message=self.prompt,
            is_bot=False
        )
        if widget is not None:
            QTimer.singleShot(100, widget.set_content_size)
            QTimer.singleShot(100, self.scroll_to_bottom)

        self.clear_prompt()
        self.start_progress_bar()
        self.emit_signal(
            SignalCode.LLM_TEXT_GENERATE_REQUEST_SIGNAL,
            {
                "llm_request": True,
                "request_data": {
                    "action": self.action,
                    "prompt": prompt,
                    "llm_request": LLMRequest.from_default()
                }
            }
        )

    def on_token_signal(self, val):
        self.handle_token_signal(val)

    def showEvent(self, event):
        super().showEvent(event)
        if not self.registered:
            self.registered = True

        # handle return pressed on QPlainTextEdit
        # there is no returnPressed signal for QPlainTextEdit
        # so we have to use the keyPressEvent
        self.promptKeyPressEvent = self.ui.prompt.keyPressEvent

        # Override the method
        self.ui.prompt.keyPressEvent = self.handle_key_press

        if not self.chat_loaded:
            self.disable_send_button()

        try:
            self.ui.conversation.hide()
        except RuntimeError as e:
            self.logger.warning(f"Error hiding conversation: {e}")

        try:
            self.ui.chat_container.show()
        except RuntimeError as e:
            if AIRUNNER_ART_ENABLED:
                self.logger.warning(f"Error setting SD status text: {e}")

    def llm_action_changed(self, val: str):
        if val == "Chat":
            llm_action_value = LLMActionType.CHAT
        elif val == "Image":
            llm_action_value = LLMActionType.GENERATE_IMAGE
        elif val == "RAG":
            llm_action_value = LLMActionType.PERFORM_RAG_SEARCH
        elif val == "Auto":
            llm_action_value = LLMActionType.APPLICATION_COMMAND
        else:
            llm_action_value = LLMActionType.APPLICATION_COMMAND
        self.update_llm_generator_settings("action", llm_action_value.name)

    def prompt_text_changed(self):
        self.prompt = self.ui.prompt.toPlainText()

    def clear_prompt(self):
        self.ui.prompt.setPlainText("")

    def start_progress_bar(self):
        self.ui.progressBar.setRange(0, 0)
        self.ui.progressBar.setValue(0)

    def stop_progress_bar(self):
        self.ui.progressBar.setRange(0, 1)
        self.ui.progressBar.setValue(1)
        self.ui.progressBar.reset()

    def disable_send_button(self):
        self.ui.send_button.setEnabled(False)
        self._disabled = True

    def enable_send_button(self):
        self.ui.send_button.setEnabled(True)
        self._disabled = False

    def response_text_changed(self):
        pass

    def handle_key_press(self, event):
        if event.key() == Qt.Key.Key_Return:
            if not self._disabled and event.modifiers() != Qt.KeyboardModifier.ShiftModifier:
                self.do_generate()
                return
        # Call the original method
        if self.originalKeyPressEvent is not None and self.originalKeyPressEvent != self.handle_key_press:
            self.originalKeyPressEvent(event)

    def hide_action_menu(self):
        self.action_menu_displayed = False
        self.ui.action_menu.hide()

    def display_action_menu(self):
        self.action_menu_displayed = True
        self.ui.action_menu.show()

    def insert_newline(self):
        self.ui.prompt.insertPlainText("\n")

    def add_message_to_conversation(
        self,
        name: str,
        message: str,
        is_bot: bool, 
        first_message: bool = True,
        _message_id: Optional[int] = None
    ):
        message = strip_names_from_message(
            message.lstrip() if first_message else message, 
            self.user.username, 
            self.chatbot.botname
        )
        if not first_message:
            # get the last widget from the scrollAreaWidgetContents.layout()
            # and append the message to it. must be a MessageWidget object
            # must start at the end of the layout and work backwards
            for i in range(self.ui.scrollAreaWidgetContents.layout().count() - 1, -1, -1):
                current_widget = self.ui.scrollAreaWidgetContents.layout().itemAt(i).widget()
                if isinstance(current_widget, MessageWidget):
                    if current_widget.is_bot:
                        if message != "":
                            current_widget.update_message(message)
                        self.scroll_to_bottom()
                        return
                    break

        self.remove_spacer()

        widget = None
        if message != "":
            total_widgets = self.ui.scrollAreaWidgetContents.layout().count() - 1
            if total_widgets < 0:
                total_widgets = 0
            widget = MessageWidget(
                name=name,
                message=message,
                is_bot=is_bot,
                message_id=total_widgets,
                conversation_id=self.conversation_id
            )
            self.ui.scrollAreaWidgetContents.layout().addWidget(widget)

        self.add_spacer()

        # automatically scroll to the bottom of the scrollAreaWidgetContents
        self.scroll_to_bottom()

        return widget

    def remove_spacer(self):
        if self.spacer is not None:
            self.ui.scrollAreaWidgetContents.layout().removeItem(self.spacer)

    def add_spacer(self):
        if self.spacer is None:
            self.spacer = QSpacerItem(20, 40, QSizePolicy.Policy.Minimum, QSizePolicy.Policy.Expanding)
        self.ui.scrollAreaWidgetContents.layout().addItem(self.spacer)

    def message_type_text_changed(self, val):
        self.update_llm_generator_settings("message_type", val)

    def action_button_clicked_generate_characters(self):
        pass

    def scroll_to_bottom(self):
        if self.scroll_bar is None:
            self.scroll_bar = self.ui.chat_container.verticalScrollBar()

        if self.scroll_animation is None:
            self.scroll_animation = QPropertyAnimation(self.scroll_bar, b"value")
            self.scroll_animation.setDuration(500)

        # Stop any ongoing animation
        if self.scroll_animation and self.scroll_animation.state() == QPropertyAnimation.State.Running:
            self.scroll_animation.stop()

        self.scroll_animation.setStartValue(self.scroll_bar.value())
        self.scroll_animation.setEndValue(self.scroll_bar.maximum())
        self.scroll_animation.start()
