# Copyright (c) 2026, UChicago Argonne, LLC
# BSD OPEN SOURCE LICENSE. Full license can be found in LICENSE
# Copyright (c) 2026, UChicago Argonne, LLC
# BSD OPEN SOURCE LICENSE. Full license can be found in LICENSE
# Adapted from the AequilibraE plugin
import os
import platform
import qgis
import subprocess
from functools import partial
from pathlib import Path
from polaris import Polaris
from polaris.network.network import Network
from polaris.utils.database.db_utils import commit_and_close
from qgis.PyQt.QtCore import Qt, QSize
from qgis.PyQt.QtGui import QIcon
from qgis.PyQt.QtWidgets import QCheckBox, QTabWidget, QApplication, QComboBox, QLabel
from qgis.PyQt.QtWidgets import QWidget, QDockWidget, QAction, QVBoxLayout, QToolBar, QToolButton, QMenu
from qgis.core import QgsApplication
from qgis.core import QgsDataSourceUri, QgsVectorLayer, QgsProject, QgsVectorFileWriter, QgsExpressionContextUtils
from sqlite3 import Connection
from uuid import uuid4

from QPolaris import set_polaris_menu_instance
from QPolaris.modules.checker_interface import CheckerDialog
from QPolaris.modules.cluster.check_cluster_access import is_cluster_accessible
from QPolaris.modules.common_tools import FeaturePicker
from QPolaris.modules.menu_actions import break_long_lines, rebuild_location_links, rebuild_intersections_job, plt_kpis
from QPolaris.modules.menu_actions import clear_intersection_display, help, close_project, last_folder, load_log
from QPolaris.modules.menu_actions import load_project_from_path, change_iteration, link_revert, sys_perf, supply_setter
from QPolaris.modules.menu_actions import new_demand, create_link, active_rebuilder, location_navigator, gtfs_importer
from QPolaris.modules.menu_actions import rebuild_location_parking, traffic_diagnostics, global_geo_consistency
from QPolaris.modules.menu_actions import run_consistency, new_supply, load_supply_file, break_link, json_editor
from QPolaris.modules.menu_actions import tnc_navigator, open_project, display_traffic, fixes_zones_no_locs
from QPolaris.modules.menu_actions import traffic_shortest_path, path_demand_db, multimodal_path, popsynth, run_model
from QPolaris.modules.menu_actions import transit_navigator, upgrade_databases, load_skim_matrices, cluster_status
from QPolaris.modules.menu_actions.show_dialog import show_diag
from QPolaris.modules.processing_provider.provider import Provider
from QPolaris.modules.style_loader import intersec_ctrl, load_link_lane_style
from QPolaris.modules.style_loader import load_editor_styles
from QPolaris.modules.visualization import LocationDialog

if hasattr(Qt, "AA_EnableHighDpiScaling"):
    QApplication.setAttribute(Qt.AA_EnableHighDpiScaling, True)

if hasattr(Qt, "AA_UseHighDpiPixmaps"):
    QApplication.setAttribute(Qt.AA_UseHighDpiPixmaps, True)


