######################################################################################################################
# Copyright (C) 2017-2022 Spine project consortium
# Copyright Spine Items contributors
# This file is part of Spine Items.
# Spine Items is free software: you can redistribute it and/or modify it under the terms of the GNU Lesser General
# Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option)
# any later version. This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY;
# without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General
# Public License for more details. You should have received a copy of the GNU Lesser General Public License along with
# this program. If not, see <http://www.gnu.org/licenses/>.
######################################################################################################################

"""Contains :class:`PreviewUpdater`."""
from copy import deepcopy
from dataclasses import dataclass
from multiprocessing import Process, Queue
import queue
from time import monotonic
from PySide6.QtCore import QItemSelectionModel, QModelIndex, Qt, QTimer, Slot
from PySide6.QtWidgets import QFileDialog
from spinedb_api import DatabaseMapping, SpineDBAPIError, SpineDBVersionError
from spinedb_api.export_mapping.export_mapping import ExportMapping
from spinedb_api.export_mapping.group_functions import GroupFunction
from spinedb_api.spine_io.exporters.writer import write
from ..mvcmodels.full_url_list_model import FullUrlListModel
from ..mvcmodels.mappings_table_model import MappingsTableModel
from ..mvcmodels.preview_table_model import PreviewTableModel
from ..mvcmodels.preview_tree_model import PreviewTreeModel
from ..preview_table_writer import TableWriter


