# -*- coding: utf-8 -*-
"""
/***************************************************************************
 OMRAT
                                 A QGIS plugin
 This is an open source implementation of Pedersens equations from 1995
 Generated by Plugin Builder: http://g-sherman.github.io/Qgis-Plugin-Builder/
                              -------------------
        begin                : 2022-12-22
        git sha              : $Format:%H$
        copyright            : (C) 2022 by Axel Hörteborn
        email                : axel.horteborn@ri.se
 ***************************************************************************/

/***************************************************************************
 *                                                                         *
 *   This program is free software; you can redistribute it and/or modify  *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 2 of the License, or     *
 *   any later version.                                   *
 *                                                                         *
 ***************************************************************************/
"""
from typing import Any, Callable
from functools import partial
from qgis.PyQt.QtCore import QSettings, QTranslator, QCoreApplication, QMetaType, QUrl
from qgis.PyQt.QtGui import QIcon, QAction, QCloseEvent, QDesktopServices
from qgis.PyQt.QtWidgets import QMenuBar, QWidget, QFileDialog, QToolBar, QMessageBox
from qgis._core import QgsVectorDataProvider
from qgis.core import (QgsVectorLayer, QgsFeature, QgsGeometry, QgsLineString, QgsPoint, QgsProject,
                       QgsField, QgsFields)
from qgis.core import QgsPointXY, QgsMessageLog, Qgis, QgsApplication

from qgis.gui import QgsMapToolPan, QgisInterface
import copy
import gc
import sys
import os
import os.path
import time

sys.path.append('.')
# Initialize Qt resources from file resources.py
from resources import *

# Import the code for the DockWidget
from compute.run_calculations import Calculation
from compute.calculation_task import CalculationTask
from geometries.handle_qgis_iface import HandleQGISIface
from omrat_utils.causation_factors import CausationFactors
from omrat_utils.handle_ais import AIS
from omrat_utils.handle_settings import DriftSettings
from omrat_utils.handle_traffic import Traffic
from omrat_utils.handle_distributions import Distributions
from omrat_utils.storage import Storage
from omrat_utils.handle_object import OObject
from omrat_utils.handle_ship_cat import ShipCategories
from omrat_utils.handle_consequence import Consequence
from omrat_utils.handle_junctions import Junctions
from omrat_utils.gather_data import GatherData
from omrat_widget import OMRATMainWidget
from geometries.drift import DriftCorridorGenerator
from geometries.drift_corridor_task_v2 import DriftCorridorTask


from omrat_utils.compare_mixin import CompareMixin
from omrat_utils.accident_results_mixin import AccidentResultsMixin
from omrat_utils.drift_analysis_mixin import DriftAnalysisMixin
from omrat_utils.iwrap_io_mixin import IwrapIOMixin
from omrat_utils.run_history_mixin import RunHistoryMixin