class PolarisMenu:
    def __init__(self, iface):
        set_polaris_menu_instance(self)
        self._network = Network()
        self.translator = None
        self.iface = iface
        self.path = last_folder()
        self.polaris_project = Polaris()
        self.alternative_project = Polaris()
        self.open_mode = None
        self.provider = None

        self.conn = None  # type: Connection
        self.layers = {}
        self.dock = QDockWidget("Polaris")
        self.manager = QWidget()
        # The self.toolbar will hold everything
        self.toolbar = QToolBar()
        self.set_font(self.toolbar)

        self.toolbar.setOrientation(2)

        self.menuActions = {
            "Load and Unload": {"root": []},
            "Result analysis": [],
            "Traffic": {"root": [], "Editing": [], "Visualization": [], "Styling": []},
            "TNC": [],
            "Transit and active": {"root": [], "Editing": [], "Visualization": []},
            "Model building": [],
            "Data": {"root": []},
            "Polaris": [],
        }

        inside_anl = is_cluster_accessible()
        if inside_anl:
            self.menuActions["Systems Management"] = {"root": []}

        # # #######################   PROJECT SUB-MENU   ############################
        self.add_menu_action("Load and Unload", "Open Project", partial(open_project, self), sub="root")
        self.add_menu_action("Load and Unload", "Load Supply", partial(load_supply_file, self), sub="root")
        self.add_menu_action("Load and Unload", "Reveal on explorer", self.open_explorer, sub="root")
        self.add_menu_action("Load and Unload", "Close all", partial(close_project, self), sub="root")

        # ###################       Traffic Sub-menu       ########################
        mmenu = "Result analysis"
        self.add_menu_action(mmenu, "Map link results", partial(display_traffic, self))
        self.add_menu_action(mmenu, "Show simulated path", partial(path_demand_db, self))
        self.add_menu_action(mmenu, "Plot KPIs", partial(plt_kpis, self))

        # ###################       Traffic Sub-menu       ########################
        mmenu = "Traffic"
        self.add_menu_action(mmenu, "Explore intersection", self.show_specific_intersection, sub="Editing")
        self.add_menu_action(mmenu, "Diagnostics", partial(traffic_diagnostics, self), sub="Editing")
        self.add_menu_action(mmenu, "Rebuild Intersections", partial(rebuild_intersections_job, self), sub="Editing")
        self.add_menu_action(mmenu, "Computing Paths", partial(traffic_shortest_path, self), sub="Visualization")

        self.add_menu_action(mmenu, "Link directional capacity", partial(load_link_lane_style, self), sub="Styling")
        self.add_menu_action(mmenu, "Intersection control", partial(intersec_ctrl, self), sub="Styling")

        # # ############################   TNC MENU   ##################################
        self.add_menu_action("TNC", "Visualize", partial(tnc_navigator, self))

        # ###################       Transit Sub-menu       ########################
        mmenu = "Transit and active"
        self.add_menu_action(mmenu, "Explore Transit", partial(transit_navigator, self), sub="Visualization")
        self.add_menu_action(mmenu, "Multimodal Paths", partial(multimodal_path, self), sub="Visualization")
        self.add_menu_action(mmenu, "Import GTFS", partial(gtfs_importer, self), sub="Editing")
        self.add_menu_action(mmenu, "Rebuild active networks", partial(active_rebuilder, self), sub="Editing")

        # ###################      Global Editor Sub-menu     ########################
        mmenu = "Model building"
        # self.add_menu_action(mmenu, "Population Synthesis", partial(popsynth, self), sub="root")

        self.add_menu_action(mmenu, "Global Geo-consistency", partial(global_geo_consistency, self))
        self.add_menu_action(mmenu, "Break long links", partial(break_long_lines, self))
        self.add_menu_action(mmenu, "Rebuild Location Links", partial(rebuild_location_links, self))
        self.add_menu_action(mmenu, "Rebuild Location Parking", partial(rebuild_location_parking, self))
        self.add_menu_action(mmenu, "Zones with no Locations", partial(fixes_zones_no_locs, self))

        # # ######################   DATA MENU   ############################
        self.add_menu_action("Data", "Display skims", partial(load_skim_matrices, self), sub="root")

        self.add_menu_action("Data", "Edit Parameters", partial(json_editor, self), sub="root")

        # # #######################     LOOSE STUFF      ###########################
        self.add_menu_action("Polaris", "Help", help)
        # if not extra_packages:
        #     self.add_menu_action('Polaris', 'Install extra packages', self.install_extra_packages)

        self.add_menu_action("Polaris", "Run model", partial(run_model, self))

        # ###################       Cluster Sub-menu       ########################
        if inside_anl:
            self.add_menu_action("Systems Management", "Cluster status", cluster_status, sub="root")
            self.add_menu_action("Systems Management", "Performance history", sys_perf, sub="root")

        self.build_menu()
        # ########################################################################
        # #################        PROJECT MANAGER       #########################

        lbl = QLabel("Model iteration")
        lbl.setAlignment(Qt.AlignCenter)
        self.toolbar.addWidget(lbl)
        self.cob_iter = QComboBox()
        self.toolbar.addWidget(self.cob_iter)

        self.projectManager = QTabWidget()
        self.toolbar.addWidget(self.projectManager)

        # # # ########################################################################
        self.tabContents = []
        self.toolbar.setIconSize(QSize(16, 16))

        p1_vertical = QVBoxLayout()
        p1_vertical.setContentsMargins(0, 0, 0, 0)
        p1_vertical.addWidget(self.toolbar)
        self.manager.setLayout(p1_vertical)
        self.manager.setMinimumHeight(430)
        self.dock.setWidget(self.manager)
        self.dock.setAllowedAreas(Qt.LeftDockWidgetArea | Qt.RightDockWidgetArea)
        self.iface.addDockWidget(Qt.LeftDockWidgetArea, self.dock)
        self.manager.setMinimumHeight(370)
        self.manager.resize(150, 430)
        QgsProject.instance().layerRemoved.connect(self.layerRemoved)
        self.actions = []
        self.menu = "&Polaris"

        self.menubar = self.iface.addToolBar("PolarisClick")
        self.add_action("polaris.png", "Polaris Help", help)
        self.add_action("open_file.jpg", "Open Supply", partial(load_supply_file, self))
        self.add_action("new.png", "New Supply", partial(new_supply, self.iface))
        self.add_action("demand.jpg", "New Demand", new_demand)
        self.add_action("check_mark.png", "Run consistency", partial(run_consistency, self))
        self.add_action("log_file.png", "Show log file", partial(load_log, self.iface))
        self.add_action("scissors.jpg", "Break Link", partial(break_link, self))
        self.add_action("invert.svg", "Invert link direction", partial(link_revert, self))
        self.add_action("link_from_template.png", "Create Link from template", partial(create_link, self))
        self.add_action("intersection.jpg", "Show intersection", self.show_specific_intersection)
        self.add_action("location.svg", "Show location", self.show_specific_location)
        self.add_action("clear_intersection.jpg", "Clear display", partial(clear_intersection_display, self))
        self.add_action("check_for_errors.png", "Check for errors", self.run_database_checker)
        self.add_action("fast-forward.png", "Upgrade Databases", partial(upgrade_databases, self))
        self.add_action("folder.png", "Open project on explorer", self.open_explorer)

        # # # ########################################################################
        # ############   SAVING PROJECT CONFIGS (copied from QAequilibraE)   #########
        QgsProject.instance().readProject.connect(self.reload_project)

        temp_saving = self.iface.mainWindow().findChild(QAction, "mActionSaveProject")
        if temp_saving:
            temp_saving.triggered.connect(self.save_in_project)

    def open_explorer(self):
        if self.open_mode is None:
            return
        path = Path(self.supply_path).parent
        if platform.system() == "Windows":
            subprocess.run(["explorer", path])
        elif platform.system() == "Darwin":
            subprocess.call(["open", "-R", path])
        else:
            subprocess.Popen(["xdg-open", path])

    def set_font(self, obj):
        f = obj.font()
        f.setPointSize(11)
        obj.setFont(f)

    def add_menu_action(self, main_menu: str, text: str, function, sub=None):
        if main_menu == "Polaris":
            action = QToolButton()
            action.setText(text)
            action.clicked.connect(function)
        else:
            action = QAction(text, self.manager)
            action.triggered.connect(function)
        if sub is None:
            self.menuActions[main_menu].append(action)
        else:
            self.menuActions[main_menu][sub].append(action)

    def build_menu(self):
        for menu, actions in self.menuActions.items():
            if menu == "Polaris":
                for action in actions:
                    self.toolbar.addWidget(action)
                continue
            itemMenu = QMenu()
            self.set_font(itemMenu)
            if isinstance(actions, dict):
                for submenu, mini_actions in actions.items():
                    if submenu == "root":
                        for mini_action in mini_actions:
                            itemMenu.addAction(mini_action)
                    else:
                        new_sub_menu = itemMenu.addMenu(submenu)
                        self.set_font(new_sub_menu)
                        for mini_action in mini_actions:
                            new_sub_menu.addAction(mini_action)
            else:
                for action in actions:
                    itemMenu.addAction(action)
            itemButton = QToolButton()
            itemButton.setText(menu)
            itemButton.setPopupMode(2)
            itemButton.setMenu(itemMenu)

            self.toolbar.addWidget(itemButton)

    def add_action(self, icon_name: str, name, function):
        """Add a toolbar icon to the toolbar."""
        pth = os.path.dirname(__file__)
        icon = QIcon(os.path.join(pth, "icons", icon_name))
        action = QAction(icon, name, self.iface.mainWindow())
        action.triggered.connect(function)
        action.setEnabled(True)

        self.menubar.addAction(action)
        self.iface.addPluginToMenu(self.menu, action)
        self.actions.append(action)

    def layerRemoved(self, layer):
        layers_to_re_create = [key for key, val in self.layers.items() if val[1] == layer]

        # Clears the pool of layers
        self.layers = {key: val for key, val in self.layers.items() if val[1] != layer}

        # Re-creates in memory only the layer that was destroyed
        for layer_name in layers_to_re_create:
            self.create_layer_by_name(layer_name)

    def show_specific_intersection(self):
        if not self.open_mode:
            self.show_message_no_supply()
            return

        self.dlg1 = FeaturePicker(self, "node", "node")

    def show_specific_location(self):
        if not self.open_mode:
            self.show_message_no_supply()
            return

        dlg2 = LocationDialog(self)
        show_diag(dlg2)

    def run_database_checker(self):
        if not self.open_mode:
            self.show_message_no_supply()
            return
        dlg2 = CheckerDialog(self)
        show_diag(dlg2)

    def unload(self):
        del self.dock
        for action in self.actions:
            self.iface.removePluginMenu("&PolarisClick", action)
            self.iface.removeToolBarIcon(action)
        self.menubar.close()
        set_polaris_menu_instance(None)

    def initProcessing(self):
        self.provider = Provider()
        QgsApplication.processingRegistry().addProvider(self.provider)

    def initGui(self):
        self.initProcessing()

    def show_message_no_supply(self):
        self.show_error_message(message="You need to load a project or supply first", level=2)

    def show_message_geoconsistency_issues(self):
        self.show_error_message(message="Your project may have internal inconsistencies", level=2)

    def show_message_no_results(self):
        self.show_error_message(message="No results file available on this folder", level=2)

    def message_no_project(self):
        f = "" if self.open_mode is None else "It is not sufficient to load a supply."
        self.show_error_message(message=f"You need to load a project. {f}", level=2)

    def show_message_close_first(self):
        self.show_error_message(message=f"Close your open {self.open_mode} first", level=2)

    def show_error_message(self, message, duration=5, level=2, message_type="Error"):
        msg = self.iface.messageBar()
        msg.pushMessage(message_type, message, level=level, duration=duration)

    def load_geo_layer(self):
        sel = self.geo_layers_table.selectedItems()
        lyr = [s.text() for s in sel][0]
        self.load_layer_by_name(lyr)

    def load_layer_by_name(self, layer_name: str):
        if layer_name.lower() not in self.layers:
            self.show_error_message(f"Layer {layer_name} was not found, which is weird")
            self.create_layer_by_name(layer_name)
        layer = self.layers[layer_name.lower()][0]
        QgsProject.instance().addMapLayer(layer)
        qgis.utils.iface.mapCanvas().refresh()

    def create_layer_by_name(self, layer_name: str):
        layer = self.create_loose_layer(layer_name)
        self.layers[layer_name.lower()] = [layer, layer.id()]
        load_editor_styles(layer, layer_name.lower())

    def create_loose_layer(self, layer_name: str) -> QgsVectorLayer:
        uri = QgsDataSourceUri()
        uri.setDatabase(self.supply_path)
        uri.setDataSource("", layer_name, "geo")
        layer = QgsVectorLayer(uri.uri(), layer_name, "spatialite")
        # layer.editingStopped.connect(partial(run_consistency, self))
        return layer

    def save_in_project(self):
        """Saves temporary layers to the project using QGIS saving buttons."""
        if not self.open_mode:
            return
        elif self.open_mode == "project":
            folder = self.polaris_project.model_path
            QgsExpressionContextUtils.setProjectVariable(
                QgsProject.instance(), "polaris_path", str(self.polaris_project.model_path)
            )
        elif self.open_mode == "supply":
            folder = Path(self.supply_path).parent
            QgsExpressionContextUtils.setProjectVariable(QgsProject.instance(), "polaris_supply", self.supply_path)

        else:
            raise ValueError("Invalid open mode")

        QgsExpressionContextUtils.setProjectVariable(QgsProject.instance(), "polaris_iter", self.cob_iter.currentText())

        output_file_path = os.path.join(str(folder), "qgis_temp_layers.sqlite")

        file_exists = True if os.path.isfile(output_file_path) else False

        for layer in QgsProject.instance().mapLayers().values():
            if layer.isTemporary():
                layer_name = layer.name() + f"_{uuid4().hex}"
                options = QgsVectorFileWriter.SaveVectorOptions()
                options.driverName = "SQLite"
                options.layerName = layer_name
                if file_exists:
                    options.actionOnExistingFile = QgsVectorFileWriter.CreateOrOverwriteLayer

                transform_context = QgsProject.instance().transformContext()

                error = QgsVectorFileWriter.writeAsVectorFormatV3(layer, output_file_path, transform_context, options)

                if error[0] == QgsVectorFileWriter.NoError:
                    layer.setDataSource(output_file_path + f"|layername={layer_name}", layer.name(), "ogr")

                file_exists = True
        QgsProject.instance().write()

    def reload_project(self):
        """Opens Polaris project when opening a QGIS project containing a Polaris model."""
        # Copied from QAequilibraE

        # Checks if project contains a Polaris model
        path = QgsProject.instance().customVariables()
        if "polaris_path" in path:
            load_project_from_path(self, Path(str(path["polaris_path"])))

            if "polaris_iter" in path:
                self.cob_iter.setCurrentText(str(path["polaris_iter"]))
                change_iteration(self)
        elif "polaris_supply" in path:
            supply_setter(self, str(path["polaris_supply"]))
        else:
            return

        # Checks if the layers in the project have the same database path as the Polaris project layers.
        # if so, we replace the path in self.layers
        for lyr in QgsProject.instance().mapLayers().values():
            if "sqlite" not in lyr.source():
                continue
            pth = Path(lyr.source().replace("\\\\", "/").split("'")[1])
            if pth.parent != Path(self.supply_path).parent or pth.name != Path(self.supply_path).name:
                print(1)
                continue
            self.layers[str(lyr.name()).lower()] = [lyr, lyr.id()]

    @property
    def supply_path(self):
        if self.open_mode == "project":
            return str(self.polaris_project.supply_file)

        if self.open_mode == "supply":
            return str(self._network.path_to_file)

    @property
    def freight_path(self):
        if self.open_mode == "project":
            return self.polaris_project.freight_file
        return None

    @property
    def result_path(self):
        if self.open_mode == "project":
            return self.polaris_project.result_file
        return None

    @property
    def result_h5_path(self):
        if self.open_mode == "project":
            return self.polaris_project.result_h5_file
        return None

    @property
    def demand_path(self):
        if self.open_mode == "project":
            return self.polaris_project.demand_file
        return None

    @property
    def network(self):
        if self.open_mode == "project":
            return self.polaris_project.network

        if self.open_mode == "supply":
            return self._network

    @property
    def supply_conn(self):
        return commit_and_close(self.supply_path, spatial=True)
