# -*- 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, QUrl
from qgis.PyQt.QtGui import QIcon, QAction, QDesktopServices
from qgis.PyQt.QtWidgets import QMenuBar, QWidget, QFileDialog, QToolBar, QMessageBox
from qgis._core import QgsVectorDataProvider
from qgis.core import (QgsVectorLayer, QgsFeature, QgsLineString, QgsPoint, QgsProject)
from qgis.core import 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 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
from omrat_utils.notifier import MessageBarNotifier


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__)
        # Notifier (message-bar) wired up early so even early-init
        # failures surface through the QGIS message bar rather than a
        # blocking modal.  Tracker URL is read from metadata.txt so the
        # "Let us know" button always points at the live issue tracker.
        self.notifier = MessageBarNotifier(
            iface=self.iface,
            plugin_name='OMRAT',
            tracker_url=self._read_tracker_url(),
            parent=self.iface.mainWindow() if not self.testing else None,
        )
        if not self.testing:
            self._init_locale()

        # 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')

        self.pluginIsActive = 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 _init_locale(self) -> None:
        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)

    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 _remove_drift_corridor_layers(self) -> None:
        """Remove all drift corridor layers from the QGIS project and clear the list."""
        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()

    def _read_tracker_url(self) -> str:
        """Read the `tracker` field from metadata.txt for the Let-us-know button."""
        try:
            import configparser
            cp = configparser.ConfigParser()
            cp.read(os.path.join(self.plugin_dir, 'metadata.txt'), encoding='utf-8')
            return cp.get('general', 'tracker', fallback='')
        except Exception:
            return ''

    def _restore_opentopo_api_key(self) -> None:
        """Populate the OpenTopography API key field from QSettings.

        The key lives in ``omrat/opentopo_api_key`` -- QGIS-profile
        scope, so it survives plugin reloads, QGIS restarts, and
        plugin upgrades.  Each user enters it once via the Depths tab;
        we never commit a personal key into the plugin source.
        """
        le = getattr(self.main_widget, 'LEOpenTopoAPIKey', None)
        if le is None:
            return
        try:
            saved = QSettings().value('omrat/opentopo_api_key', '', type=str)
        except Exception:
            saved = ''
        if saved:
            # Block textChanged while we restore so we don't trigger a
            # spurious re-save of the same value.
            blocker = le.blockSignals(True)
            try:
                le.setText(saved)
            finally:
                le.blockSignals(blocker)
        try:
            le.textChanged.connect(self._on_opentopo_api_key_changed)
        except Exception:
            pass

    def _on_opentopo_api_key_changed(self, text: str) -> None:
        """Persist the API key to QSettings on every edit."""
        try:
            QSettings().setValue('omrat/opentopo_api_key', text or '')
        except Exception:
            pass

    def show_error_popup(self, message: str, function_name: str) -> None:
        """Surface an error in the QGIS message bar.

        Kept as a method (not a free function) so the ~50 legacy call
        sites in mixins/utilities continue to work.  Internally the bar
        notifier handles rendering, Details/Open-logs/Let-us-know
        buttons, and logging -- no more blocking QMessageBox.
        """
        # Stash function_name in the log payload so log readers can still
        # trace where the error came from.
        log_message = f"{function_name}: {message}"
        if getattr(self, 'notifier', None) is not None:
            from omrat_utils.errors import OmratError
            self.notifier.display_exception(
                OmratError(log_message=log_message, user_message=message)
            )
            return
        # Fallback for the (unlikely) pre-notifier window during __init__.
        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_layers(self) -> None:
        self.qgis_geoms.clear()
        self.object.clear()
        self._remove_drift_corridor_layers()
        self._remove_history_layers_from_project()
        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)

    def _reset_data_structures(self) -> None:
        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
        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
        from omrat_utils.causation_factors import CausationFactors
        self.causation_f.data = CausationFactors(self).data
        self.causation_values = self.causation_f.data
        self.distributions.last_id = '1'
        self.distributions.new_id = '1'
        self.traffic.run_update = False
        self.traffic.traffic_data = self.traffic_data
        self.traffic.c_seg = ''

    def _clear_ui_widgets(self) -> None:
        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()
        for w in [
            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,
        ]:
            w.setText('0')
        self.main_widget.sbUniformP1.setValue(0)
        self.main_widget.sbUniformP2.setValue(0)
        for field in [
            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,
        ]:
            field.setText('')
        self.main_widget.LEModelName.setText('')
        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('')

    def _clear_distribution_canvas(self) -> None:
        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

    def clear_model(self) -> None:
        """Reset all model state and UI to defaults."""
        QgsMessageLog.logMessage("Clearing current model", "OMRAT", Qgis.Info)
        self._clear_layers()
        self._reset_data_structures()
        self._clear_ui_widgets()
        self._clear_distribution_canvas()
        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 _remove_calc_result_layers(self) -> None:
        if self.calc is None:
            return
        for attr in (
            'allision_result_layer', 'grounding_result_layer',
            'collision_line_layer', 'collision_point_layer',
            'powered_grounding_layer', 'powered_allision_layer',
        ):
            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)

    def _unload_modules(self) -> None:
        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]

    def unload(self):
        """Removes the plugin menu item and icon from QGIS GUI."""
        try:
            self.notifier.dismiss_all()
        except Exception:
            pass
        try:
            QgsProject.instance().cleared.disconnect(self._on_project_cleared)
        except Exception:
            pass
        self.traffic.unload()
        self.ais.unload()
        self.object.unload()
        self.qgis_geoms.unload()
        self._remove_drift_corridor_layers()
        self._remove_history_layers_from_project()
        self._remove_calc_result_layers()
        for action in self.actions:
            self.iface.removePluginMenu(self.tr(u'&Open Maritime Risk Analysis Tool'), action)
            self.iface.removeToolBarIcon(action)
        if hasattr(self, 'toolbar'):
            del self.toolbar
        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()
        self._unload_modules()
        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 _init_segment_layer(self, seg_data: dict, crs: str):
        route_id = seg_data['Route_Id']
        seg_id = seg_data['Segment_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 = vl.dataProvider()
        if provider is None:
            return vl, fields, None, 0
        provider.addAttributes(fields.toList())
        vl.updateFields()
        try:
            fid = int(str(seg_id))
        except Exception:
            fid = self.segment_id
        return vl, fields, provider, fid

    def _build_segment_feature(self, seg_data: dict, fields, fid: int, crs: str = ''):
        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']}",
        ])
        fet.setId(fid)
        return fet

    def _validate_and_wire_segment(self, vl, fet, fid: int, seg_data: dict) -> bool:
        if not fet.geometry().isGeosValid():
            msg = (
                f"Invalid geometry for segment {seg_data.get('Segment_Id')} "
                f"(route {seg_data.get('Route_Id')}): "
                f"start={seg_data.get('Start_Point')!r} end={seg_data.get('End_Point')!r}. Layer not added."
            )
            print(msg)
            try:
                QgsMessageLog.logMessage(msg, "OMRAT", Qgis.Warning)
            except Exception:
                pass
            try:
                QgsProject.instance().removeMapLayer(vl.id())
            except Exception:
                pass
            return False
        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()
        return True

    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():
            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)
            vl, fields, provider, fid = self._init_segment_layer(seg_data, crs)
            if provider is None:
                return
            fet = self._build_segment_feature(seg_data, fields, fid, crs)
            self.qgis_geoms.style_layer(vl)
            provider.addFeature(fet)
            self.qgis_geoms.label_layer(vl)
            vl.updateExtents()
            if not self._validate_and_wire_segment(vl, fet, fid, seg_data):
                continue
            self.qgis_geoms.vector_layers.append(vl)
            try:
                self.qgis_geoms.leg_dirs[str(seg_data["Segment_Id"])] = list(seg_data.get("Dirs", []))
            except Exception:
                pass
        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
        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:
        # Step 1 (main thread, needs UI dialogs): validate route geometry —
        # snap close waypoints, split crossing legs.
        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,
            )

        # Steps 2+3 (background task): fetch AIS traffic per leg, then
        # rebuild junction transition matrices.  The task disables
        # pbUpdateAIS for its duration and re-enables it in finished().
        self.ais.update_legs()
    
    # 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._connect_toolbar_buttons()
            fileMenu = self._build_menubar()
            self._connect_geometry_buttons()
            self.main_widget.pbRunModel.clicked.connect(self.run_calculation)
            self._wire_report_path_widgets()
            self._wire_drift_button()
            self._init_results_ui(fileMenu)
            self.reset_route_table()
            self.run_traffic_module()
            self.main_widget.show()

    def _connect_toolbar_buttons(self) -> None:
        self.main_widget.pbAddRoute.clicked.connect(self.qgis_geoms.add_new_route)
        self.main_widget.pbStopRoute.clicked.connect(self.stop_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)

    def _build_menubar(self):
        menubar = QMenuBar(self.main_widget)
        menubar.setMinimumSize(320, 20)
        fileMenu = menubar.addMenu('File')
        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)
        sm = menubar.addMenu('Settings')
        sm.addAction("Drift settings", self.open_drift)
        sm.addAction("Ship Categories", self.open_ship_categories)
        sm.addAction("Causation Factors", self.open_causation_factors)
        sm.addAction("AIS connection settings", self.ais_settings)
        sm.addAction("Database setup wizard...", self.open_db_setup_wizard)
        sm.addAction("Junction transition matrix...", self.open_junction_dialog)
        cm = menubar.addMenu('Consequence')
        cm.addAction("Maximum oil onboard...", self.consequence.run_oil_onboard_dialog)
        cm.addAction("Spill probability per accident...", self.consequence.run_spill_probability_dialog)
        cm.addAction("Spill fraction per accident...", self.consequence.run_spill_fraction_dialog)
        cm.addAction("Catastrophe levels...", self.consequence.run_catastrophe_levels_dialog)
        hm = menubar.addMenu('Help')
        hm.addAction("Documentation home", partial(self._open_docs_url, ""))
        hm.addAction("Quickstart", partial(self._open_docs_url, "quickstart.html"))
        hm.addAction("User guide", partial(self._open_docs_url, "user_guide.html"))
        hm.addAction("Database setup && AIS ingestion", partial(self._open_docs_url, "database_setup.html"))
        hm.addAction("Theory", partial(self._open_docs_url, "theory.html"))
        hm.addSeparator()
        hm.addAction("Report a bug (GitHub issues)", partial(self._open_external_url, "https://github.com/axelande/OMRAT/issues"))
        return fileMenu

    def _connect_geometry_buttons(self) -> None:
        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)

    def _wire_report_path_widgets(self) -> None:
        try:
            if hasattr(self.main_widget, 'pbReportPath'):
                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()
            self._restore_opentopo_api_key()
        except Exception as exc:
            QgsMessageLog.logMessage(f'Could not wire output-folder picker: {exc}', 'OMRAT', Qgis.Warning)

    def _wire_drift_button(self) -> None:
        try:
            if hasattr(self.main_widget, 'pbRunDriftAnalysis'):
                self.main_widget.pbRunDriftAnalysis.clicked.connect(self.run_drift_analysis)
        except Exception:
            pass

    def _init_results_ui(self, fileMenu) -> None:
        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)
        try:
            self._setup_accident_results_table()
        except Exception as e:
            QgsMessageLog.logMessage(f'Could not initialise accident-results table: {e}', 'OMRAT', Qgis.Warning)
        try:
            self._setup_compare_tab()
        except Exception as e:
            QgsMessageLog.logMessage(f'Could not initialise Compare tab: {e}', 'OMRAT', Qgis.Warning)
        try:
            fileMenu.addSeparator()
            fileMenu.addAction('Manage previous runs...', self.open_previous_runs_dialog)
        except Exception:
            pass