class OMRAT(
    IwrapIOMixin, CompareMixin, DriftAnalysisMixin, RunHistoryMixin,
    AccidentResultsMixin,
):
    """QGIS Plugin Implementation.

    Methods are progressively being extracted into mixins under
    ``omrat_utils/`` to keep this module readable.  Each mixin owns a
    cohesive group of slots / helpers; see the ``omrat_utils`` package
    docstrings for what lives where.
    """
    def __init__(self, iface:QgisInterface, testing:bool=False):
        """Constructor.

        :param iface: An interface instance that will be passed to this class
            which provides the hook by which you can manipulate the QGIS
            application at run time.
        :type iface: QgsInterface
        """
        # Save reference to the QGIS interface
        self.iface = iface
        self.testing = testing
        # initialize plugin directory
        self.plugin_dir = os.path.dirname(__file__)
        if not self.testing:
            # initialize locale
            locale: str = QSettings().value('locale/userLocale')[0:2]
            locale_path: str = os.path.join(
                self.plugin_dir,
                'i18n',
                f'OMRAT_{locale}.qm')

            if os.path.exists(locale_path):
                self.translator: QTranslator = QTranslator()
                self.translator.load(locale_path)
                QCoreApplication.installTranslator(self.translator)

        # Declare instance attributes
        self.actions: list[Any] = []
        self.menu: str = self.tr(u'&Open Maritime Risk Analysis Tool')
        self.toolbar: QToolBar = self.iface.addToolBar(u'OMRAT')
        self.toolbar.setObjectName(u'OMRAT')

            #print "** INITIALIZING OMRAT"
        self.pluginIsActive = False
        self.been_closed = False
        self.segment_id = 0
        self.traffic_data: dict[str, dict[str, dict[str, Any]]] = {}
        self.segment_data: dict[str, Any] = {}
        # Project-level traffic scaling -- global spinbox value + per-row
        # "follow global" bool list.  Compute reads the per-cell matrix
        # in ``traffic_data[...]['Scaling (%)']`` (kept in sync by the
        # broadcast logic in ``omrat_utils.handle_traffic``); this dict
        # is only the UI's bookkeeping for the global control.
        from omrat_utils.handle_traffic import _default_traffic_scaling
        self.traffic_scaling: dict[str, Any] = _default_traffic_scaling()
        # Write-once snapshot of segment endpoints at first creation.
        # Audit report diffs the live ``segment_data`` against this to
        # surface waypoint moves that no other input reveals.
        self.segments_imported: dict[str, dict[str, str]] = {}
        self.main_widget = OMRATMainWidget(plugin=self)
        self.qgis_geoms: HandleQGISIface = HandleQGISIface(self)
        self.drift_settings: DriftSettings = DriftSettings(self)
        self.causation_f: CausationFactors = CausationFactors(self)
        self.calc: Calculation | None = Calculation(self)
        self.ship_cat = ShipCategories(self)
        self.consequence = Consequence(self)
        self.junctions = Junctions(self)
        self.ais = AIS(self)
        self.traffic = Traffic(self, self.main_widget)
        self.distributions = Distributions(self)
        self.object = OObject(self)
        self.drift_values = self.drift_settings.drift_values
        self.causation_values = self.causation_f.data
        self.drift_corridor_layers: list[QgsVectorLayer] = []  # Track corridor layers for cleanup
        # Layers loaded from a saved per-run GeoPackage via the
        # "Add results to map" button.  Tracked here so they get
        # removed when the plugin is unloaded or the project is
        # cleared -- otherwise they linger in the layer panel after
        # a Plugin Reloader cycle.
        self._history_layers: list[QgsVectorLayer] = []
        # Strong reference to the active CalculationTask while it runs.
        # Cleared from ``_on_calculation_finished`` /
        # ``_on_calculation_failed``.  See ``run_calculation`` for why.
        self._current_task = None

    def tr(self, message:str):
        """Get the translation for a string using Qt translation API.

        We implement this ourselves since we do not inherit QObject.

        :param message: String for translation.
        :type message: str, QString

        :returns: Translated version of message.
        :rtype: QString
        """
        return QCoreApplication.translate('OMRAT', message)

    def add_action(
        self,
        icon_path:str,
        text:str,
        callback: Callable[[], None],
        enabled_flag:bool=True,
        add_to_menu:bool=True,
        add_to_toolbar:bool=True,
        status_tip:str|None=None,
        whats_this:str|None=None,
        parent:QWidget|None=None) -> QAction:
        """Add a toolbar icon to the toolbar.

        :param icon_path: Path to the icon for this action. Can be a resource
            path (e.g. ':/plugins/foo/bar.png') or a normal file system path.
        :type icon_path: str

        :param text: Text that should be shown in menu items for this action.
        :type text: str

        :param callback: Function to be called when the action is triggered.
        :type callback: function

        :param enabled_flag: A flag indicating if the action should be enabled
            by default. Defaults to True.
        :type enabled_flag: bool

        :param add_to_menu: Flag indicating whether the action should also
            be added to the menu. Defaults to True.
        :type add_to_menu: bool

        :param add_to_toolbar: Flag indicating whether the action should also
            be added to the toolbar. Defaults to True.

        :param status_tip: Optional text to show in a popup when mouse pointer
            hovers over the action.
        :type status_tip: str

        :param parent: Parent widget for the new action. Defaults None.
        :type parent: QWidget

        :param whats_this: Optional text to show in the status bar when the
            mouse pointer hovers over the action.

        :returns: The action that was created. Note that the action is also
            added to self.actions list.
        :rtype: QAction
        """

        icon = QIcon(icon_path)
        action = QAction(icon, text, parent)
        action.triggered.connect(callback)
        action.setEnabled(enabled_flag)

        if status_tip is not None:
            action.setStatusTip(status_tip)

        if whats_this is not None:
            action.setWhatsThis(whats_this)

        if add_to_toolbar:
            self.toolbar.addAction(action)

        if add_to_menu:
            self.iface.addPluginToMenu(
                self.menu,
                action)

        self.actions.append(action)

        return action

    def initGui(self):
        """Create the menu entries and toolbar icons inside the QGIS GUI."""

        icon_path = self.plugin_dir + '/img/icon.png'
        self.add_action(
            icon_path,
            text=self.tr(u'Omrat'),
            callback=self.run,
            parent=self.iface.mainWindow())

        # Connect to project cleared signal to clean up drift corridors
        QgsProject.instance().cleared.connect(self._on_project_cleared)

    #--------------------------------------------------------------------------

    def _on_project_cleared(self):
        """Called when the QGIS project is cleared/reloaded.

        Clears drift corridor layers and any cached data.
        """
        QgsMessageLog.logMessage("Project cleared - cleaning up drift corridors", "OMRAT", Qgis.Info)
        self.drift_corridor_layers.clear()
        # The QGIS project drop also wipes any history-loaded layers,
        # so just drop our own references; no removeMapLayer needed.
        self._history_layers.clear()
        # Reset any task reference
        if hasattr(self, '_drift_task'):
            self._drift_task = None

    def _remove_history_layers_from_project(self) -> None:
        """Remove every history-loaded layer from the QGIS project.

        Called from ``clear_model`` and ``unload``.  Layers are
        appended to ``self._history_layers`` by
        :meth:`_add_selected_run_to_map`.
        """
        for layer in self._history_layers:
            try:
                project = QgsProject.instance()
                if project is not None and layer is not None:
                    project.removeMapLayer(layer.id())
            except Exception:
                pass
        self._history_layers.clear()

    def show_error_popup(self, message: str, function_name:str) -> None:
        msg_box = QMessageBox()
        msg_box.setIcon(QMessageBox.Icon.Critical)
        msg_box.setWindowTitle(self.tr(f"Failure in {function_name}"))
        msg_box.setText(self.tr("A failure occurred:"))
        msg_box.setInformativeText(message)
        msg_box.exec()

    def _confirm_before_load(self, action_label: str) -> str:
        """Show a confirmation dialog before loading/importing data.

        Returns 'clear', 'merge', or 'cancel'.
        If the current model is empty the dialog is skipped and 'clear' is returned.
        """
        has_data = (
            len(self.segment_data) > 0
            or len(self.traffic_data) > 0
            or self.main_widget.twRouteList.rowCount() > 0
            or self.main_widget.twDepthList.rowCount() > 0
            or self.main_widget.twObjectList.rowCount() > 0
        )
        if not has_data:
            return 'clear'

        msg_box = QMessageBox(self.main_widget)
        msg_box.setIcon(QMessageBox.Icon.Question)
        msg_box.setWindowTitle(self.tr(f'{action_label} Project'))
        msg_box.setText(self.tr(
            'The current model contains data.\n\n'
            'Do you want to clear all current data before loading,\n'
            'or merge the new data into the existing model?'
        ))
        btn_clear = msg_box.addButton(self.tr('Clear && Load'), QMessageBox.ButtonRole.AcceptRole)
        btn_merge = msg_box.addButton(self.tr('Merge'), QMessageBox.ButtonRole.ActionRole)
        msg_box.addButton(QMessageBox.StandardButton.Cancel)
        msg_box.setDefaultButton(btn_clear)
        msg_box.exec()

        clicked = msg_box.clickedButton()
        if clicked == btn_clear:
            return 'clear'
        elif clicked == btn_merge:
            return 'merge'
        return 'cancel'

    def clear_model(self) -> None:
        """Reset all model state and UI to defaults."""
        QgsMessageLog.logMessage("Clearing current model", "OMRAT", Qgis.Info)

        # --- 1. Clear QGIS layers ---
        self.qgis_geoms.clear()
        self.object.clear()

        # Remove drift corridor layers
        for layer in self.drift_corridor_layers:
            try:
                project = QgsProject.instance()
                if project is not None and layer is not None:
                    project.removeMapLayer(layer.id())
            except Exception:
                pass
        self.drift_corridor_layers.clear()

        # Remove any layers loaded from previous-run GeoPackages.
        self._remove_history_layers_from_project()

        # Reset cached Calculation result-layer attributes.  The plugin
        # no longer auto-adds memory layers to the canvas during a run --
        # results are persisted to per-run GeoPackages instead -- so
        # there is nothing to remove from the project here.  Just clear
        # the attribute references so a re-run starts clean.
        if self.calc is not None:
            for attr in (
                'allision_result_layer', 'grounding_result_layer',
                'collision_line_layer', 'collision_point_layer',
                'powered_grounding_layer', 'powered_allision_layer',
            ):
                setattr(self.calc, attr, None)

        # --- 2. Clear internal data structures ---
        self.traffic_data = {}
        self.segment_data = {}
        from omrat_utils.handle_traffic import _default_traffic_scaling
        self.traffic_scaling = _default_traffic_scaling()
        self.segments_imported = {}
        self.segment_id = 0

        # Reset drift settings to defaults
        from compute.iwrap_defaults import default_drift_values
        default_drift = default_drift_values()
        self.drift_values = default_drift
        self.drift_settings.drift_values = default_drift

        # Reset causation factors to defaults.  Use the full default
        # set defined in ``CausationFactors.__init__`` so collision
        # keys (``headon``, ``overtaking``, ...) survive Clear.
        from omrat_utils.causation_factors import CausationFactors
        self.causation_f.data = CausationFactors(self).data
        self.causation_values = self.causation_f.data

        # Reset distributions tracking
        self.distributions.last_id = '1'
        self.distributions.new_id = '1'

        # Prevent traffic table updates during clearing
        self.traffic.run_update = False
        self.traffic.traffic_data = self.traffic_data
        self.traffic.c_seg = ''

        # --- 3. Clear UI widgets ---
        self.reset_route_table()

        self.main_widget.twDepthList.setRowCount(0)
        self.main_widget.twObjectList.setRowCount(0)

        self.main_widget.cbTrafficSelectSeg.clear()
        self.main_widget.cbTrafficDirectionSelect.clear()
        self.traffic.set_table_headings()

        # Distribution fields
        dist_widgets = [
            self.main_widget.leNormMean1_1, self.main_widget.leNormMean1_2,
            self.main_widget.leNormMean1_3, self.main_widget.leNormMean2_1,
            self.main_widget.leNormMean2_2, self.main_widget.leNormMean2_3,
            self.main_widget.leNormStd1_1, self.main_widget.leNormStd1_2,
            self.main_widget.leNormStd1_3, self.main_widget.leNormStd2_1,
            self.main_widget.leNormStd2_2, self.main_widget.leNormStd2_3,
            self.main_widget.leNormWeight1_1, self.main_widget.leNormWeight1_2,
            self.main_widget.leNormWeight1_3, self.main_widget.leNormWeight2_1,
            self.main_widget.leNormWeight2_2, self.main_widget.leNormWeight2_3,
            self.main_widget.leUniformMin1, self.main_widget.leUniformMax1,
            self.main_widget.leUniformMin2, self.main_widget.leUniformMax2,
        ]
        for w in dist_widgets:
            w.setText('0')
        self.main_widget.sbUniformP1.setValue(0)
        self.main_widget.sbUniformP2.setValue(0)

        # Result fields
        result_fields = [
            self.main_widget.LEPDriftAllision,
            self.main_widget.LEPPoweredAllision,
            self.main_widget.LEPDriftingGrounding,
            self.main_widget.LEPPoweredGrounding,
            self.main_widget.LEPOvertakingCollision,
            self.main_widget.LEPHeadOnCollision,
            self.main_widget.LEPCrossingCollision,
            self.main_widget.LEPMergingCollision,
        ]
        for field in result_fields:
            field.setText('')

        self.main_widget.LEModelName.setText('')
        # The Previous-runs table reflects the master run-history DB,
        # not the currently-loaded project, so reload from the DB
        # instead of wiping it -- the rows survive a project clear /
        # load.
        try:
            self.refresh_previous_runs_table()
        except Exception:
            self.main_widget.TWPreviousRuns.setRowCount(0)
        self.main_widget.laDir1.setText('')
        self.main_widget.laDir2.setText('')

        if hasattr(self.main_widget, 'label_drift_status'):
            self.main_widget.label_drift_status.setText('')

        # Remove distribution plot canvas if present
        if self.distributions.canvas is not None:
            try:
                self.main_widget.DistributionWidget.removeWidget(self.distributions.canvas)
                self.distributions.canvas.deleteLater()
                self.distributions.canvas = None
            except Exception:
                pass

        # Re-enable traffic updates
        self.traffic.run_update = True

        self.iface.mapCanvas().refresh()
        QgsMessageLog.logMessage("Model cleared successfully", "OMRAT", Qgis.Info)

    def onClosePlugin(self):
        """Cleanup necessary items here when plugin main_widget is closed"""

        # Disconnect signals
        if self.main_widget.is_active:
            self.main_widget.pbAddRoute.clicked.disconnect()
            self.main_widget.pbStopRoute.clicked.disconnect()
            self.main_widget.pbLoadRoute.clicked.disconnect()
            self.main_widget.pbRemoveRoute.clicked.disconnect()
            self.main_widget.pbAddSimpleDepth.clicked.disconnect()
            self.main_widget.pbLoadDepth.clicked.disconnect()
            self.main_widget.pbRemoveDepth.clicked.disconnect()
            self.main_widget.pbAddSimpleObject.clicked.disconnect()
            self.main_widget.pbLoadObject.clicked.disconnect()
            self.main_widget.pbRemoveObject.clicked.disconnect()
            self.main_widget.pbRunModel.clicked.disconnect()

        # Set plugin state to inactive
        self.pluginIsActive = False
        self.main_widget.is_active = False

    def unload(self):
        """Removes the plugin menu item and icon from QGIS GUI."""
        # Disconnect project cleared signal
        try:
            QgsProject.instance().cleared.disconnect(self._on_project_cleared)
        except Exception:
            pass

        # Call the unload method of Traffic
        self.traffic.unload()

        # Call the unload method of AIS
        self.ais.unload()

        # Unload geometries
        self.object.unload()
        self.qgis_geoms.unload()

        # Remove drift corridor layers
        for layer in self.drift_corridor_layers:
            try:
                project = QgsProject.instance()
                if project is not None and layer is not None:
                    project.removeMapLayer(layer.id())
            except Exception:
                pass
        self.drift_corridor_layers.clear()

        # Remove any layers loaded from previous-run GeoPackages.
        self._remove_history_layers_from_project()

        # Remove result layers (allision/grounding/collision/powered visualization)
        if self.calc is not None:
            result_layer_attrs = [
                'allision_result_layer', 'grounding_result_layer',
                'collision_line_layer', 'collision_point_layer',
                'powered_grounding_layer', 'powered_allision_layer',
            ]
            for attr in result_layer_attrs:
                layer = getattr(self.calc, attr, None)
                try:
                    project = QgsProject.instance()
                    if project is not None and layer is not None:
                        project.removeMapLayer(layer.id())
                except Exception:
                    pass
                setattr(self.calc, attr, None)

        # Remove actions from menu and toolbar
        for action in self.actions:
            self.iface.removePluginMenu(
                self.tr(u'&Open Maritime Risk Analysis Tool'),
                action)
            self.iface.removeToolBarIcon(action)

        # Remove the toolbar
        if hasattr(self, 'toolbar'):
            del self.toolbar    

        # Disconnect signals and clean up dock widget
        try:
            self.onClosePlugin()
        except Exception as e:
            print(f"Error during onClosePlugin: {e}")
        try:
            self.main_widget.pbGetGebcoDephts.clicked.disconnect()
            self.main_widget.PBUpdateDepthIntervals.clicked.disconnect()
            self.main_widget.deleteLater()
        except TypeError:
            pass
        gc.collect()
        # Remove all plugin modules from sys.modules to ensure fresh code on reload
        for folder in ['omrat_utils', 'compute', 'geometries']:
            to_remove = [m for m in sys.modules if m.startswith(folder)]
            for m in to_remove:
                QgsMessageLog.logMessage(f"{m} unloaded", "OMRAT", Qgis.Info)
                del sys.modules[m]
        self.been_close = True
        QgsMessageLog.logMessage("Plugin unloaded", "OMRAT", Qgis.Info)
        
    def point4326_from_wkt(self, coord_str:str, crs:str) -> QgsPoint:
        if '(' in coord_str:
            coord_str = coord_str.split('(')[1].split(')')[0]
        coords_str: list[str] = coord_str.split(' ')
        coords: list[float] = [float(coord) for coord in coords_str]
        q_point = QgsPoint(coords[0], coords[1])
        return q_point
        
    def load_lines(self, data:dict[str, dict[str, Any]]) -> None:
        self.segment_id = 0
        crs = self.iface.mapCanvas().mapSettings().destinationCrs().authid()
        for key, seg_data in data["segment_data"].items():
            # Freshly-drawn legs may not have Segment_Id / Route_Id in the
            # dict yet (those are filled in by GatherData.get_segment_tbl
            # on save).  Fall back to the dict key and route 1 so the
            # validation-pass reload doesn't KeyError mid-rebuild and
            # wipe the canvas + route table.
            seg_id = seg_data.get('Segment_Id', key)
            route_id = seg_data.get('Route_Id', 1)
            seg_data.setdefault('Segment_Id', seg_id)
            seg_data.setdefault('Route_Id', route_id)
            name = f"Segment {route_id} - {seg_id}"
            vl = QgsVectorLayer(f"LineString?crs=EPSG:4326", name, "memory")
            QgsProject.instance().addMapLayer(vl)
            
            self.segment_id += 1
            if not vl.isEditable():
                vl.startEditing()
            fields = self.qgis_geoms.create_fields()
            provider: QgsVectorDataProvider | None = vl.dataProvider()
            if provider is None:
                return
            provider.addAttributes(fields.toList())
            vl.updateFields()

            # Create feature
            fet = QgsFeature(fields)
            start = self.point4326_from_wkt(seg_data["Start_Point"], crs)
            end = self.point4326_from_wkt(seg_data["End_Point"], crs)
            fet.setGeometry(QgsLineString([start, end]))
            fet.setAttributes([seg_data["Segment_Id"], seg_data["Route_Id"], seg_data["Start_Point"],
                               seg_data["End_Point"], f"LEG_{seg_data['Route_Id']}_{seg_data['Segment_Id']}"])
            self.qgis_geoms.style_layer(vl)
            # Ensure feature id matches Segment_Id so geometry edits update correct row
            try:
                fid = int(str(seg_data["Segment_Id"]))
            except Exception:
                fid = self.segment_id
            fet.setId(fid)
            provider.addFeature(fet)
            self.qgis_geoms.label_layer(vl)
            vl.updateExtents()
            
            # Validate geometry
            if not fet.geometry().isGeosValid():
                print(f"Invalid geometry for segment {seg_data['Segment_Id']}")
                continue
            # Keep layer editable and connect geometry change handler
            edit_buffer = vl.editBuffer()
            if edit_buffer is not None:
                edit_buffer.geometryChanged.connect(partial(self.qgis_geoms.on_geometry_changed_wrapper, fid))
            vl.triggerRepaint()
            self.qgis_geoms.vector_layers.append(vl)

            # Preserve direction labels mapping for AIS
            try:
                self.qgis_geoms.leg_dirs[str(seg_data["Segment_Id"])]=list(seg_data.get("Dirs", []))
            except Exception:
                pass

        # After loading, ensure subsequent new segments use a higher id
        try:
            max_id = max(int(str(v.get('Segment_Id'))) for v in data['segment_data'].values())
        except Exception:
            max_id = self.segment_id
        self.qgis_geoms.segment_id = max_id
        self.segment_id = max_id
        # Refresh map canvas
        self.iface.mapCanvas().refresh()    

    def reset_route_table(self) -> None:
        self.main_widget.twRouteList.setColumnCount(7)
        self.main_widget.twRouteList.setHorizontalHeaderLabels(['Segment_Id', 'Route_Id', 'Leg_name',
                                                            'Start_Point', 'End_Point', 'Width', 'Update AIS'])
        self.main_widget.twRouteList.setColumnWidth(0, 75)
        self.main_widget.twRouteList.setColumnWidth(1, 75)
        self.main_widget.twRouteList.setColumnWidth(2, 75)
        self.main_widget.twRouteList.setColumnWidth(3, 125)
        self.main_widget.twRouteList.setColumnWidth(4, 125)
        self.main_widget.twRouteList.setColumnWidth(5, 75)
        self.main_widget.twRouteList.setColumnWidth(6, 75)
        self.main_widget.twRouteList.setRowCount(0)
    
    def run_traffic_module(self) -> None:
        self.traffic.traffic_data = self.traffic_data
        self.traffic.fill_cbTrafficSelectSeg()
        self.traffic.update_direction_select()
        
    def save_work(self):
        store = Storage(self)
        store.store_all()

    def load_work(self):
        store = Storage(self)
        file_path = store.select_file()
        if not file_path:
            return
        choice = self._confirm_before_load('Load')
        if choice == 'cancel':
            return
        if choice == 'clear':
            self.clear_model()
        store.load_from_path(file_path)
               
    def run_calculation(self):
        """Run calculation in a background QgsTask for better UI responsiveness."""
        gd = GatherData(self)
        data = gd.get_all_for_save()
        if self.calc is None:
            return

        # Refuse to run unless an output folder AND a model name are
        # configured.  The button should already be disabled, but defend
        # against direct calls / keyboard shortcuts.
        out_dir = self._get_output_dir()
        try:
            run_name = (self.main_widget.LEModelName.text() or '').strip()
        except Exception:
            run_name = ''
        if out_dir is None or not run_name:
            missing: list[str] = []
            if not run_name:
                missing.append('Name of the model')
            if out_dir is None:
                missing.append('Output folder (File path)')
            QMessageBox.information(
                self.main_widget,
                self.tr('Quick start: missing inputs'),
                self.tr(
                    "Before running the model, fill in the following on the "
                    "Run Analysis tab:\n\n  - {fields}\n\n"
                    "Each run writes a result GeoPackage, an input snapshot "
                    "(.omrat) and a markdown summary to that folder, named "
                    "after the model."
                ).format(fields="\n  - ".join(missing)),
            )
            return

        # Create and configure the background task
        task = CalculationTask(
            description='OMRAT: Starting calculations...',
            calc_object=self.calc,
            data=data
        )

        # Connect signals for progress updates and completion
        task.progress_updated.connect(self._on_calculation_progress)
        task.calculation_finished.connect(self._on_calculation_finished)
        task.calculation_failed.connect(self._on_calculation_failed)

        # CRITICAL: keep a strong Python reference to the task object.
        # QGIS's task manager only holds a C++ pointer; without a Python
        # reference, the QObject wrapper can be garbage-collected before
        # ``finished()`` fires, which silently breaks our signals.
        self._current_task = task

        # Track elapsed time for the run-history "Duration" column.
        self._run_started_at = time.monotonic()

        # Snapshot the input dict at kick-off so the .omrat archive
        # written after a successful run reflects exactly what the
        # calc consumed, not whatever the user has typed since.
        self._run_input_snapshot = copy.deepcopy(data)

        # Add task to QGIS task manager (runs in background)
        QgsApplication.taskManager().addTask(task)

        QgsMessageLog.logMessage(
            'Calculation started in background. Check task manager for progress.',
            'OMRAT',
            Qgis.Info
        )

    def _on_calculation_progress(self, completed: int, total: int, message: str) -> None:
        """Handle progress updates from the calculation task."""
        # Log progress to QGIS message log
        QgsMessageLog.logMessage(
            f"Progress: {completed}/{total} - {message}",
            'OMRAT',
            Qgis.Info
        )

    def _on_calculation_finished(self, calc_object: Any) -> None:
        """Handle successful completion of calculation."""
        QgsMessageLog.logMessage(
            'Calculation completed successfully!',
            'OMRAT',
            Qgis.Success
        )
        # Catastrophe-level exceedance table on the Run Analysis tab.
        try:
            self._populate_catastrophe_results_table(
                getattr(calc_object, 'consequence_result', None),
            )
        except Exception as e:
            QgsMessageLog.logMessage(
                f'Failed to populate catastrophe-results table: {e}',
                'OMRAT', Qgis.Warning,
            )
        # Auto-save the run to the history GeoPackage so the user can
        # reload / compare it later.  Failures here are non-fatal --
        # the on-canvas results and line-edits are already populated.
        try:
            self._auto_save_run(calc_object)
        except Exception as e:
            import traceback
            QgsMessageLog.logMessage(
                f'Failed to save run history: {e}\n{traceback.format_exc()}',
                'OMRAT', Qgis.Warning,
            )
        try:
            self.refresh_previous_runs_table()
            # Select the just-completed run (always the top row of the
            # history table) so the View buttons default to the latest
            # run.  Without this the user would see the "select a run"
            # popup right after their first model run.
            self._select_latest_previous_run()
        except Exception as e:
            import traceback
            QgsMessageLog.logMessage(
                f'Failed to refresh previous-runs table: {e}\n'
                f'{traceback.format_exc()}',
                'OMRAT', Qgis.Warning,
            )
        # Drop the strong reference now that the task has finished.
        self._current_task = None

    def _on_calculation_failed(self, error_msg: str) -> None:
        """Handle calculation failure (overrides earlier definition).

        Keeps the message-box behaviour but also drops the strong task
        reference so the next run starts clean.
        """
        QgsMessageLog.logMessage(
            f'Calculation failed: {error_msg}',
            'OMRAT', Qgis.Critical,
        )
        QMessageBox.critical(
            None, 'Calculation Failed',
            f'The calculation failed with the following error:\n\n{error_msg}',
        )
        self._current_task = None

    # Run-history flow (auto-save, sidecars, markdown report) is
    # provided by ``RunHistoryMixin`` (see
    # ``omrat_utils.run_history_mixin``).  These methods used to live
    # here.

    # Compare-tab slots are provided by ``CompareMixin`` (see
    # ``omrat_utils.compare_mixin``).  They used to live here.

    # ------------------------------------------------------------------
    # Accident-results QTable (replaces the legacy LEP* line-edit grid)
    # ------------------------------------------------------------------
    # (Accident type label, LEP* widget name on main_widget, View pushbutton
    #  name on main_widget, slot on self).  Order is the row order of the
    #  table.
    # Accident-results table setup, ``LEP*`` legacy widgets, and the
    # per-row View dispatcher all live in ``AccidentResultsMixin``
    # (see ``omrat_utils.accident_results_mixin``).

    # ------------------------------------------------------------------
    # Output folder + Run Model gating
    # ------------------------------------------------------------------
    # ------------------------------------------------------------------
    # Previous-runs table, output-folder gating, save-run flow are all
    # provided by ``RunHistoryMixin`` (see
    # ``omrat_utils.run_history_mixin``).
    # ------------------------------------------------------------------

    def choose_report_path(self) -> None:
        """Open a Save File dialog and store path in LEReportPath."""
        # Initial directory & filename from current field if present
        initial = ''
        try:
            if hasattr(self.main_widget, 'LEReportPath') and self.main_widget.LEReportPath is not None:
                initial = self.main_widget.LEReportPath.text() or ''
        except Exception:
            initial = ''
        filename, _ = QFileDialog.getSaveFileName(
            self.main_widget,
            self.tr('Select report file'),
            initial,
            'Markdown Files (*.md);;All Files (*.*)'
        )
        if filename:
            try:
                self.main_widget.LEReportPath.setText(filename)
            except Exception:
                pass

    # IWRAP import / export slots are provided by ``IwrapIOMixin``;
    # accident-results table + per-row View dispatcher live in
    # ``AccidentResultsMixin``.  See ``omrat_utils/`` for the split.

    def open_drift(self):
        self.drift_settings.run()

    def open_ship_categories(self):
        self.ship_cat.run()

    def open_causation_factors(self):
        self.causation_f.run()

    def stop_route(self):
        """Stop the current route, commit changes, and set the active tool to 'Pan Map'."""
        # Disable the "Stop Route" button
        self.main_widget.pbStopRoute.setEnabled(False)

        for layer in self.qgis_geoms.vector_layers:
            if layer.isEditable():
                layer.commitChanges()  # Save changes
            layer.triggerRepaint()  # Ensure the layer is refreshed

        # Advance the route id on the leg-drawing handler -- that's the
        # attribute the next "Add route" pass actually reads when it
        # stamps Route_Id / Leg_name on freshly-drawn legs.
        self.qgis_geoms.cur_route_id += 1
        self.qgis_geoms.current_start_point = None

        # Set the active tool to "Pan Map"
        canvas: QgsMapCanvas | None = self.iface.mapCanvas()
        map_tool_pan = QgsMapToolPan(canvas)
        if canvas is not None:
            canvas.setMapTool(map_tool_pan)
        actionPan: QAction | None = self.iface.actionPan()
        if actionPan is not None:
            actionPan.trigger()
        
    def ais_settings(self) -> None:
        self.ais.run()

    def open_db_setup_wizard(self) -> None:
        # Lazy import: keeps plugin startup unaffected even if the wizard
        # module pulls in something heavyweight in the future.
        from omrat_utils.handle_db_setup import DbSetupWizard
        wiz = DbSetupWizard(self.main_widget)
        wiz.exec()

    def open_junction_dialog(self) -> None:
        """Open the junction transition-matrix editor."""
        from omrat_utils.junction_matrix_dialog import open_junction_dialog
        # Make sure the registry reflects the current segment set before
        # the user starts editing — newly added legs would otherwise be
        # invisible.
        if getattr(self, 'junctions', None) is not None:
            self.junctions.rebuild_from_segments(prefer_user=True)
        open_junction_dialog(self)

    _DOCS_BASE_URL = "https://axelande.github.io/OMRAT/"

    def _open_docs_url(self, page: str) -> None:
        """Open a page of the published Sphinx docs in the user's browser."""
        QDesktopServices.openUrl(QUrl(self._DOCS_BASE_URL + page))

    def _open_external_url(self, url: str) -> None:
        QDesktopServices.openUrl(QUrl(url))


    def update_ais(self) -> None:
        # The "Update all distributions" button does three things in
        # sequence: (1) validate routes (snap close waypoints, split
        # crossing legs); (2) refresh per-leg AIS traffic; (3) refresh
        # the junction transition matrices, preferring AIS-derived
        # shares where available.
        try:
            from omrat_utils.route_validation_ui import run_validation_pass
            outcome = run_validation_pass(self)
            if outcome.merges_applied or outcome.splits_applied:
                QMessageBox.information(
                    self.main_widget,
                    self.tr("Route validation"),
                    self.tr(
                        f"Merged {outcome.merges_applied} waypoint pair(s); "
                        f"split {outcome.splits_applied} crossing(s).  "
                        "The junction registry has been refreshed."
                    ),
                )
        except Exception as exc:
            QgsMessageLog.logMessage(
                f"Route validation skipped: {exc}", "OMRAT", Qgis.Warning,
            )

        self.ais.update_legs()

        # Refresh junction transition matrices.  When a DB connection is
        # configured we prefer AIS-derived shares; otherwise the handler
        # falls back to its existing geometry/user defaults.
        try:
            handler = getattr(self, 'junctions', None)
            if handler is not None:
                handler.rebuild_from_segments(prefer_user=True)
                if getattr(self.ais, 'db', None) is not None:
                    counts = self.ais.compute_junction_transitions()
                    if counts:
                        handler.apply_ais_counts(counts)
        except Exception as exc:
            QgsMessageLog.logMessage(
                f"Junction transition refresh skipped: {exc}",
                "OMRAT",
                Qgis.Warning,
            )
    
    def remove_route(self)-> None:
        # To implement
        pass
    
    def load_route(self)-> None:
        # To implement
        pass

    # Drift-analysis slots are provided by ``DriftAnalysisMixin``
    # (see ``omrat_utils.drift_analysis_mixin``); they used to live here.

    #--------------------------------------------------------------------------

    def run(self):
        """Run method that loads and starts the plugin"""
        if not self.pluginIsActive:
            self.pluginIsActive = True
            self.main_widget.pbAddRoute.clicked.connect(self.qgis_geoms.add_new_route)
            self.main_widget.pbStopRoute.clicked.connect(self.stop_route)
            self.main_widget.pbRemoveRoute.clicked.connect(self.remove_route)
            self.main_widget.pbLoadRoute.clicked.connect(self.load_route)
            self.main_widget.pbUpdateAIS.clicked.connect(self.update_ais)
            self.main_widget.pbGetGebcoDephts.clicked.connect(self.object.obtain_gebco_data)
            self.main_widget.PBUpdateDepthIntervals.clicked.connect(self.object.update_depth_intervals)
            
            menubar = QMenuBar(self.main_widget)
            menubar.setMinimumSize(320,20)

            fileMenu = menubar.addMenu('File')
            SettingMenu = menubar.addMenu('Settings')
            ConsequenceMenu = menubar.addMenu('Consequence')
            HelpMenu = menubar.addMenu('Help')
            
            # mnuSub1 = viewMenu.addMenu('Sub-menu')
            fileMenu.addAction("Save", self.save_work)
            fileMenu.addAction("Load", self.load_work)
            fileMenu.addSeparator()
            fileMenu.addAction("Export to IWRAP XML", self.export_to_iwrap)
            fileMenu.addAction("Import from IWRAP XML", self.import_from_iwrap)
            SettingMenu.addAction("Drift settings", self.open_drift)
            SettingMenu.addAction("Ship Categories", self.open_ship_categories)
            SettingMenu.addAction("Causation Factors", self.open_causation_factors)
            SettingMenu.addAction("AIS connection settings", self.ais_settings)
            SettingMenu.addAction("Database setup wizard...", self.open_db_setup_wizard)
            SettingMenu.addAction(
                "Junction transition matrix...", self.open_junction_dialog,
            )
            ConsequenceMenu.addAction(
                "Maximum oil onboard...", self.consequence.run_oil_onboard_dialog,
            )
            ConsequenceMenu.addAction(
                "Spill probability per accident...",
                self.consequence.run_spill_probability_dialog,
            )
            ConsequenceMenu.addAction(
                "Spill fraction per accident...",
                self.consequence.run_spill_fraction_dialog,
            )
            ConsequenceMenu.addAction(
                "Catastrophe levels...",
                self.consequence.run_catastrophe_levels_dialog,
            )
            HelpMenu.addAction(
                "Documentation home",
                partial(self._open_docs_url, ""),
            )
            HelpMenu.addAction(
                "Quickstart",
                partial(self._open_docs_url, "quickstart.html"),
            )
            HelpMenu.addAction(
                "User guide",
                partial(self._open_docs_url, "user_guide.html"),
            )
            HelpMenu.addAction(
                "Database setup && AIS ingestion",
                partial(self._open_docs_url, "database_setup.html"),
            )
            HelpMenu.addAction(
                "Theory",
                partial(self._open_docs_url, "theory.html"),
            )
            HelpMenu.addSeparator()
            HelpMenu.addAction(
                "Report a bug (GitHub issues)",
                partial(self._open_external_url, "https://github.com/axelande/OMRAT/issues"),
            )
            #self.main_widget.actionSave_project.clicked.connect(self.save_work)
            #self.main_widget.actionOpen_project.clicked.connect(self.load_work)
            self.main_widget.pbAddSimpleDepth.clicked.connect(self.object.add_simple_depth)
            self.main_widget.pbLoadDepth.clicked.connect(self.object.load_depths)
            self.main_widget.pbRemoveDepth.clicked.connect(self.object.remove_depth)
            self.main_widget.pbAddSimpleObject.clicked.connect(self.object.add_simple_object)
            self.main_widget.pbLoadObject.clicked.connect(self.object.load_objects)
            self.main_widget.pbRemoveObject.clicked.connect(self.object.remove_object)
            # The legacy ``pbView*`` push-buttons used to live in the
            # .ui file as part of the LEP*+View grid; they were
            # removed when the accident-results QTable was added.
            # ``_setup_accident_results_table`` now creates per-row
            # ``QPushButton`` cell widgets inside the table and wires
            # them to ``show_*`` slots, so no ``clicked.connect`` is
            # needed here.
            self.main_widget.pbRunModel.clicked.connect(self.run_calculation)
            # ``LEReportPath`` is the per-run output folder (one .gpkg
            # is written per Run Model into this directory).  The
            # picker button opens a folder-only dialog; live edits to
            # the line-edit re-evaluate the gate on the Run Model
            # button.
            try:
                if hasattr(self.main_widget, 'pbReportPath'):
                    # Replace the legacy "save report to file" wiring.
                    try:
                        self.main_widget.pbReportPath.clicked.disconnect()
                    except Exception:
                        pass
                    self.main_widget.pbReportPath.clicked.connect(
                        self.choose_output_folder,
                    )
                if hasattr(self.main_widget, 'LEReportPath'):
                    self.main_widget.LEReportPath.textChanged.connect(
                        lambda *_: self._update_run_model_enabled(),
                    )
                if hasattr(self.main_widget, 'LEModelName'):
                    self.main_widget.LEModelName.textChanged.connect(
                        lambda *_: self._update_run_model_enabled(),
                    )
                self._restore_output_dir()
                self._update_run_model_enabled()
            except Exception as exc:
                QgsMessageLog.logMessage(
                    f'Could not wire output-folder picker: {exc}',
                    'OMRAT', Qgis.Warning,
                )
            # Connect drift analysis button
            try:
                if hasattr(self.main_widget, 'pbRunDriftAnalysis'):
                    self.main_widget.pbRunDriftAnalysis.clicked.connect(self.run_drift_analysis)
            except Exception:
                pass
            # Wire up the "Previous runs" history table.
            try:
                self._setup_previous_runs_table()
                self.refresh_previous_runs_table()
            except Exception as e:
                QgsMessageLog.logMessage(
                    f'Could not initialise previous-runs table: {e}',
                    'OMRAT', Qgis.Warning,
                )
            # Wire up the accident-results QTable that replaces the
            # legacy LEP* line-edit grid.
            try:
                self._setup_accident_results_table()
            except Exception as e:
                QgsMessageLog.logMessage(
                    f'Could not initialise accident-results table: {e}',
                    'OMRAT', Qgis.Warning,
                )
            # Wire up the Compare tab.
            try:
                self._setup_compare_tab()
            except Exception as e:
                QgsMessageLog.logMessage(
                    f'Could not initialise Compare tab: {e}',
                    'OMRAT', Qgis.Warning,
                )
            # Add a "Manage previous runs" item to the File menu so
            # users can find load / compare / delete actions even if
            # they ignore the table on the Results tab.
            try:
                fileMenu.addSeparator()
                fileMenu.addAction(
                    'Manage previous runs...', self.open_previous_runs_dialog,
                )
            except Exception:
                pass

            self.reset_route_table()
            self.run_traffic_module()
            self.main_widget.show()