class PreviewUpdater:
    def __init__(
        self, window, ui, url_model, mappings_table_model, mappings_proxy_model, mapping_editor_table_model, project_dir
    ):
        """
        Args:
            window (SpecificationEditorWindow): specification editor's window
            ui (Ui_Form): specification editor's UI
            url_model (FullUrlListModel, optional): URL model
            mappings_table_model (MappingsTableModel): mappings table model
            mappings_proxy_model (QAbstractProxyModel): proxy model for mappings table
            mapping_editor_table_model (MappingEditorTableModel): mapping table model
            project_dir (str): path to initial directory from which to load databases
        """
        self._current_url = None
        self._window = window
        self._ui = ui
        self._project_dir = project_dir
        if url_model is None:
            url_model = FullUrlListModel(self._ui.database_url_combo_box)
        self._url_model = url_model
        self._url_model.rowsInserted.connect(self._enable_controls_after_url_insertion)
        self._url_model.modelReset.connect(self._enable_controls)
        self._url_model.destroyed.connect(self._forget_url_model)
        self._ui.database_url_combo_box.setModel(self._url_model)
        self._mappings_table_model = mappings_table_model
        self._set_expect_removals_and_inserts(False)
        self._mappings_table_model.write_order_about_to_change.connect(
            lambda: self._set_expect_removals_and_inserts(True)
        )
        self._mappings_table_model.write_order_changed.connect(lambda: self._set_expect_removals_and_inserts(False))
        self._mappings_table_model.dataChanged.connect(self._rename_mapping)
        self._mappings_table_model.dataChanged.connect(self._enable_mapping)
        self._mappings_table_model.dataChanged.connect(self._update_changed_tables)
        self._mappings_proxy_model = mappings_proxy_model
        self._mapping_editor_table_model = mapping_editor_table_model
        self._mapping_editor_table_model.dataChanged.connect(lambda *_: self._update_current_mappings_tables())
        self._expect_mapping_editor_resets(False)
        self._window.current_mapping_about_to_change.connect(lambda: self._expect_mapping_editor_resets(True))
        self._window.current_mapping_changed.connect(lambda: self._expect_mapping_editor_resets(False))
        self._preview_tree_model = PreviewTreeModel()
        self._preview_tree_model.dataChanged.connect(self._update_table)
        self._preview_tree_model.dataChanged.connect(self._expand_tree_after_table_change)
        self._preview_tree_model.rowsInserted.connect(self._expand_tree_after_table_insert)
        self._preview_table_model = PreviewTableModel()
        self._stamps = {}
        self._writer_sender = Queue()
        self._writer_receiver = Queue()
        self._writer_process = Process(target=write_task_loop, args=(self._writer_receiver, self._writer_sender))
        self._writer_timer = QTimer(window)
        self._writer_timer.setInterval(100)
        self._writer_timer.timeout.connect(self._communicate)
        self._mapping_tables = {}
        self._ui.preview_tree_view.setModel(self._preview_tree_model)
        self._ui.preview_tree_view.selectionModel().currentChanged.connect(self._change_table)
        self._ui.preview_table_view.setModel(self._preview_table_model)
        self._ui.database_url_combo_box.currentTextChanged.connect(self._reload_preview)
        self._ui.live_preview_check_box.stateChanged.connect(self._ui.frame_preview.setEnabled)
        self._ui.live_preview_check_box.clicked.connect(self._reload_preview)
        self._ui.max_preview_rows_spin_box.valueChanged.connect(self._reload_preview)
        self._ui.max_preview_tables_spin_box.valueChanged.connect(self._reload_preview)
        self._ui.load_url_from_fs_button.clicked.connect(self._load_url_from_filesystem)
        self._ui.mappings_table.selectionModel().currentChanged.connect(self._change_selected_table)
        self._enable_controls()
        if self._url_model.rowCount() > 0:
            self._reload_preview()

    @Slot()
    def _reload_preview(self):
        """Sets the current url and reloads preview."""
        current_url = self._ui.database_url_combo_box.currentText()
        if not current_url:
            return
        self._current_url = current_url
        self._stamps.clear()
        for row in range(self._mappings_table_model.rowCount()):
            index = self._mappings_table_model.index(row, 0)
            if index.data(Qt.ItemDataRole.CheckStateRole) == Qt.CheckState.Checked:
                self._load_preview_data(index.data())

    @Slot(QModelIndex, QModelIndex)
    def _change_table(self, current, previous):
        """
        Changes preview table data.

        Args:
            current (QModelIndex): index to the currently selected table on the list
            previous (QModelIndex): index to the previously selected table on the list
        """
        current_parent = self._preview_tree_model.parent(current)
        if not current_parent.isValid():
            current_parent = current
        previous_parent = self._preview_tree_model.parent(previous)
        if not previous_parent.isValid():
            previous_parent = previous
        if current_parent != previous_parent:
            name = self._preview_tree_model.data(current_parent)
            index = self._mappings_proxy_model.mapFromSource(self._mappings_table_model.index_of(name))
            self._ui.mappings_table.selectionModel().currentChanged.disconnect(self._change_selected_table)
            self._ui.mappings_table.setCurrentIndex(index)
            self._ui.mappings_table.selectionModel().currentChanged.connect(self._change_selected_table)
        table = current.data(PreviewTreeModel.TABLE_ROLE)
        if table is not None:
            table_name = current.data()
            mapping_name = current.parent().data()
            mapping_colors = self._mapping_editor_table_model.mapping_colors()
            self._preview_table_model.reset(mapping_name, table_name, table, mapping_colors)
        else:
            self._preview_table_model.clear()

    @Slot(QModelIndex, QModelIndex)
    def _change_selected_table(self, current, previous):
        """
        Changes selected table on the tree view after user has selected another mapping.

        Args:
            current (QModelIndex): index to selected mapping on mapping list
            previous (QModelIndex): index to previously selected mapping on mapping list
        """
        if not current.isValid():
            return
        current = self._mappings_proxy_model.mapToSource(current)
        mappings_index = self._mappings_table_model.index(current.row(), 0)
        if mappings_index.data(Qt.ItemDataRole.CheckStateRole) == Qt.CheckState.Unchecked.value:
            return
        current_preview = self._ui.preview_tree_view.selectionModel().currentIndex()
        if current_preview.parent().row() >= self._preview_tree_model.rowCount():
            # May happen when mappings at the end of the list are removed.
            return
        mapping_name = mappings_index.data()
        if current_preview.parent().data() == mapping_name:
            return
        self._set_current_table(mapping_name)

    def _set_current_table(self, name):
        """Sets the current index of preview tree.

        Args:
            name (str): mapping's name
        """
        for row in range(self._preview_tree_model.rowCount()):
            index = self._preview_tree_model.index(row, 0)
            if index.data() == name:
                if self._preview_tree_model.rowCount(index) > 0:
                    table_index = self._preview_tree_model.index(0, 0, index)
                    self._ui.preview_tree_view.selectionModel().setCurrentIndex(
                        table_index, QItemSelectionModel.ClearAndSelect
                    )
                    return
                self._ui.preview_tree_view.selectionModel().setCurrentIndex(index, QItemSelectionModel.ClearAndSelect)
                return

    @Slot(QModelIndex, QModelIndex, list)
    def _expand_tree_after_table_change(self, top_left, bottom_right, roles):
        """
        Expands preview tree's root items if their tables have received updates.

        Args:
            top_left (QModelIndex): top left corner of changed preview tree data
            bottom_right (QModelIndex): bottom right corner of changed preview tree data
            roles (list of int): roles
        """
        if PreviewTreeModel.TABLE_ROLE not in roles:
            return
        root_index = top_left.parent()
        if not root_index.isValid():
            return
        self._ui.preview_tree_view.setExpanded(root_index, True)

    @Slot(QModelIndex, QModelIndex, list)
    def _expand_tree_after_table_insert(self, parent, first, last):
        """
        Expands preview tree's root items if tables have been added/removed.

        Args:
            parent (QModelIndex): parent index of items that have been added/removed
            first (int): first row index of affected items
            last (int): last row index of affected items
        """
        if not parent.isValid():
            return
        self._ui.preview_tree_view.setExpanded(parent, True)

    @Slot(QModelIndex, QModelIndex, list)
    def _update_table(self, top_left, bottom_right, roles):
        """
        Updates the preview table if its data has changed.

        Args:
            top_left (QModelIndex): top left corner of changed preview tree data
            bottom_right (QModelIndex): bottom right corner of changed preview tree data
            roles (list of int): roles
        """
        if PreviewTreeModel.TABLE_ROLE not in roles:
            return
        current_parent = top_left.parent()
        current_table_name = self._preview_table_model.table_name()
        current_mapping_name = self._preview_table_model.mapping_name()
        for row in range(top_left.row(), bottom_right.row() + 1):
            index = self._preview_tree_model.index(row, 0, current_parent)
            table_name = index.data()
            mapping_name = index.parent().data()
            if current_table_name != table_name or current_mapping_name != mapping_name:
                continue
            table = index.data(PreviewTreeModel.TABLE_ROLE)
            mapping_colors = self._mapping_editor_table_model.mapping_colors()
            self._preview_table_model.reset(mapping_name, table_name, table, mapping_colors)
            break

    def _set_expect_removals_and_inserts(self, except_removals):
        """Connects and disconnects signals with regarding mapping removals.

        Args:
            except_removals (bool): True means that a mapping removal is expected
        """
        if except_removals:
            self._mappings_table_model.rowsAboutToBeRemoved.disconnect(self._remove_mappings)
            self._mappings_table_model.rowsInserted.disconnect(self._add_mappings)
        else:
            self._mappings_table_model.rowsAboutToBeRemoved.connect(self._remove_mappings)
            self._mappings_table_model.rowsInserted.connect(self._add_mappings)

    @Slot(int, int, QModelIndex)
    def _remove_mappings(self, parent, first, last):
        """
        Removes mappings from the preview models.

        Args:
            parent (QModelIndex): parent index, ignored
            first (int): first mapping list row to be removed
            last (int): last mapping list row to be removed
        """
        if self._current_url is None:
            return
        for row in range(last, first - 1, -1):
            index = self._mappings_table_model.index(row, 0)
            if index.data(Qt.ItemDataRole.CheckStateRole) == Qt.CheckState.Unchecked.value:
                continue
            removed_name = index.data()
            self._preview_tree_model.remove_mapping(removed_name)
            self._stamps.pop((self._current_url, removed_name), None)

    @Slot(QModelIndex, QModelIndex, list)
    def _rename_mapping(self, top_left, bottom_right, roles):
        """
        Renames preview models' mappings.

        Args:
            top_left (QModelIndex): top left corner of modified mappings' in mapping list model
            bottom_right (QModelIndex): bottom right corner of modified mappings' in mapping list model
            roles (list of int): changed data's role
        """
        if Qt.ItemDataRole.DisplayRole not in roles or self._current_url is None:
            return
        make_index = self._mappings_table_model.index
        indexes = [make_index(row, 0) for row in range(self._mappings_table_model.rowCount())]
        names = [
            index.data() for index in indexes if index.data(Qt.ItemDataRole.CheckStateRole) == Qt.CheckState.Checked
        ]
        old_name, new_name = self._preview_tree_model.rename_mappings(names)
        if not old_name:
            return
        self._stamps.pop((self._current_url, old_name), None)
        self._set_current_table(new_name)

    @Slot(QModelIndex, QModelIndex, list)
    def _enable_mapping(self, top_left, bottom_right, roles):
        """
        Removes/adds mappings as they are enabled or disabled.

        Args:
            top_left (QModelIndex): top left corner of modified mappings' in mapping list model
            bottom_right (QModelIndex): bottom right corner of modified mappings' in mapping list model
            roles (list of int): changed data's role
        """
        if Qt.ItemDataRole.CheckStateRole.value not in roles or self._current_url is None:
            return
        first = top_left.row()
        last = bottom_right.row()
        for row in range(first, last + 1):
            index = self._mappings_table_model.index(row, 0)
            enabled = index.data(Qt.ItemDataRole.CheckStateRole) == Qt.CheckState.Checked
            name = index.data()
            if not enabled and self._preview_tree_model.has_name(name):
                self._preview_tree_model.remove_mapping(name)
                self._stamps.pop((self._current_url, name), None)
            elif not self._preview_tree_model.has_name(name):
                self._load_preview_data(name)

    @Slot(QModelIndex, QModelIndex, list)
    def _update_changed_tables(self, top_left, bottom_right, roles):
        """
        Updates mappings whenever the always write header flag or fixed table name changes.

        Args:
            top_left (QModelIndex): top left corner of modified mappings' in mapping list model
            bottom_right (QModelIndex): bottom right corner of modified mappings' in mapping list model
            roles (list of int): changed data's role
        """
        if not {
            MappingsTableModel.ALWAYS_EXPORT_HEADER_ROLE,
            MappingsTableModel.FIXED_TABLE_NAME_ROLE,
            MappingsTableModel.GROUP_FN_ROLE,
        } & set(roles):
            return
        row = self._mappings_proxy_model.mapToSource(self._ui.mappings_table.currentIndex()).row()
        index = self._mappings_table_model.index(row, 0)
        if index.data(Qt.ItemDataRole.CheckStateRole) == Qt.CheckState.Unchecked.value:
            return
        name = index.data()
        self._load_preview_data(name)

    @Slot()
    def _update_current_mappings_tables(self):
        """Reloads mapping that is currently selected on mapping list."""
        if self._current_url is None:
            return
        current_index = self._ui.mappings_table.currentIndex()
        if not current_index.isValid():
            return
        row = self._mappings_proxy_model.mapToSource(current_index).row()
        index = self._mappings_table_model.index(row, 0)
        if index.data(Qt.ItemDataRole.CheckStateRole) == Qt.CheckState.Unchecked.value:
            return
        mapping_name = index.data()
        self._load_preview_data(mapping_name)

    @Slot(QModelIndex, int, int)
    def _add_mappings(self, parent, first, last):
        """Adds new mapping for preview.

        Args:
            parent (QModelIndex): parent index - unused
            first (int): first row added to mappings table
            last (int): last row added to mappings table
        """
        if self._current_url is None:
            return
        for _ in range(first, last + 1):
            index = self._mappings_table_model.index(first, 0)
            if index.data(Qt.ItemDataRole.CheckStateRole) == Qt.CheckState.Unchecked.value:
                continue
            mapping_name = index.data()
            self._load_preview_data(mapping_name)

    def _expect_mapping_editor_resets(self, expect_resets):
        """Connects and disconnects signals that handle mapping editor resets.

        Args:
            expect_resets (bool): True to except editor resets
        """
        if expect_resets:
            self._mapping_editor_table_model.modelReset.disconnect(self._update_current_mappings_tables)
        else:
            self._mapping_editor_table_model.modelReset.connect(self._update_current_mappings_tables)

    @Slot(QModelIndex, int, int)
    def _enable_controls_after_url_insertion(self, parent, first, last):
        """Enables controls after new preview database URL has been inserted."""
        self._enable_controls()

    @Slot()
    def _enable_controls(self):
        """Enables and disables widgets as needed."""
        urls_available = self._url_model.rowCount() > 0
        self._ui.max_preview_rows_spin_box.setEnabled(urls_available)
        self._ui.max_preview_tables_spin_box.setEnabled(urls_available)

    def _load_preview_data(self, mapping_name):
        """
        Loads preview data from database into the preview tables if the url is set and live previews are enabled.

        Args:
            mapping_name (str): mapping's name
        """
        if self._current_url is None or not self._ui.live_preview_check_box.isChecked() or not mapping_name:
            return
        mapping_spec = self._mappings_table_model.mapping_specification(mapping_name)
        max_tables = self._ui.max_preview_tables_spin_box.value()
        max_rows = self._ui.max_preview_rows_spin_box.value()
        id_ = (self._current_url, mapping_name)
        stamp = monotonic()
        self._stamps[id_] = stamp
        if not self._writer_process.is_alive():
            self._writer_process.start()
        self._writer_sender.put(
            WriteTableTask(
                self._current_url,
                mapping_name,
                stamp,
                deepcopy(mapping_spec.root),
                mapping_spec.always_export_header,
                max_tables,
                max_rows,
                mapping_spec.group_fn,
            )
        )
        if not self._writer_timer.isActive():
            self._writer_timer.start()
        else:
            # This ensures we don't end up in a deadlock where we keep sending
            # new tasks while the pipe is stuck with finished ones.
            self._communicate()

    @Slot()
    def _communicate(self):
        """Periodically communicates with the writer process."""
        self._writer_timer.stop()
        if self._writer_process is None:
            return
        try:
            message = self._writer_receiver.get_nowait()
        except queue.Empty:
            pass
        else:
            self._add_or_update_data(*message)
        if self._stamps:
            self._writer_timer.setInterval(50)
            self._writer_timer.start()
        else:
            self._writer_timer.setInterval(100)

    @Slot(tuple, str, object, float)
    def _add_or_update_data(self, worker_id, mapping_name, data, stamp):
        """
        Sets preview data.

        Args:
            worker_id (tuple): a worker identifier
            mapping_name (str): mapping's name
            data (dict): mapping from table name to table
            stamp (float): task's time stamp
        """
        current_stamp = self._stamps.pop(worker_id, None)
        if stamp != current_stamp:
            if current_stamp is not None:
                self._stamps[worker_id] = current_stamp
            return
        self._preview_tree_model.add_or_update_tables(mapping_name, data)
        mapping_index = self._ui.mappings_table.selectionModel().currentIndex()
        self._change_selected_table(mapping_index, mapping_index)

    @Slot(bool)
    def _load_url_from_filesystem(self, _):
        """Shows a file dialog and opens selected database in the preview widget."""
        path = self._browse_database()
        if not path:
            return
        url = "sqlite:///" + path
        self._url_model.append(url)
        self._ui.database_url_combo_box.setCurrentText(url)

    def _browse_database(self):
        """
        Queries a database file from the user.

        Returns:
            str: path to database file
        """
        return QFileDialog.getOpenFileName(self._window, "Select database", self._project_dir, "sqlite (*.sqlite)")[0]

    @Slot()
    def _forget_url_model(self):
        """Replaces current URL model with a mock one."""
        self._url_model = FullUrlListModel(self._ui.database_url_combo_box)
        self._url_model.rowsInserted.connect(self._enable_controls_after_url_insertion)
        self._url_model.modelReset.connect(self._enable_controls)
        self._url_model.destroyed.connect(self._forget_url_model)
        self._ui.database_url_combo_box.setModel(self._url_model)
        self._reload_preview()

    def tear_down(self):
        """Stops the writer process and cleans up."""
        self._writer_sender.put("quit")
        self._writer_timer.stop()
        self._stamps.clear()
        self._url_model.rowsInserted.disconnect(self._enable_controls_after_url_insertion)
        self._url_model.modelReset.disconnect(self._enable_controls)
        self._url_model.destroyed.disconnect(self._forget_url_model)
        if self._writer_process.is_alive():
            finished = False
            while not finished:
                # Drain the pipe, otherwise writer process's send() may block and the join() below hangs.
                finished = self._writer_receiver.get() == "finished"
            self._writer_process.join()
        self._writer_process = None


