#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# ================================================== #
# This file is a part of PYGPT package               #
# Website: https://pygpt.net                         #
# GitHub:  https://github.com/szczyglis-dev/py-gpt   #
# MIT License                                        #
# Created By  : Marcin Szczygliński                  #
# Updated Date: 2024.03.18 03:00:00                  #
# ================================================== #

import datetime
import os
from pathlib import Path

from PySide6.QtCore import Qt, QModelIndex, QDir
from PySide6.QtGui import QAction, QIcon, QCursor, QResizeEvent
from PySide6.QtWidgets import QTreeView, QMenu, QWidget, QVBoxLayout, QFileSystemModel, QLabel, QHBoxLayout, \
    QPushButton, QSizePolicy

from pygpt_net.ui.widget.element.button import ContextMenuButton
from pygpt_net.ui.widget.element.labels import HelpLabel
from pygpt_net.utils import trans
import pygpt_net.icons_rc


class FileExplorer(QWidget):
    def __init__(self, window, directory, index_data):
        """
        File explorer widget

        :param window: Window instance
        :param directory: directory to explore
        :param index_data: index data
        """
        super().__init__()

        self.window = window
        self.index_data = index_data
        self.directory = directory
        self.model = IndexedFileSystemModel(self.window, self.index_data)
        self.model.setRootPath(directory)
        self.model.setFilter(self.model.filter() | QDir.Hidden)  # show hidden files

        self.treeView = QTreeView()
        self.treeView.setModel(self.model)
        self.treeView.setRootIndex(self.model.index(directory))

        header = QHBoxLayout()

        self.btn_open = QPushButton(trans('action.open'))
        self.btn_open.setMaximumHeight(40)
        self.btn_open.clicked.connect(
                lambda: self.action_open(self.directory)
        )
        self.btn_open.setSizePolicy(QSizePolicy.Minimum, QSizePolicy.Minimum)

        self.btn_upload = QPushButton(trans('files.local.upload'))
        self.btn_upload.setMaximumHeight(40)
        self.btn_upload.clicked.connect(self.window.controller.files.upload_local)
        self.btn_upload.setSizePolicy(QSizePolicy.Minimum, QSizePolicy.Minimum)

        self.btn_idx = ContextMenuButton(trans('idx.btn.index_all'), self)
        self.btn_idx.action = self.idx_context_menu
        self.btn_idx.setMaximumHeight(40)
        self.btn_idx.setSizePolicy(QSizePolicy.Minimum, QSizePolicy.Minimum)

        self.btn_clear = ContextMenuButton(trans('idx.btn.clear'), self)
        self.btn_clear.action = self.clear_context_menu
        self.btn_clear.setMaximumHeight(40)
        self.btn_clear.setSizePolicy(QSizePolicy.Minimum, QSizePolicy.Minimum)

        self.path_label = QLabel(self.directory)
        self.path_label.setMaximumHeight(40)
        self.path_label.setAlignment(Qt.AlignCenter | Qt.AlignVCenter)

        header.addWidget(self.btn_open)
        header.addWidget(self.btn_upload)
        header.addStretch()
        header.addWidget(self.path_label)
        header.addStretch()

        header.addWidget(self.btn_idx)
        header.addWidget(self.btn_clear)
        layout = QVBoxLayout()

        self.window.ui.nodes['tip.output.tab.files'] = HelpLabel(trans('tip.output.tab.files'), self.window)

        layout.addWidget(self.treeView)
        layout.addWidget(self.window.ui.nodes['tip.output.tab.files'])
        layout.addLayout(header)
        layout.setContentsMargins(0, 0, 0, 0)
        self.setLayout(layout)

        self.treeView.setContextMenuPolicy(Qt.CustomContextMenu)
        self.treeView.customContextMenuRequested.connect(self.openContextMenu)
        self.treeView.setColumnWidth(0, self.width() / 2)

        self.header = self.treeView.header()
        self.header.setStretchLastSection(True)
        self.header.setContentsMargins(0, 0, 0, 0)

        self.column_proportion = 0.3
        self.adjustColumnWidths()

        self.header.setStyleSheet("""
           QHeaderView::section {
               text-align: center;
               vertical-align: middle;
           }
       """)

    def idx_context_menu(self, parent, pos):
        """
        Index all btn context menu

        :param parent: parent widget
        :param pos: mouse  position
        """
        menu = QMenu(self)
        idx_list = self.window.core.config.get('llama.idx.list')
        if len(idx_list) > 0:
            for idx in idx_list:
                id = idx['id']
                name = idx['name'] + " (" + idx['id'] + ")"
                action = menu.addAction("IDX: " + name)
                action.triggered.connect(
                    lambda checked=False,
                           id=id: self.window.controller.idx.indexer.index_all_files(id)
                )
        menu.exec_(parent.mapToGlobal(pos))

    def clear_context_menu(self, parent, pos):
        """
        Clear btn context menu

        :param parent: parent widget
        :param pos: mouse position
        """
        menu = QMenu(self)
        idx_list = self.window.core.config.get('llama.idx.list')
        if len(idx_list) > 0:
            for idx in idx_list:
                id = idx['id']
                name = idx['name'] + " (" + idx['id'] + ")"
                action = menu.addAction("IDX: " + name)
                action.triggered.connect(
                    lambda checked=False,
                           id=id: self.window.controller.idx.indexer.clear(id)
                )
        menu.exec_(parent.mapToGlobal(pos))

    def adjustColumnWidths(self):
        """Adjust column widths"""
        total_width = self.treeView.width()
        first_column_width = int(total_width * self.column_proportion)
        self.treeView.setColumnWidth(0, first_column_width)
        for column in range(1, self.model.columnCount()):
            self.treeView.setColumnWidth(column, (total_width - first_column_width) // (self.model.columnCount() - 1))

    def resizeEvent(self, event: QResizeEvent):
        """
        Resize event

        :param event: Event object
        """
        super().resizeEvent(event)
        self.adjustColumnWidths()

    def openContextMenu(self, position):
        """
        Open context menu

        :param position: mouse position
        """
        indexes = self.treeView.selectedIndexes()
        if indexes:
            index = indexes[0]
            path = self.model.filePath(index)
            actions = {}

            """
            # play
            if self.window.core.filesystem.is_video(path):
                actions['play_video'] = QAction(
                    QIcon(":/icons/video.svg"),
                    trans('action.video.play'),
                    self,
                )
                actions['play_video'].triggered.connect(
                    lambda: self.window.controller.video.play(path),
                )
            """

            # open file
            actions['open'] = QAction(QIcon(":/icons/view.svg"), trans('action.open'), self)
            actions['open'].triggered.connect(
                lambda: self.action_open(path),
            )

            # open in file manager
            actions['open_dir'] = QAction(QIcon(":/icons/folder_filled.svg"), trans('action.open_dir'), self)
            actions['open_dir'].triggered.connect(
                lambda: self.action_open_dir(path),
            )

            # download
            actions['download'] = QAction(QIcon(":/icons/download.svg"), trans('action.download'), self)
            actions['download'].triggered.connect(
                lambda: self.window.controller.files.download_local(path),
            )

            # rename
            actions['rename'] = QAction(QIcon(":/icons/edit.svg"), trans('action.rename'), self)
            actions['rename'].triggered.connect(
                lambda: self.action_rename(path),
            )

            # duplicate
            actions['duplicate'] = QAction(QIcon(":/icons/stack.svg"), trans('action.duplicate'), self)
            actions['duplicate'].triggered.connect(
                lambda: self.window.controller.files.duplicate_local(path, ""),
            )

            if os.path.isdir(path):
                parent = path
            else:
                parent = os.path.dirname(path)

            # touch file
            actions['touch'] = QAction(QIcon(":/icons/add.svg"), trans('action.touch'), self)
            actions['touch'].triggered.connect(
                lambda: self.window.controller.files.touch_file(parent),
            )

            # make dir in dir
            actions['mkdir'] = QAction(QIcon(":/icons/add_folder.svg"), trans('action.mkdir'), self)
            actions['mkdir'].triggered.connect(
                lambda: self.action_make_dir(parent),
            )

            # refresh
            actions['refresh'] = QAction(QIcon(":/icons/reload.svg"), trans('action.refresh'), self)
            actions['refresh'].triggered.connect(
                lambda: self.window.controller.files.update_explorer(),
            )

            # upload to dir
            actions['upload'] = QAction(QIcon(":/icons/upload.svg"), trans('action.upload'), self)
            actions['upload'].triggered.connect(
                lambda: self.window.controller.files.upload_local(parent),
            )

            # delete
            actions['delete'] = QAction(QIcon(":/icons/delete.svg"), trans('action.delete'), self)
            actions['delete'].triggered.connect(
                lambda: self.action_delete(path),
            )

            menu = QMenu(self)
            """
            if "play_video" in actions:
                menu.addAction(actions['play_video'])
            """
            menu.addAction(actions['open'])
            menu.addAction(actions['open_dir'])

            # use
            use_menu = QMenu(trans('action.use'), self)

            # use as attachment
            if not os.path.isdir(path):
                actions['use_attachment'] = QAction(
                    QIcon(":/icons/attachment.svg"),
                    trans('action.use.attachment'),
                    self,
                )
                actions['use_attachment'].triggered.connect(
                    lambda: self.window.controller.files.use_attachment(path),
                )

                # by type
                if self.window.core.filesystem.is_image(path):
                    actions['use_as_image'] = QAction(
                        QIcon(":/icons/brush.svg"),
                        trans('action.use.image'),
                        self,
                    )
                    actions['use_as_image'].triggered.connect(
                        lambda: self.window.controller.painter.open_external(path),
                    )

            # copy work path
            actions['use_copy_work_path'] = QAction(
                QIcon(":/icons/copy.svg"),
                trans('action.use.copy_work_path'),
                self,
            )
            actions['use_copy_work_path'].triggered.connect(
                lambda: self.window.controller.files.copy_work_path(path),
            )

            # copy sys path
            actions['use_copy_sys_path'] = QAction(
                QIcon(":/icons/copy.svg"),
                trans('action.use.copy_sys_path'),
                self,
            )
            actions['use_copy_sys_path'].triggered.connect(
                lambda: self.window.controller.files.copy_sys_path(path),
            )

            # use read cmd
            actions['use_read_cmd'] = QAction(QIcon(":/icons/view.svg"), trans('action.use.read_cmd'), self)
            actions['use_read_cmd'].triggered.connect(
                lambda: self.window.controller.files.make_read_cmd(path),
            )

            # add actions to menu
            if not os.path.isdir(path):
                use_menu.addAction(actions['use_attachment'])
            if "use_as_image" in actions:
                use_menu.addAction(actions['use_as_image'])
            use_menu.addAction(actions['use_copy_work_path'])
            use_menu.addAction(actions['use_copy_sys_path'])
            use_menu.addAction(actions['use_read_cmd'])
            menu.addMenu(use_menu)

            # remove from index
            file_id = self.window.core.idx.files.get_id(path)
            remove_actions = []
            for idx in self.index_data:
                items = self.index_data[idx]
                if file_id in items:
                    action = QAction(QIcon(":/icons/delete.svg"), trans("action.idx.remove") + ": " + idx, self)
                    action.triggered.connect(
                        lambda checked=False,
                               idx=idx,
                               file_id=file_id: self.action_idx_remove(file_id, idx)  # by file_id, not path
                    )
                    remove_actions.append(action)

            # add to index if allowed
            if self.window.core.idx.indexing.is_allowed(path):
                idx_menu = QMenu(trans('action.idx'), self)
                idx_list = self.window.core.config.get('llama.idx.list')
                if len(idx_list) > 0 or len(remove_actions) > 0:
                    # add
                    if len(idx_list) > 0:
                        for idx in idx_list:
                            id = idx['id']
                            name = idx['name'] + " (" + idx['id'] + ")"
                            action = QAction(QIcon(":/icons/search.svg"), "IDX: " + name, self)
                            action.triggered.connect(
                                lambda checked=False,
                                       id=id,
                                       path=path: self.action_idx(path, id)
                            )
                            idx_menu.addAction(action)

                # remove
                if len(remove_actions) > 0:
                    idx_menu.addSeparator()
                    for action in remove_actions:
                        idx_menu.addAction(action)

                menu.addMenu(idx_menu)

            menu.addAction(actions['download'])
            menu.addAction(actions['touch'])
            menu.addAction(actions['mkdir'])
            menu.addAction(actions['upload'])
            menu.addAction(actions['refresh'])

            menu.addAction(actions['rename'])
            menu.addAction(actions['duplicate'])
            menu.addAction(actions['delete'])

            menu.exec(QCursor.pos())
        else:
            # no item selected
            actions = {}

            # touch file
            actions['touch'] = QAction(QIcon(":/icons/add.svg"), trans('action.touch'), self)
            actions['touch'].triggered.connect(
                lambda: self.window.controller.files.touch_file(self.directory),
            )

            # open in file manager
            actions['open_dir'] = QAction(QIcon(":/icons/folder_filled.svg"), trans('action.open_dir'), self)
            actions['open_dir'].triggered.connect(
                lambda: self.action_open_dir(self.directory),
            )

            # make dir in dir
            actions['mkdir'] = QAction(QIcon(":/icons/add_folder.svg"), trans('action.mkdir'), self)
            actions['mkdir'].triggered.connect(
                lambda: self.action_make_dir(self.directory),
            )

            # upload in dir
            actions['upload'] = QAction(QIcon(":/icons/upload.svg"), trans('action.upload'), self)
            actions['upload'].triggered.connect(
                lambda: self.window.controller.files.upload_local(),
            )

            menu = QMenu(self)
            menu.addAction(actions['touch'])
            menu.addAction(actions['open_dir'])
            menu.addAction(actions['mkdir'])
            menu.addAction(actions['upload'])
            menu.exec(QCursor.pos())

    def action_open(self, path):
        """
        Open action handler

        :param path: path to open
        """
        self.window.controller.files.open(path)

    def action_idx(self, path: str, idx: str):
        """
        Index file or dir handler

        :param path: path to open
        :param idx: index ID to use (name)
        """
        self.window.controller.idx.indexer.index_file(path, idx)

    def action_idx_remove(self, path: str, idx: str):
        """
        Remove file or dir from index handler

        :param path: path to open
        :param idx: index ID to use (name)
        """
        self.window.controller.idx.indexer.index_file_remove(path, idx)

    def action_open_dir(self, path: str):
        """
        Open in directory action handler

        :param path: path to open
        """
        self.window.controller.files.open_dir(path, True)

    def action_make_dir(self, path: str = None):
        """
        Make directory action handler

        :param path: parent path
        """
        self.window.controller.files.make_dir_dialog(path)

    def action_rename(self, path: str):
        """
        Rename action handler

        :param path: path to rename
        """
        self.window.controller.files.rename(path)

    def action_delete(self, path: str):
        """
        Delete action handler

        :param path: path to delete
        """
        self.window.controller.files.delete(path)


class IndexedFileSystemModel(QFileSystemModel):
    def __init__(self, window, index_dict, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.window = window
        self.index_dict = index_dict
        self.directoryLoaded.connect(self.refresh_path)

    def refresh_path(self, path):
        index = self.index(path)
        if index.isValid():
            self.dataChanged.emit(index, index)

    def columnCount(self, parent=QModelIndex()) -> int:
        """
        Return column count

        :param parent: parent
        :return: column count
        """
        return super().columnCount(parent) + 1

    def data(self, index, role=Qt.DisplayRole) -> any:
        """
        Data handler

        :param index: row index
        :param role: role
        :return: data
        """
        # index status
        if index.column() == self.columnCount() - 1:
            if role == Qt.DisplayRole:
                file_path = self.filePath(index.siblingAtColumn(0))
                status = self.get_index_status(file_path)   # get index status
                if status['indexed']:
                    ts = datetime.datetime.fromtimestamp(status['last_index_at'])
                    # check if today
                    if ts.date() == datetime.date.today():
                        dt = ts.strftime("%H:%M")
                    else:
                        dt = ts.strftime("%Y-%m-%d %H:%M")
                    content = ''
                    content += dt
                    content += ' (' + ",".join(status['indexed_in']) + ')'
                else:
                    content = '-'  # if file not indexed
                return content
        # modified date
        elif index.column() == self.columnCount() - 2:
            if role == Qt.DisplayRole:
                dt = self.lastModified(index)
                data = dt.toString("yyyy-MM-dd HH:mm:ss")
                ts = dt.toSecsSinceEpoch()
                dt_from_ts = datetime.datetime.fromtimestamp(ts)
                if dt_from_ts.date() == datetime.date.today():
                    data = dt_from_ts.strftime("%H:%M")
                else:
                    data = dt_from_ts.strftime("%Y-%m-%d %H:%M")
                file_path = self.filePath(index.siblingAtColumn(0))
                status = self.get_index_status(file_path)  # get index status
                if status['indexed']:
                    # if modified date is newer, mark file with *
                    if 'last_index_at' in status and status['last_index_at'] < dt.toSecsSinceEpoch():
                        data += '*'
                return data

        return super().data(index, role)

    def get_index_status(self, file_path) -> dict:
        """
        Get index status

        :param file_path: file path
        :return: file index status
        """
        file_id = self.window.core.idx.files.get_id(file_path)
        indexed_in = []
        indexed_timestamps = {}
        last_index_at = 0
        for idx in self.index_dict:
            items = self.index_dict[idx]
            if file_id in items:
                indexed_in.append(idx)  # append idx where file is indexed
                indexed_timestamps[idx] = items[file_id]['indexed_ts']
                if items[file_id]['indexed_ts'] > last_index_at:
                    last_index_at = items[file_id]['indexed_ts']

        # sort indexed_in by timestamp DESC
        indexed_in = sorted(
            indexed_in,
            key=lambda x: indexed_timestamps[x],
            reverse=True,
        )
        if len(indexed_in) > 0:
            return {
                'indexed': True,
                'indexed_in': indexed_in,
                'last_index_at': last_index_at,
            }
        else:
            return {
                'indexed': False,
            }

    def headerData(self, section, orientation, role=Qt.DisplayRole) -> str:
        """
        Prepare Header data (append Indexed column)

        :param section: Section
        :param orientation: Orientation
        :param role: Role
        :return: Header data
        """
        if orientation == Qt.Horizontal and role == Qt.DisplayRole:
            if section == 0:  # name
                return trans('files.explorer.header.name')
            elif section == 1:  # size
                return trans('files.explorer.header.size')
            elif section == 2:  # type
                return trans('files.explorer.header.type')
            elif section == 3:  # modified
                return trans('files.explorer.header.modified')
            elif section == 4:  # indexed
                return trans('files.explorer.header.indexed')
        return super().headerData(section, orientation, role)

    def update_idx_status(self, idx_data):
        """
        Update index data status

        :param idx_data: new index data dict
        """
        self.index_dict = idx_data
        top_left_index = self.index(0, 0)
        bottom_right_index = self.index(self.rowCount() - 1, self.columnCount() - 1)
        path = self.rootPath()
        self.dataChanged.emit(top_left_index, bottom_right_index, [Qt.DisplayRole])
        self.setRootPath("")
        self.setRootPath(path)
        self.layoutChanged.emit()
