# Standard Packages
from enum import Enum
from pathlib import Path
from copy import deepcopy
import webbrowser

# External Packages
from PyQt6 import QtGui, QtWidgets
from PyQt6.QtCore import Qt, QThread, QObject, pyqtSignal

# Internal Packages
from khoj.configure import configure_server
from khoj.interface.desktop.file_browser import FileBrowser
from khoj.interface.desktop.labelled_text_field import LabelledTextField
from khoj.utils import constants, state, yaml as yaml_utils
from khoj.utils.cli import cli
from khoj.utils.config import SearchType, ProcessorType
from khoj.utils.helpers import merge_dicts, resolve_absolute_path


class MainWindow(QtWidgets.QMainWindow):
    """Create Window to Configure Khoj
    Allow user to
    1. Configure content types to search
    2. Configure conversation processor
    3. Save the configuration to khoj.yml
    """

    def __init__(self, config_file: Path):
        super(MainWindow, self).__init__()
        self.config_file = config_file
        # Set regenerate flag to regenerate embeddings everytime user clicks configure
        if state.cli_args:
            state.cli_args += ["--regenerate"]
        else:
            state.cli_args = ["--regenerate"]

        # Load config from existing config, if exists, else load from default config
        if resolve_absolute_path(self.config_file).exists():
            self.first_run = False
            self.current_config = yaml_utils.load_config_from_file(self.config_file)
        else:
            self.first_run = True
            self.current_config = deepcopy(constants.default_config)
        self.new_config = self.current_config

        # Initialize Configure Window
        self.setWindowTitle("Khoj")
        self.setFixedWidth(600)

        # Set Window Icon
        icon_path = constants.web_directory / "assets/icons/favicon-128x128.png"
        self.setWindowIcon(QtGui.QIcon(f"{icon_path.absolute()}"))

        # Initialize Configure Window Layout
        self.wlayout = QtWidgets.QVBoxLayout()

        # Add Settings Panels for each Search Type to Configure Window Layout
        self.search_settings_panels = []
        for search_type in SearchType:
            current_content_config = self.current_config["content-type"].get(
                search_type, None
            ) or self.get_default_config(search_type=search_type)
            self.search_settings_panels += [self.add_settings_panel(current_content_config, search_type)]
        # Add Conversation Processor Panel to Configure Screen
        self.processor_settings_panels = []
        conversation_type = ProcessorType.Conversation
        if self.current_config["processor"] and conversation_type in self.current_config["processor"]:
            current_conversation_config = self.current_config["processor"][conversation_type]
        else:
            current_conversation_config = self.get_default_config(processor_type=conversation_type)
        self.processor_settings_panels += [self.add_processor_panel(current_conversation_config, conversation_type)]

        # Add Action Buttons Panel
        self.add_action_panel()

        # Set the central widget of the Window. Widget will expand
        # to take up all the space in the window by default.
        self.config_window = QtWidgets.QWidget()
        self.config_window.setLayout(self.wlayout)
        self.setCentralWidget(self.config_window)
        self.position_window()

    def add_settings_panel(self, current_content_config: dict, search_type: SearchType):
        "Add Settings Panel for specified Search Type. Toggle Editable Search Types"
        # Get current files from config for given search type
        if search_type == SearchType.Image:
            current_content_files = current_content_config.get("input-directories", [])
            file_input_text = f"{search_type.name} Folders"
        elif search_type == SearchType.Github:
            return self.add_github_settings_panel(current_content_config, SearchType.Github)
        else:
            current_content_files = current_content_config.get("input-files", [])
            file_input_text = f"{search_type.name} Files"

        # Create widgets to display settings for given search type
        search_type_settings = QtWidgets.QWidget()
        search_type_layout = QtWidgets.QVBoxLayout(search_type_settings)
        enable_search_type = SearchCheckBox(f"Search {search_type.name}", search_type)
        # Add file browser to set input files for given search type
        input_files = FileBrowser(file_input_text, search_type, current_content_files or [])

        # Set enabled/disabled based on checkbox state
        enable_search_type.setChecked(current_content_files is not None and len(current_content_files) > 0)
        input_files.setEnabled(enable_search_type.isChecked())
        enable_search_type.stateChanged.connect(lambda _: input_files.setEnabled(enable_search_type.isChecked()))  # type: ignore[attr-defined]

        # Add setting widgets for given search type to panel
        search_type_layout.addWidget(enable_search_type)
        search_type_layout.addWidget(input_files)
        self.wlayout.addWidget(search_type_settings)

        return search_type_settings

    def add_github_settings_panel(self, current_content_config: dict, search_type: SearchType):
        search_type_settings = QtWidgets.QWidget()
        search_type_layout = QtWidgets.QVBoxLayout(search_type_settings)
        enable_search_type = SearchCheckBox(f"Search {search_type.name}", search_type)
        # Add labelled text input field
        input_fields = []

        fields = ["pat-token", "repo-name", "repo-owner", "repo-branch"]
        active = False
        for field in fields:
            field_value = current_content_config.get(field, None)
            input_field = LabelledTextField(field, search_type=search_type, default_value=field_value)
            search_type_layout.addWidget(input_field)
            input_fields += [input_field]
            if field_value:
                active = True

        # Set enabled/disabled based on checkbox state
        enable_search_type.setChecked(active)
        for input_field in input_fields:
            input_field.setEnabled(enable_search_type.isChecked())
        enable_search_type.stateChanged.connect(lambda _: [input_field.setEnabled(enable_search_type.isChecked()) for input_field in input_fields])  # type: ignore[attr-defined]

        # Add setting widgets for given search type to panel
        search_type_layout.addWidget(enable_search_type)
        for input_field in input_fields:
            search_type_layout.addWidget(input_field)
        self.wlayout.addWidget(search_type_settings)

        return search_type_settings

    def add_processor_panel(self, current_conversation_config: dict, processor_type: ProcessorType):
        "Add Conversation Processor Panel"
        # Get current settings from config for given processor type
        current_openai_api_key = current_conversation_config.get("openai-api-key", None)

        # Create widgets to display settings for given processor type
        processor_type_settings = QtWidgets.QWidget()
        processor_type_layout = QtWidgets.QVBoxLayout(processor_type_settings)
        enable_conversation = ProcessorCheckBox(f"Conversation", processor_type)
        # Add file browser to set input files for given processor type
        input_field = LabelledTextField(
            "OpenAI API Key", processor_type=processor_type, default_value=current_openai_api_key
        )

        # Set enabled/disabled based on checkbox state
        enable_conversation.setChecked(current_openai_api_key is not None)
        input_field.setEnabled(enable_conversation.isChecked())
        enable_conversation.stateChanged.connect(lambda _: input_field.setEnabled(enable_conversation.isChecked()))  # type: ignore[attr-defined]

        # Add setting widgets for given processor type to panel
        processor_type_layout.addWidget(enable_conversation)
        processor_type_layout.addWidget(input_field)
        self.wlayout.addWidget(processor_type_settings)

        return processor_type_settings

    def add_action_panel(self):
        "Add Action Panel"
        # Button to Save Settings
        action_bar = QtWidgets.QWidget()
        action_bar_layout = QtWidgets.QHBoxLayout(action_bar)

        self.configure_button = QtWidgets.QPushButton("Configure", clicked=self.configure_app)
        self.search_button = QtWidgets.QPushButton(
            "Search", clicked=lambda: webbrowser.open(f"http://{state.host}:{state.port}/")
        )
        self.search_button.setEnabled(not self.first_run)

        action_bar_layout.addWidget(self.configure_button)
        action_bar_layout.addWidget(self.search_button)
        self.wlayout.addWidget(action_bar)

    def get_default_config(self, search_type: SearchType = None, processor_type: ProcessorType = None):
        "Get default config"
        config = constants.default_config
        if search_type:
            return config["content-type"][search_type]  # type: ignore
        elif processor_type:
            return config["processor"][processor_type]  # type: ignore
        else:
            return config

    def add_error_message(self, message: str):
        "Add Error Message to Configure Screen"
        # Remove any existing error messages
        for message_prefix in ErrorType:
            for i in reversed(range(self.wlayout.count())):
                current_widget = self.wlayout.itemAt(i).widget()
                if isinstance(current_widget, QtWidgets.QLabel) and current_widget.text().startswith(
                    message_prefix.value
                ):
                    self.wlayout.removeWidget(current_widget)
                    current_widget.deleteLater()

        # Add new error message
        if message:
            error_message = QtWidgets.QLabel()
            error_message.setWordWrap(True)
            error_message.setText(message)
            error_message.setStyleSheet("color: red")
            self.wlayout.addWidget(error_message)

    def update_search_settings(self):
        "Update config with search settings from UI"
        for settings_panel in self.search_settings_panels:
            for child in settings_panel.children():
                if not isinstance(child, (SearchCheckBox, FileBrowser, LabelledTextField)):
                    continue
                if isinstance(child, SearchCheckBox):
                    # Search Type Disabled
                    if not child.isChecked() and child.search_type in self.new_config["content-type"]:
                        del self.new_config["content-type"][child.search_type]
                    # Search Type (re)-Enabled
                    if child.isChecked():
                        current_search_config = self.current_config["content-type"].get(child.search_type, {})
                        if current_search_config == None:
                            current_search_config = {}
                        default_search_config = self.get_default_config(search_type=child.search_type)
                        self.new_config["content-type"][child.search_type.value] = merge_dicts(
                            current_search_config, default_search_config
                        )
                elif isinstance(child, FileBrowser) and child.search_type in self.new_config["content-type"]:
                    if child.search_type.value == SearchType.Image:
                        self.new_config["content-type"][child.search_type.value]["input-directories"] = (
                            child.getPaths() if child.getPaths() != [] else None
                        )
                    else:
                        self.new_config["content-type"][child.search_type.value]["input-files"] = (
                            child.getPaths() if child.getPaths() != [] else None
                        )
                elif isinstance(child, LabelledTextField):
                    self.new_config["content-type"][child.search_type.value][
                        child.label.text()
                    ] = child.input_field.toPlainText()

    def update_processor_settings(self):
        "Update config with conversation settings from UI"
        for settings_panel in self.processor_settings_panels:
            for child in settings_panel.children():
                if not isinstance(child, (ProcessorCheckBox, LabelledTextField)):
                    continue
                if isinstance(child, ProcessorCheckBox):
                    # Processor Type Disabled
                    if not child.isChecked() and child.processor_type in self.new_config["processor"]:
                        del self.new_config["processor"][child.processor_type]
                    # Processor Type (re)-Enabled
                    if child.isChecked():
                        current_processor_config = self.current_config["processor"].get(child.processor_type, {})
                        default_processor_config = self.get_default_config(processor_type=child.processor_type)
                        self.new_config["processor"][child.processor_type.value] = merge_dicts(
                            current_processor_config, default_processor_config
                        )
                elif isinstance(child, LabelledTextField) and child.processor_type in self.new_config["processor"]:
                    if child.processor_type == ProcessorType.Conversation:
                        self.new_config["processor"][child.processor_type.value]["openai-api-key"] = (
                            child.input_field.toPlainText() if child.input_field.toPlainText() != "" else None
                        )

    def save_settings_to_file(self) -> bool:
        "Save validated settings to file"
        # Validate config before writing to file
        try:
            yaml_utils.parse_config_from_string(self.new_config)
        except Exception as e:
            print(f"Error validating config: {e}")
            self.add_error_message(f"{ErrorType.ConfigValidationError.value}: {e}")
            return False

        # Save the config to app config file
        self.add_error_message(None)
        yaml_utils.save_config_to_file(self.new_config, self.config_file)
        return True

    def load_updated_settings(self):
        "Hot swap to use the updated config from config file"
        # Load parsed, validated config from app config file
        args = cli(state.cli_args)
        self.current_config = self.new_config

        # Configure server with loaded config
        configure_server(args, required=True)

    def configure_app(self):
        "Save the new settings to khoj.yml. Reload app with updated settings"
        self.update_search_settings()
        self.update_processor_settings()
        if self.save_settings_to_file():
            # Setup thread to load updated settings in background
            self.thread = QThread()
            self.settings_loader = SettingsLoader(self.load_updated_settings)
            self.settings_loader.moveToThread(self.thread)

            # Connect slots and signals for thread
            self.thread.started.connect(self.settings_loader.run)
            self.settings_loader.finished.connect(self.thread.quit)
            self.settings_loader.finished.connect(self.settings_loader.deleteLater)
            self.settings_loader.error.connect(self.add_error_message)
            self.thread.finished.connect(self.thread.deleteLater)

            # Start thread
            self.thread.start()

            # Disable Save Button
            self.search_button.setEnabled(False)
            self.configure_button.setEnabled(False)
            self.configure_button.setText("Configuring...")

            # Reset UI
            self.thread.finished.connect(lambda: self.configure_button.setText("Configure"))
            self.thread.finished.connect(lambda: self.configure_button.setEnabled(True))
            self.thread.finished.connect(lambda: self.search_button.setEnabled(True))

    def position_window(self):
        "Position the window at center of X axis and near top on Y axis"
        window_rectangle = self.geometry()
        screen_center = self.screen().availableGeometry().center()
        window_rectangle.moveCenter(screen_center)
        self.move(window_rectangle.topLeft().x(), 25)

    def show_on_top(self):
        "Bring Window on Top"
        self.show()
        self.setWindowState(Qt.WindowState.WindowActive)
        self.activateWindow()  # For Bringing to Top on Windows
        self.raise_()  # For Bringing to Top from Minimized State on OSX


class SettingsLoader(QObject):
    "Load Settings Thread"
    finished = pyqtSignal()
    error = pyqtSignal(str)

    def __init__(self, load_settings_func):
        super(SettingsLoader, self).__init__()
        self.load_settings_func = load_settings_func

    def run(self):
        "Load Settings"
        try:
            self.load_settings_func()
        except FileNotFoundError as e:
            self.error.emit(f"{ErrorType.ConfigLoadingError.value}: {e}")
        else:
            self.error.emit(None)
        self.finished.emit()


class SearchCheckBox(QtWidgets.QCheckBox):
    def __init__(self, text, search_type: SearchType, parent=None):
        self.search_type = search_type
        super(SearchCheckBox, self).__init__(text, parent=parent)


class ProcessorCheckBox(QtWidgets.QCheckBox):
    def __init__(self, text, processor_type: ProcessorType, parent=None):
        self.processor_type = processor_type
        super(ProcessorCheckBox, self).__init__(text, parent=parent)


class ErrorType(Enum):
    "Error Types"
    ConfigLoadingError = "Config Loading Error"
    ConfigValidationError = "Config Validation Error"
