plugin.py

# GAIA-alpha Plugin for QGIS
# Copyright (C) 2026 GAIA
#
# 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
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.

import os

from typing import Any, Optional

from qgis.core import (
    Qgis,
    QgsApplication,
    QgsLayerTreeNode,
    QgsMapLayer,
    QgsWkbTypes,
    QgsMessageLog,
)
from qgis.PyQt.QtCore import Qt
from qgis.PyQt.QtGui import QIcon
from qgis.PyQt.QtWidgets import QAction, QMenu

from .ai_dock import AiAssistantDock
from .arcgeek_provider import ARCGEEK_ALGORITHMS, GaiaArcGeekProvider
from .scripts.basemap_manager import BasemapManager
from .scripts.go_to_xy import GoToXYDialog
from .scripts.screen_capture import run_screen_capture
from .geoai_panels import ensure_geoai_plugin_importable


class GaiaAlphaPlugin:
    """Main plugin class — AI Assistant dock + ArcGeek Calculator tools."""

    def __init__(self, iface):
        self.iface = iface
        self.plugin_dir = os.path.dirname(__file__)

        # AI dock
        self.dock: Optional[AiAssistantDock] = None
        self.dock_action: Optional[QAction] = None

        # ArcGeek state
        self.provider: Optional[GaiaArcGeekProvider] = None
        self.arcgeek_actions: list[QAction] = []
        self.context_menu_actions: list[QAction] = []
        self.go_to_xy_dialog: Optional[GoToXYDialog] = None

        # Log collection
        self.log_history: list[dict[str, Any]] = []
        QgsApplication.messageLog().messageReceived.connect(self._on_message_received)

    # ------------------------------------------------------------------
    # Icon helper
    # ------------------------------------------------------------------

    def _arcgeek_icon(self, filename: str) -> QIcon:
        path = os.path.join(self.plugin_dir, "icons", "arcgeek", filename)
        return QIcon(path) if os.path.isfile(path) else QIcon()

    # ------------------------------------------------------------------
    # initGui
    # ------------------------------------------------------------------

    def initGui(self):
        # ---- AI Assistant dock toggle ----
        icon_path = os.path.join(self.plugin_dir, "icons", "icon.svg")
        icon = QIcon(icon_path) if os.path.isfile(icon_path) else QIcon()

        self.dock_action = QAction(icon, self.tr("GAIA-alpha"), self.iface.mainWindow())
        self.dock_action.setObjectName("GaiaAlphaAction")
        self.dock_action.setCheckable(True)
        self.dock_action.setToolTip(self.tr("Show or hide the AI assistant dock"))
        self.dock_action.triggered.connect(self._toggle_dock)
        self.iface.addPluginToMenu(self.tr("&GAIA"), self.dock_action)

        self.dock = AiAssistantDock(self.iface, self.iface.mainWindow())
        self.dock.setAllowedAreas(
            Qt.DockWidgetArea.LeftDockWidgetArea | Qt.DockWidgetArea.RightDockWidgetArea
        )
        self.iface.addDockWidget(Qt.DockWidgetArea.RightDockWidgetArea, self.dock)
        self.dock.visibilityChanged.connect(self._dock_visibility_changed)
        self.dock.hide()

        # ---- ArcGeek Processing Provider ----
        self.provider = GaiaArcGeekProvider()
        QgsApplication.processingRegistry().addProvider(self.provider)

        # ---- ArcGeek menu items (submenu under &GAIA) ----
        self._build_arcgeek_menu()

        # ---- Context menus ----
        version = Qgis.QGIS_VERSION_INT
        try:
            self.iface.layerTreeView().contextMenuAboutToShow.disconnect(
                self._add_layer_menu_items
            )
        except Exception:
            pass

        if version >= 30000:
            self.iface.layerTreeView().contextMenuAboutToShow.connect(
                self._add_layer_menu_items
            )
        else:
            self.iface.layerTreeView().layerTreeContextMenuAboutToShow.connect(
                self._add_layer_menu_items
            )
        self.iface.mapCanvas().contextMenuAboutToShow.connect(self._add_map_menu_items)

    # ------------------------------------------------------------------
    # ArcGeek submenu builder
    # ------------------------------------------------------------------

    def _build_arcgeek_menu(self):
        menu_label = self.tr("&GAIA")

        # Helper to create & register an action
        def _add(label: str, callback, icon_file: str):
            action = QAction(
                self._arcgeek_icon(icon_file), label, self.iface.mainWindow()
            )
            action.triggered.connect(callback)
            self.iface.addPluginToMenu(menu_label, action)
            self.arcgeek_actions.append(action)

        def _sep():
            sep = QAction(self.iface.mainWindow())
            sep.setSeparator(True)
            self.iface.addPluginToMenu(menu_label, sep)
            self.arcgeek_actions.append(sep)

        # --- Geometric & Coordinate Tools ---
        _add("Calculate Point Coordinates", self._run("coordinate"), "calculate_xy.png")
        _add("Calculate Line Geometry", self._run("line"), "calculate_length.png")
        _add("Calculate Polygon Geometry", self._run("polygon"), "calculate_area.png")
        _add("Calculate Angles", self._run("angles"), "calculate_angles.png")
        _add(
            "Extract Ordered Points from Polygons",
            self._run("polygon_to_points"),
            "order_point.png",
        )
        _add(
            "Lines to Ordered Points",
            self._run("lines_to_ordered_points"),
            "lines_to_points.png",
        )
        _add(
            "Azimuth and Distance from Coordinates and Table",
            self._run("calculate_line"),
            "calculate_line.png",
        )
        _add(
            "Export to CSV (Excel compatible)",
            self._run("export_to_csv"),
            "export_csv.png",
        )
        _add(
            "Tree Planting Pattern Generator",
            self._run("tree_planting_pattern"),
            "tree_planting.png",
        )
        _sep()

        # --- Hydrological Analysis ---
        _add(
            "Stream Network with Order",
            self._run("watershed_stream"),
            "watershed_network.png",
        )
        _add(
            "Watershed Basin Delineation",
            self._run("watershed_basin"),
            "watershed_basin.png",
        )
        _add(
            "Watershed Morphometric Analysis",
            self._run("basin_analysis"),
            "watershed_morfo.png",
        )
        _add(
            "Multiple Basin Delineation by Points",
            self._run("multiple_basin_delineation"),
            "watershed_basin.png",
        )
        _add("Global Curve Number", self._run("global_cn"), "global_cn.png")
        _sep()

        # --- Land Analysis & Raster ---
        _add(
            "Land Use Change Detection",
            self._run("land_use_change"),
            "land_use_change.png",
        )
        _add("Weighted Sum", self._run("weighted_sum"), "weighted_sum.png")
        _add(
            "Least Cost Path Finder",
            self._run("least_cost_path"),
            "least_cost_path.png",
        )
        _add("Dam Flood Simulation", self._run("dam_flood_simulation"), "dam_flood.png")
        _add("Kriging Analysis", self._run("kriging_analysis"), "kriging.png")
        _sep()

        # --- Terrain & CAD ---
        _add(
            "Export Contours to 3D CAD",
            self._run("contour_export"),
            "contour_export3DCAD.png",
        )
        _add(
            "Optimized Parcel Division",
            self._run("optimized_parcel_division"),
            "parcel_division.png",
        )
        _sep()

        # --- Visualization & Utilities ---
        _add(
            "Topographic Profile", self._run("topographic_profile"), "topo_profile.png"
        )
        _add(
            "Manage Basemaps (Google, Bing, Esri)",
            self._run_basemap_manager,
            "basemap.png",
        )
        _add("Screen Capture", self._run_screen_capture, "screen_capture.png")
        _add(
            "Satellite Index Calculator",
            self._run("satellite_index"),
            "satellite_index.png",
        )
        _add(
            "Enhanced Image Classification",
            self._run("enhanced_classification"),
            "classification.png",
        )
        _sep()

        # --- Navigation ---
        _add("Go to XY", self._run_go_to_xy, "gotoXY.png")
        _sep()

    # ------------------------------------------------------------------
    # Algorithm launcher (opens the Processing dialog)
    # ------------------------------------------------------------------

    def _run(self, algorithm_key: str):
        """Return a callback that opens the Processing dialog for the given algorithm."""

        def callback():
            from qgis import processing

            alg_class = ARCGEEK_ALGORITHMS[algorithm_key]["class"]
            processing.execAlgorithmDialog(alg_class())

        return callback

    # ------------------------------------------------------------------
    # Standalone launchers
    # ------------------------------------------------------------------

    def _run_basemap_manager(self):
        dialog = BasemapManager(self.iface)
        try:
            dialog.exec()
        except AttributeError:
            dialog.exec_()

    def _run_screen_capture(self):
        run_screen_capture(self.iface)

    def _run_go_to_xy(self):
        if self.go_to_xy_dialog is None:
            self.go_to_xy_dialog = GoToXYDialog(self.iface, self.iface.mainWindow())
        self.go_to_xy_dialog.show()

    # ------------------------------------------------------------------
    # AI dock helpers
    # ------------------------------------------------------------------

    def _dock_visibility_changed(self, visible: bool):
        if self.dock_action is not None:
            self.dock_action.blockSignals(True)
            self.dock_action.setChecked(visible)
            self.dock_action.blockSignals(False)

    def _toggle_dock(self, checked: bool):
        if self.dock is None:
            return
        if checked:
            self.dock.show()
            self.dock.raise_()
        else:
            self.dock.hide()

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


    def _toggle_geoai_dock(self, dock_key: str):
        """Open or focus a GeoAI dock widget, loading dependencies if needed."""
        # 1. Check dependencies
        if not ensure_dependencies(self.iface, dock_key):
            # Dependency installer was triggered instead of opening the dock
            return

        try:
            # 2. Get or create the dock
            if dock_key not in self._geoai_docks:
                ensure_geoai_plugin_importable()
                # Instantiate the specific dock widget based on the key
                if dock_key == "moondream":
                    from geoai.dialogs.moondream import MoondreamDockWidget
                    self._geoai_docks[dock_key] = MoondreamDockWidget(self.iface, self.iface.mainWindow())
                elif dock_key == "segmentation":
                    from geoai.dialogs.segmentation import SegmentationDockWidget
                    self._geoai_docks[dock_key] = SegmentationDockWidget(self.iface, self.iface.mainWindow())
                elif dock_key == "instance_segmentation":
                    from geoai.dialogs.instance_segmentation import InstanceSegmentationDockWidget
                    self._geoai_docks[dock_key] = InstanceSegmentationDockWidget(self.iface, self.iface.mainWindow())
                elif dock_key == "samgeo":
                    from geoai.dialogs.samgeo import SamgeoDockWidget
                    self._geoai_docks[dock_key] = SamgeoDockWidget(self.iface, self.iface.mainWindow())
                elif dock_key == "deepforest":
                    from geoai.dialogs.deepforest_panel import DeepForestDockWidget
                    self._geoai_docks[dock_key] = DeepForestDockWidget(self.iface, self.iface.mainWindow())
                elif dock_key == "water_segmentation":
                    from geoai.dialogs.water_segmentation import WaterSegmentationDockWidget
                    self._geoai_docks[dock_key] = WaterSegmentationDockWidget(self.iface, self.iface.mainWindow())
                else:
                    return

                self.iface.addDockWidget(Qt.DockWidgetArea.RightDockWidgetArea, self._geoai_docks[dock_key])
                self._geoai_docks[dock_key].hide()

            dock = self._geoai_docks[dock_key]

            # 3. Toggle visibility
            if dock.isVisible():
                dock.hide()
            else:
                dock.show()
                dock.raise_()
        except Exception as e:
            from qgis.core import QgsMessageLog, Qgis
            QgsMessageLog.logMessage(f"Failed to toggle {dock_key} panel: {str(e)}", "GAIA", Qgis.MessageLevel.Critical)
            if self.iface:
                self.iface.messageBar().pushMessage("Error", f"Failed to open {dock_key} panel: {str(e)}", level=Qgis.MessageLevel.Critical, duration=5)

    def _clear_geoai_gpu_memory(self):
        """Clear GPU memory gracefully by freeing unused cache and objects."""
        from qgis.PyQt.QtWidgets import QMessageBox

        cleared_items = []
        import gc

        ensure_geoai_plugin_importable()
        try:
            import torch
        except ImportError:
            torch = None
        
        # We only really need to clean up SAM and DeepForest models which might be lingering in dock attributes
        sam_dock = self._geoai_docks.get("samgeo")
        df_dock = self._geoai_docks.get("deepforest")

        if sam_dock and hasattr(sam_dock, "sam"):
            sam_dock.sam = None
            cleared_items.append("Segment Anything model")
        if df_dock and hasattr(df_dock, "deepforest"):
            df_dock.deepforest = None
            cleared_items.append("Tree Segmentation model")

        # Force garbage collection
        for _ in range(5):
            gc.collect()

        memory_info = ""
        if torch is not None and torch.cuda.is_available():
            try:
                torch.cuda.synchronize()
                torch.cuda.empty_cache()
                if hasattr(torch.cuda, "ipc_collect"):
                    torch.cuda.ipc_collect()
                gc.collect()
                torch.cuda.empty_cache()
                torch.cuda.synchronize()
                cleared_items.append("CUDA cache")

                allocated = torch.cuda.memory_allocated() / 1024**2
                reserved = torch.cuda.memory_reserved() / 1024**2
                memory_info = f"\n\nGPU Memory:\n  Allocated: {allocated:.1f} MB\n  Reserved: {reserved:.1f} MB"
            except Exception as e:
                memory_info = f"\n\nError clearing CUDA: {str(e)}"
        elif torch is not None and hasattr(torch, "backends") and hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
            try:
                if hasattr(torch, "mps") and hasattr(torch.mps, "synchronize"):
                    torch.mps.synchronize()
                if hasattr(torch, "mps") and hasattr(torch.mps, "empty_cache"):
                    torch.mps.empty_cache()
                gc.collect()
                if hasattr(torch, "mps") and hasattr(torch.mps, "synchronize"):
                    torch.mps.synchronize()
                cleared_items.append("MPS cache")
                memory_info = "\n\nApple Metal (MPS) cache cleared."
            except Exception as e:
                memory_info = f"\n\nError clearing MPS cache: {str(e)}"
        elif torch is None:
            memory_info = "\n\nPyTorch not loaded. Cache is already clear."
        else:
            memory_info = "\n\nNo CUDA or MPS accelerator available."

        if cleared_items:
            message = f"Cleared: {', '.join(cleared_items)}{memory_info}"
        else:
            message = f"No active models found to clear.{memory_info}"

        self.iface.statusBarIface().showMessage("Accelerator memory cleared", 3000)
        QMessageBox.information(
            self.iface.mainWindow(),
            "Clear GPU Memory",
            message,
        )

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

    # Context menus
    # ------------------------------------------------------------------

    def _add_layer_menu_items(self, menu):
        # Clean up previous actions
        for action in self.context_menu_actions[:]:
            try:
                if action in menu.actions():
                    menu.removeAction(action)
                self.context_menu_actions.remove(action)
            except RuntimeError:
                self.context_menu_actions.remove(action)
            except Exception:
                pass

        current_node = self.iface.layerTreeView().currentNode()
        if not (
            isinstance(current_node, QgsLayerTreeNode)
            and current_node.nodeType() == QgsLayerTreeNode.NodeLayer
        ):
            return

        layer = self.iface.layerTreeView().currentLayer()
        if layer is None or layer.type() != QgsMapLayer.VectorLayer:
            return

        geom_type = layer.geometryType()

        def _ctx_action(label, alg_key, icon_file):
            action = QAction(self._arcgeek_icon(icon_file), label, menu)
            action.triggered.connect(lambda: self._run(alg_key)())
            self._insert_context_action(menu, action)

        if geom_type == QgsWkbTypes.PointGeometry:
            _ctx_action("Calculate XY Coordinates", "coordinate", "calculate_xy.png")
        elif geom_type == QgsWkbTypes.LineGeometry:
            _ctx_action("Calculate Length", "line", "calculate_length.png")
        elif geom_type == QgsWkbTypes.PolygonGeometry:
            _ctx_action("Calculate Area and Perimeter", "polygon", "calculate_area.png")

    def _insert_context_action(self, menu, action):
        insert_pos = 0
        for i, existing in enumerate(menu.actions()):
            if existing.isSeparator() or existing.menu():
                insert_pos = i
                break
        actions = menu.actions()
        menu.insertAction(
            actions[insert_pos] if insert_pos < len(actions) else None, action
        )
        self.context_menu_actions.append(action)

    def _add_map_menu_items(self, menu):
        action = QAction(self._arcgeek_icon("gotoXY.png"), "Go to XY", menu)
        action.triggered.connect(self._run_go_to_xy)
        menu.addAction(action)

    # ------------------------------------------------------------------
    # tr helper (for plugins without QObject)
    # ------------------------------------------------------------------

    @staticmethod
    def tr(msg: str) -> str:
        from qgis.PyQt.QtCore import QCoreApplication

        return QCoreApplication.translate("GaiaAlphaPlugin", msg)

    # ------------------------------------------------------------------
    # Log listener
    # ------------------------------------------------------------------

    def _on_message_received(self, message, tag, level):
        from datetime import datetime
        
        # Keep only the last 500 messages to avoid memory issues
        if len(self.log_history) > 500:
            self.log_history.pop(0)
            
        self.log_history.append({
            "timestamp": datetime.now().isoformat(),
            "message": message,
            "tag": tag,
            "level": level
        })

    def get_log_history(self) -> list[dict[str, Any]]:
        return self.log_history

    # ------------------------------------------------------------------
    # Unload
    # ------------------------------------------------------------------

    def unload(self):
        # AI dock
        if self.dock is not None:
            self.dock.visibilityChanged.disconnect(self._dock_visibility_changed)
            self.iface.removeDockWidget(self.dock)
            self.dock.deleteLater()
            self.dock = None

        if self.dock_action is not None:
            self.iface.removePluginMenu(self.tr("&GAIA"), self.dock_action)
            self.dock_action.deleteLater()
            self.dock_action = None

        # ArcGeek actions
        menu_label = self.tr("&GAIA")
        for action in self.arcgeek_actions:
            self.iface.removePluginMenu(menu_label, action)
        self.arcgeek_actions.clear()

        # Processing provider
        # The underlying C++ object may already have been destroyed (e.g. during
        # a plugin reload or when QGIS is shutting down), in which case
        # removeProvider() raises RuntimeError. Guard against that so the rest
        # of unload() can still run cleanly.
        if self.provider is not None:
            try:
                QgsApplication.processingRegistry().removeProvider(self.provider)
            except RuntimeError:
                pass
            self.provider = None

        # GoToXY dialog
        if self.go_to_xy_dialog:
            self.go_to_xy_dialog.close()
            self.go_to_xy_dialog = None

        # Context menus
        try:
            self.iface.layerTreeView().contextMenuAboutToShow.disconnect(
                self._add_layer_menu_items
            )
        except Exception:
            pass
        try:
            self.iface.mapCanvas().contextMenuAboutToShow.disconnect(
                self._add_map_menu_items
            )
        except Exception:
            pass