@dataclass(frozen=True)
class WriteTableTask:
    url: str
    mapping_name: str
    stamp: float
    mapping: ExportMapping
    always_export_header: bool
    max_tables: int
    max_rows: int
    group_fn: GroupFunction


def write_task_loop(sender, receiver):
    """A task loop that processes table writing tasks.

    Args:
        sender (Queue): Sending queue to communicate with the loop
        receiver (Queue): Receiving queue endpoint to communicate with the loop
    """
    db_map = None
    tasks = []
    try:
        while True:
            if not receiver.empty() or not tasks:
                while True:
                    task = receiver.get()
                    if task == "quit":
                        return
                    tasks.append(task)
                    if receiver.empty():
                        break
            next_task = tasks.pop(0)
            try:
                if db_map is None or next_task.url != db_map.db_url:
                    db_map = DatabaseMapping(next_task.url)
                tables = _write_tables(db_map, next_task)
            except SpineDBVersionError:
                tables = {"error": [["unsupported database version"]]}
            except SpineDBAPIError as error:
                tables = {"error": [[str(error)]]}
            sender.put(((next_task.url, next_task.mapping_name), next_task.mapping_name, tables, next_task.stamp))
    finally:
        sender.put_nowait("finished")


def _write_tables(db_map, write_task):
    """Creates preview tables with given parameters.

    Args:
        db_map (DatabaseMapping): database mapping
        write_task (WriteTableTask): task parameters

    Returns:
        tuple: mappings from table names to tables (list of rows)
    """
    writer = TableWriter()
    write(
        db_map,
        writer,
        write_task.mapping,
        empty_data_header=write_task.always_export_header,
        max_tables=write_task.max_tables,
        max_rows=write_task.max_rows,
        group_fns=write_task.group_fn,
    )
    return writer.tables
