# -*- coding: utf-8 -*-
import math
import os
import tempfile
import time
from pathlib import Path

import numpy as np
from osgeo import gdal, ogr
from qgis.PyQt.QtWidgets import (
    QDialog, QVBoxLayout, QHBoxLayout, QLabel, QLineEdit, QPushButton,
    QFileDialog, QCheckBox, QSpinBox, QDoubleSpinBox, QProgressBar,
    QTextEdit, QGroupBox, QMessageBox, QTabWidget, QWidget, QComboBox
)
from qgis.PyQt.QtCore import QThread, pyqtSignal
from qgis.core import QgsProject, QgsMapLayerProxyModel, QgsRasterLayer
from qgis.gui import QgsMapLayerComboBox


class ViewshedWorker(QThread):
    """Worker thread for viewshed calculation."""
    progress = pyqtSignal(int, str)
    finished = pyqtSignal(bool, str)

    def __init__(self, params):
        super().__init__()
        self.params = params
        self.is_running = True

    def stop(self):
        """Request the worker loop to stop at the next safe checkpoint."""
        self.is_running = False

    def run(self):
        """Run the viewshed calculation."""
        try:
            result = self._calculate_viewshed()
            self.finished.emit(True, result)
        except Exception as e:
            self.finished.emit(False, f"Error: {str(e)}")

    def _calculate_viewshed(self):
        """Core viewshed calculation logic."""
        start_time = time.time()
        gdal.UseExceptions()

        dsm_path = self.params['dsm_path']
        viewpoints_path = self.params.get('viewpoints_path')
        viewpoints_layer_name = self.params.get('viewpoints_layer_name')
        generated_points = self.params.get('generated_points') or []
        output_path = self.params['output_path']
        max_distance = self.params['max_distance']
        observer_height = self.params['observer_height']
        raster_resolution = self.params['raster_resolution']
        chm_path = self.params.get('chm_path')
        use_chm = self.params.get('use_chm', False)

        # 1) Initialize datasets
        self.progress.emit(5, "Loading DSM...")
        dsm_ds = gdal.Open(dsm_path, gdal.GA_ReadOnly)
        if dsm_ds is None:
            raise Exception(f"Cannot open DSM: {dsm_path}")

        surface_ds = dsm_ds
        surface_band = dsm_ds.GetRasterBand(1)

        chm_band = None
        chm_inv_gt = None
        chm_nodata = None
        if use_chm and chm_path:
            self.progress.emit(10, "Loading CHM...")
            chm_ds = gdal.Open(chm_path, gdal.GA_ReadOnly)
            if chm_ds is None:
                raise Exception(f"Cannot open CHM: {chm_path}")
            chm_band = chm_ds.GetRasterBand(1)
            chm_inv_gt = gdal.InvGeoTransform(chm_ds.GetGeoTransform())
            chm_nodata = chm_band.GetNoDataValue()

        gt = surface_ds.GetGeoTransform()
        pixel_area_m2 = abs(gt[1] * gt[5])
        curvature_coeff = 0.0  # Earth curvature disabled by default

        # 2) Load viewpoints
        self.progress.emit(15, "Loading viewpoints...")
        if generated_points:
            viewpoint_coordinates = list(generated_points)
        else:
            vp_ds = ogr.Open(viewpoints_path)
            if vp_ds is None:
                raise Exception(f"Cannot open viewpoints: {viewpoints_path}")

            if viewpoints_layer_name:
                vp_layer = vp_ds.GetLayerByName(viewpoints_layer_name)
                if vp_layer is None:
                    vp_layer = vp_ds.GetLayer(0)
            else:
                vp_layer = vp_ds.GetLayer(0)

            viewpoint_coordinates = []
            vp_layer.ResetReading()
            for feature in vp_layer:
                geom = feature.GetGeometryRef()
                if geom is None:
                    continue
                viewpoint_coordinates.append((geom.GetX(), geom.GetY()))

        total_features = len(viewpoint_coordinates)

        if total_features == 0:
            raise Exception("No viewpoints found")

        points = []
        min_x = None
        max_x = None
        min_y = None
        max_y = None
        processed_count = 0
        next_progress_tenth = 1

        # 3) Process each viewpoint
        self.progress.emit(20, f"Processing {total_features} viewpoints...")
        for observer_x, observer_y in viewpoint_coordinates:
            if not self.is_running:
                break

            # CHM check (optional)
            chm_value = None
            if use_chm and chm_band is not None and chm_inv_gt is not None:
                px, py = gdal.ApplyGeoTransform(chm_inv_gt, observer_x, observer_y)
                col = int(math.floor(px))
                row = int(math.floor(py))
                if 0 <= col < chm_band.XSize and 0 <= row < chm_band.YSize:
                    arr = chm_band.ReadAsArray(col, row, 1, 1)
                    if arr is not None:
                        value = float(arr[0, 0])
                        if np.isfinite(value):
                            if chm_nodata is None or not np.isclose(value, chm_nodata):
                                chm_value = value

            # Calculate viewshed
            if chm_value is not None and chm_value > 0:
                visible_area_m2 = 0.0
            else:
                try:
                    viewshed_ds = gdal.ViewshedGenerate(
                        surface_band,
                        "MEM",
                        "",
                        [],
                        observer_x,
                        observer_y,
                        observer_height,
                        0,  # TARGET_HEIGHT_M
                        1,
                        0,
                        0,
                        0,
                        curvature_coeff,
                        gdal.GVM_Edge,
                        max_distance,
                    )

                    if viewshed_ds is None:
                        visible_area_m2 = 0.0
                    else:
                        visible_cells = np.count_nonzero(
                            viewshed_ds.GetRasterBand(1).ReadAsArray() == 1
                        )
                        visible_area_m2 = visible_cells * pixel_area_m2
                except Exception:
                    visible_area_m2 = 0.0

            points.append((observer_x, observer_y, visible_area_m2 / 10000.0))

            # Update extent
            if min_x is None or observer_x < min_x:
                min_x = observer_x
            if max_x is None or observer_x > max_x:
                max_x = observer_x
            if min_y is None or observer_y < min_y:
                min_y = observer_y
            if max_y is None or observer_y > max_y:
                max_y = observer_y

            processed_count += 1

            # Progress reporting
            if total_features > 0:
                progress_percent = 20 + int((processed_count / total_features) * 70)
                while (next_progress_tenth <= 10 and processed_count >= math.ceil((total_features * next_progress_tenth) / 10)):
                    self.progress.emit(
                        progress_percent,
                        f"Processed {processed_count}/{total_features} viewpoints"
                    )
                    next_progress_tenth += 1

        if not self.is_running:
            return "Cancelled by user"

        # 4) Aggregate and write output raster
        self.progress.emit(92, "Aggregating results and writing raster...")

        nodata_value = -9999.0
        res = float(raster_resolution)
        max_col = int(math.floor(((max_x - min_x) / res) + 0.5))
        max_row = int(math.floor(((max_y - min_y) / res) + 0.5))
        x_size = max_col + 1
        y_size = max_row + 1

        origin_x = min_x - (res / 2.0)
        origin_y = max_y + (res / 2.0)
        out_gt = (origin_x, res, 0.0, origin_y, 0.0, -res)

        out_dir = Path(output_path)
        out_dir.parent.mkdir(parents=True, exist_ok=True)

        out_arr = np.full((y_size, x_size), nodata_value, dtype=np.float32)

        cell_max = {}
        for x, y, value in points:
            col = int(math.floor(((x - min_x) / res) + 0.5))
            row = int(math.floor(((max_y - y) / res) + 0.5))
            if 0 <= col < x_size and 0 <= row < y_size:
                key = (col, row)
                prev = cell_max.get(key)
                if prev is None or value > prev:
                    cell_max[key] = value

        for (col, row), value in cell_max.items():
            out_arr[row, col] = float(value)

        # Write GeoTIFF
        driver = gdal.GetDriverByName("GTiff")
        out_raster_ds = driver.Create(output_path, x_size, y_size, 1, gdal.GDT_Float32)
        out_raster_ds.SetGeoTransform(out_gt)
        out_raster_ds.SetProjection(surface_ds.GetProjection())
        out_band = out_raster_ds.GetRasterBand(1)
        out_band.SetNoDataValue(nodata_value)
        out_band.WriteArray(out_arr)
        out_band.FlushCache()
        out_raster_ds.FlushCache()
        out_raster_ds = None

        elapsed_time = time.time() - start_time
        self.progress.emit(100, "Complete!")

        return (
            f"Successfully processed {processed_count} viewpoints\n"
            f"Output: {output_path}\n"
            f"Runtime: {elapsed_time:.2f} seconds ({elapsed_time / 3600:.2f} hours)"
        )


class ViewshedCalculatorDialog(QDialog):
    """Dialog for Viewshed Calculator plugin."""

    def __init__(self, iface):
        super().__init__()
        self.iface = iface
        self.worker = None
        self.init_ui()

    def init_ui(self):
        """Initialize the UI."""
        self.setWindowTitle('Viewshed Calculator')
        self.setMinimumWidth(600)
        self.setMinimumHeight(700)

        main_layout = QVBoxLayout()

        # DSM section
        dsm_group = QGroupBox("Digital Surface Model (DSM)")
        dsm_layout = QVBoxLayout()
        dsm_file_layout = QHBoxLayout()
        self.dsm_line = QLineEdit()
        self.dsm_line.setPlaceholderText("Path to DSM raster file")
        dsm_browse = QPushButton("Browse...")
        dsm_browse.clicked.connect(self.browse_dsm)
        dsm_file_layout.addWidget(QLabel("DSM file:"))
        dsm_file_layout.addWidget(self.dsm_line)
        dsm_file_layout.addWidget(dsm_browse)
        dsm_layout.addLayout(dsm_file_layout)

        dsm_project_layout = QHBoxLayout()
        self.dsm_layer_combo = QgsMapLayerComboBox()
        self.dsm_layer_combo.setFilters(QgsMapLayerProxyModel.RasterLayer)
        self.dsm_layer_combo.setAllowEmptyLayer(True)
        dsm_project_layout.addWidget(QLabel("Or DSM layer in project:"))
        dsm_project_layout.addWidget(self.dsm_layer_combo)
        dsm_layout.addLayout(dsm_project_layout)
        dsm_group.setLayout(dsm_layout)
        main_layout.addWidget(dsm_group)

        # CHM section
        chm_group = QGroupBox("Canopy Height Model (CHM) - Optional")
        chm_layout = QVBoxLayout()
        self.use_chm_check = QCheckBox("Use CHM for validation")
        self.use_chm_check.toggled.connect(self.toggle_chm)

        chm_file_layout = QHBoxLayout()
        self.chm_line = QLineEdit()
        self.chm_line.setPlaceholderText("Path to CHM raster file (optional)")
        self.chm_line.setEnabled(False)
        chm_browse = QPushButton("Browse...")
        chm_browse.clicked.connect(self.browse_chm)
        chm_browse.setEnabled(False)
        self.chm_browse_btn = chm_browse
        chm_file_layout.addWidget(QLabel("CHM:"))
        chm_file_layout.addWidget(self.chm_line)
        chm_file_layout.addWidget(chm_browse)

        chm_project_layout = QHBoxLayout()
        self.chm_layer_combo = QgsMapLayerComboBox()
        self.chm_layer_combo.setFilters(QgsMapLayerProxyModel.RasterLayer)
        self.chm_layer_combo.setAllowEmptyLayer(True)
        self.chm_layer_combo.setEnabled(False)
        chm_project_layout.addWidget(QLabel("Or CHM layer in project:"))
        chm_project_layout.addWidget(self.chm_layer_combo)

        chm_layout.addWidget(self.use_chm_check)
        chm_layout.addLayout(chm_file_layout)
        chm_layout.addLayout(chm_project_layout)
        chm_group.setLayout(chm_layout)
        main_layout.addWidget(chm_group)

        # Viewpoints section
        vp_group = QGroupBox("Viewpoints")
        vp_layout = QVBoxLayout()

        self.viewpoints_tabs = QTabWidget()

        # Existing viewpoints tab
        vp_existing_tab = QWidget()
        vp_existing_layout = QVBoxLayout()

        vp_file_layout = QHBoxLayout()
        self.vp_line = QLineEdit()
        self.vp_line.setPlaceholderText("Path to viewpoints shapefile")
        vp_browse = QPushButton("Browse...")
        vp_browse.clicked.connect(self.browse_viewpoints)
        vp_file_layout.addWidget(QLabel("Shapefile:"))
        vp_file_layout.addWidget(self.vp_line)
        vp_file_layout.addWidget(vp_browse)
        vp_existing_layout.addLayout(vp_file_layout)

        vp_project_layout = QHBoxLayout()
        self.vp_layer_combo = QgsMapLayerComboBox()
        self.vp_layer_combo.setFilters(QgsMapLayerProxyModel.PointLayer)
        self.vp_layer_combo.setAllowEmptyLayer(True)
        vp_project_layout.addWidget(QLabel("Or point layer in project:"))
        vp_project_layout.addWidget(self.vp_layer_combo)
        vp_existing_layout.addLayout(vp_project_layout)

        vp_res_layout = QHBoxLayout()
        vp_res_layout.addWidget(QLabel("Raster Resolution (m):"))
        self.vp_existing_resolution_spin = QDoubleSpinBox()
        self.vp_existing_resolution_spin.setValue(25.0)
        self.vp_existing_resolution_spin.setMinimum(1.0)
        self.vp_existing_resolution_spin.setMaximum(1000.0)
        vp_res_layout.addWidget(self.vp_existing_resolution_spin)
        vp_res_layout.addStretch()
        vp_existing_layout.addLayout(vp_res_layout)

        vp_existing_tab.setLayout(vp_existing_layout)

        # Generated viewpoints tab
        vp_generate_tab = QWidget()
        vp_generate_layout = QVBoxLayout()

        gen_mode_layout = QHBoxLayout()
        gen_mode_layout.addWidget(QLabel("Generate from:"))
        self.vp_generate_source_combo = QComboBox()
        self.vp_generate_source_combo.addItems(["Extent", "Polygon shapefile/layer"])
        self.vp_generate_source_combo.currentIndexChanged.connect(self.update_generation_source_ui)
        gen_mode_layout.addWidget(self.vp_generate_source_combo)
        gen_mode_layout.addStretch()
        vp_generate_layout.addLayout(gen_mode_layout)

        spacing_layout = QHBoxLayout()
        spacing_layout.addWidget(QLabel("Viewpoint spacing (map units):"))
        self.vp_spacing_spin = QDoubleSpinBox()
        self.vp_spacing_spin.setMinimum(0.1)
        self.vp_spacing_spin.setMaximum(1000000.0)
        self.vp_spacing_spin.setValue(25.0)
        self.vp_spacing_spin.setDecimals(3)
        spacing_layout.addWidget(self.vp_spacing_spin)
        spacing_layout.addStretch()
        vp_generate_layout.addLayout(spacing_layout)

        self.vp_extent_widget = QWidget()
        extent_layout = QVBoxLayout()

        extent_x_layout = QHBoxLayout()
        extent_x_layout.addWidget(QLabel("X min:"))
        self.vp_xmin_spin = QDoubleSpinBox()
        self.vp_xmin_spin.setRange(-1e12, 1e12)
        self.vp_xmin_spin.setDecimals(6)
        extent_x_layout.addWidget(self.vp_xmin_spin)
        extent_x_layout.addWidget(QLabel("X max:"))
        self.vp_xmax_spin = QDoubleSpinBox()
        self.vp_xmax_spin.setRange(-1e12, 1e12)
        self.vp_xmax_spin.setDecimals(6)
        extent_x_layout.addWidget(self.vp_xmax_spin)
        extent_layout.addLayout(extent_x_layout)

        extent_y_layout = QHBoxLayout()
        extent_y_layout.addWidget(QLabel("Y min:"))
        self.vp_ymin_spin = QDoubleSpinBox()
        self.vp_ymin_spin.setRange(-1e12, 1e12)
        self.vp_ymin_spin.setDecimals(6)
        extent_y_layout.addWidget(self.vp_ymin_spin)
        extent_y_layout.addWidget(QLabel("Y max:"))
        self.vp_ymax_spin = QDoubleSpinBox()
        self.vp_ymax_spin.setRange(-1e12, 1e12)
        self.vp_ymax_spin.setDecimals(6)
        extent_y_layout.addWidget(self.vp_ymax_spin)
        extent_layout.addLayout(extent_y_layout)

        canvas_extent_layout = QHBoxLayout()
        use_canvas_extent_btn = QPushButton("Use Current Canvas Extent")
        use_canvas_extent_btn.clicked.connect(self.set_extent_from_canvas)
        canvas_extent_layout.addWidget(use_canvas_extent_btn)
        canvas_extent_layout.addStretch()
        extent_layout.addLayout(canvas_extent_layout)

        self.vp_extent_widget.setLayout(extent_layout)
        vp_generate_layout.addWidget(self.vp_extent_widget)

        self.vp_polygon_widget = QWidget()
        polygon_layout = QVBoxLayout()

        polygon_file_layout = QHBoxLayout()
        self.vp_polygon_line = QLineEdit()
        self.vp_polygon_line.setPlaceholderText("Path to polygon shapefile")
        vp_polygon_browse = QPushButton("Browse...")
        vp_polygon_browse.clicked.connect(self.browse_viewpoints_polygon)
        polygon_file_layout.addWidget(QLabel("Polygon file:"))
        polygon_file_layout.addWidget(self.vp_polygon_line)
        polygon_file_layout.addWidget(vp_polygon_browse)
        polygon_layout.addLayout(polygon_file_layout)

        polygon_project_layout = QHBoxLayout()
        self.vp_polygon_layer_combo = QgsMapLayerComboBox()
        self.vp_polygon_layer_combo.setFilters(QgsMapLayerProxyModel.PolygonLayer)
        self.vp_polygon_layer_combo.setAllowEmptyLayer(True)
        polygon_project_layout.addWidget(QLabel("Or polygon layer in project:"))
        polygon_project_layout.addWidget(self.vp_polygon_layer_combo)
        polygon_layout.addLayout(polygon_project_layout)

        self.vp_polygon_widget.setLayout(polygon_layout)
        vp_generate_layout.addWidget(self.vp_polygon_widget)

        vp_generate_tab.setLayout(vp_generate_layout)

        self.viewpoints_tabs.addTab(vp_existing_tab, "Existing Layer")
        self.viewpoints_tabs.addTab(vp_generate_tab, "Generate Grid")
        vp_layout.addWidget(self.viewpoints_tabs)
        self.update_generation_source_ui()

        vp_group.setLayout(vp_layout)
        main_layout.addWidget(vp_group)

        # Parameters section
        params_group = QGroupBox("Parameters")
        params_layout = QVBoxLayout()

        # Observer height
        obs_h_layout = QHBoxLayout()
        obs_h_layout.addWidget(QLabel("Observer Height (m):"))
        self.observer_height_spin = QDoubleSpinBox()
        self.observer_height_spin.setValue(1.6)
        self.observer_height_spin.setMinimum(0.0)
        self.observer_height_spin.setMaximum(100.0)
        obs_h_layout.addWidget(self.observer_height_spin)
        obs_h_layout.addStretch()
        params_layout.addLayout(obs_h_layout)

        # Max distance
        max_dist_layout = QHBoxLayout()
        max_dist_layout.addWidget(QLabel("Max Distance (m):"))
        self.max_distance_spin = QSpinBox()
        self.max_distance_spin.setMinimum(100)
        self.max_distance_spin.setMaximum(100000)
        self.max_distance_spin.setValue(1200)
        max_dist_layout.addWidget(self.max_distance_spin)
        max_dist_layout.addStretch()
        params_layout.addLayout(max_dist_layout)

        params_group.setLayout(params_layout)
        main_layout.addWidget(params_group)

        # Output section
        output_group = QGroupBox("Output")
        output_layout = QHBoxLayout()
        self.output_line = QLineEdit()
        self.output_line.setPlaceholderText("Path for output GeoTIFF")
        output_browse = QPushButton("Browse...")
        output_browse.clicked.connect(self.browse_output)
        output_layout.addWidget(QLabel("Output:"))
        output_layout.addWidget(self.output_line)
        output_layout.addWidget(output_browse)
        output_group.setLayout(output_layout)
        main_layout.addWidget(output_group)

        # Progress section
        self.progress_bar = QProgressBar()
        self.progress_bar.setVisible(False)
        main_layout.addWidget(self.progress_bar)

        self.status_text = QTextEdit()
        self.status_text.setReadOnly(True)
        self.status_text.setMaximumHeight(100)
        self.status_text.setVisible(False)
        main_layout.addWidget(self.status_text)

        # Buttons
        button_layout = QHBoxLayout()
        self.run_button = QPushButton("Run Viewshed Calculation")
        self.run_button.clicked.connect(self.run_calculation)
        self.cancel_process_button = QPushButton("Cancel Process")
        self.cancel_process_button.setEnabled(False)
        self.cancel_process_button.clicked.connect(self.cancel_calculation)
        close_button = QPushButton("Close")
        close_button.clicked.connect(self.reject)
        button_layout.addWidget(self.run_button)
        button_layout.addWidget(self.cancel_process_button)
        button_layout.addWidget(close_button)
        main_layout.addLayout(button_layout)

        self.setLayout(main_layout)

    def toggle_chm(self):
        """Toggle CHM inputs."""
        enabled = self.use_chm_check.isChecked()
        self.chm_line.setEnabled(enabled)
        self.chm_browse_btn.setEnabled(enabled)
        self.chm_layer_combo.setEnabled(enabled)

    def browse_dsm(self):
        """Browse for DSM file."""
        file_path, _ = QFileDialog.getOpenFileName(
            self, "Select DSM raster", "", "GeoTIFF (*.tif *.tiff);;All files (*)"
        )
        if file_path:
            self.dsm_line.setText(file_path)

    def browse_chm(self):
        """Browse for CHM file."""
        file_path, _ = QFileDialog.getOpenFileName(
            self, "Select CHM raster", "", "GeoTIFF (*.tif *.tiff);;All files (*)"
        )
        if file_path:
            self.chm_line.setText(file_path)

    def browse_viewpoints(self):
        """Browse for viewpoints shapefile."""
        file_path, _ = QFileDialog.getOpenFileName(
            self, "Select viewpoints shapefile", "", "Shapefile (*.shp)"
        )
        if file_path:
            self.vp_line.setText(file_path)

    def browse_output(self):
        """Browse for output file."""
        file_path, _ = QFileDialog.getSaveFileName(
            self, "Save output raster as", "", "GeoTIFF (*.tif);;All files (*)"
        )
        if file_path:
            if not file_path.lower().endswith('.tif'):
                file_path += '.tif'
            self.output_line.setText(file_path)

    def browse_viewpoints_polygon(self):
        """Browse for polygon shapefile used for grid generation."""
        file_path, _ = QFileDialog.getOpenFileName(
            self, "Select polygon shapefile", "", "Shapefile (*.shp);;All files (*)"
        )
        if file_path:
            self.vp_polygon_line.setText(file_path)

    def update_generation_source_ui(self):
        """Switch generated viewpoint inputs between extent and polygon source."""
        use_extent = self.vp_generate_source_combo.currentIndex() == 0
        self.vp_extent_widget.setVisible(use_extent)
        self.vp_polygon_widget.setVisible(not use_extent)

    def set_extent_from_canvas(self):
        """Populate extent fields from the current map canvas extent."""
        canvas = self.iface.mapCanvas()
        if canvas is None:
            QMessageBox.warning(self, "Extent Error", "Map canvas not available")
            return

        extent = canvas.extent()
        self.vp_xmin_spin.setValue(extent.xMinimum())
        self.vp_xmax_spin.setValue(extent.xMaximum())
        self.vp_ymin_spin.setValue(extent.yMinimum())
        self.vp_ymax_spin.setValue(extent.yMaximum())

    def _build_points_from_extent(self, spacing):
        """Generate a regular point grid from extent values."""
        xmin = self.vp_xmin_spin.value()
        xmax = self.vp_xmax_spin.value()
        ymin = self.vp_ymin_spin.value()
        ymax = self.vp_ymax_spin.value()

        if xmax <= xmin or ymax <= ymin:
            raise ValueError("Invalid extent: max values must be greater than min values")

        points = []
        y = ymin
        while y <= ymax:
            x = xmin
            while x <= xmax:
                points.append((x, y))
                x += spacing
            y += spacing

        return points

    def _build_points_from_polygon(self, spacing):
        """Generate a regular point grid clipped by a polygon layer source."""
        polygon_path, polygon_layer_name = self._resolve_vector_source(
            self.vp_polygon_line, self.vp_polygon_layer_combo
        )
        if not polygon_path:
            raise ValueError("Please select a polygon shapefile or project polygon layer")

        polygon_ds = ogr.Open(polygon_path)
        if polygon_ds is None:
            raise ValueError(f"Cannot open polygon source: {polygon_path}")

        if polygon_layer_name:
            polygon_layer = polygon_ds.GetLayerByName(polygon_layer_name)
            if polygon_layer is None:
                polygon_layer = polygon_ds.GetLayer(0)
        else:
            polygon_layer = polygon_ds.GetLayer(0)

        if polygon_layer.GetFeatureCount() == 0:
            raise ValueError("Polygon source has no features")

        mask_geometry = None
        polygon_layer.ResetReading()
        for feature in polygon_layer:
            geom = feature.GetGeometryRef()
            if geom is None or geom.IsEmpty():
                continue
            if mask_geometry is None:
                mask_geometry = geom.Clone()
            else:
                mask_geometry = mask_geometry.Union(geom)

        if mask_geometry is None or mask_geometry.IsEmpty():
            raise ValueError("Polygon source has no valid geometry")

        xmin, xmax, ymin, ymax = mask_geometry.GetEnvelope()
        if xmax <= xmin or ymax <= ymin:
            raise ValueError("Invalid polygon extent")

        points = []
        probe = ogr.Geometry(ogr.wkbPoint)
        y = ymin
        while y <= ymax:
            x = xmin
            while x <= xmax:
                probe.Empty()
                probe.AddPoint(x, y)
                if mask_geometry.Contains(probe) or mask_geometry.Intersects(probe):
                    points.append((x, y))
                x += spacing
            y += spacing

        return points

    def _resolve_output_path(self):
        """Return output path, creating a temporary target when user leaves it blank."""
        output_path = self.output_line.text().strip()
        if output_path:
            return output_path

        fd, temp_path = tempfile.mkstemp(prefix="viewshed_", suffix=".tif")
        os.close(fd)
        # Remove placeholder file so GDAL can create the raster cleanly.
        try:
            os.unlink(temp_path)
        except OSError:
            pass

        self.output_line.setText(temp_path)
        self.status_text.append(f"No output path provided. Using temporary file: {temp_path}")
        return temp_path

    def run_calculation(self):
        """Run the viewshed calculation."""
        dsm_path = self._resolve_raster_source(self.dsm_line, self.dsm_layer_combo)
        chm_path = self._resolve_raster_source(self.chm_line, self.chm_layer_combo)
        output_path = self._resolve_output_path()
        viewpoints_path = None
        viewpoints_layer_name = None
        generated_points = []

        try:
            if self.viewpoints_tabs.currentIndex() == 0:
                viewpoints_path, viewpoints_layer_name = self._resolve_vector_source(
                    self.vp_line, self.vp_layer_combo
                )
            else:
                spacing = self.vp_spacing_spin.value()
                if spacing <= 0:
                    raise ValueError("Viewpoint spacing must be greater than zero")

                if self.vp_generate_source_combo.currentIndex() == 0:
                    generated_points = self._build_points_from_extent(spacing)
                else:
                    generated_points = self._build_points_from_polygon(spacing)

                if not generated_points:
                    raise ValueError("Generated viewpoints are empty")
        except ValueError as e:
            QMessageBox.warning(self, "Input Error", str(e))
            return

        # Validate inputs
        if not dsm_path:
            QMessageBox.warning(self, "Input Error", "Please select a DSM file or project raster layer")
            return
        if self.viewpoints_tabs.currentIndex() == 0 and not viewpoints_path:
            QMessageBox.warning(
                self,
                "Input Error",
                "Please select a viewpoints shapefile or project point layer"
            )
            return
        if self.use_chm_check.isChecked() and not chm_path:
            QMessageBox.warning(
                self,
                "Input Error",
                "Please select a CHM file/project raster layer or uncheck CHM option"
            )
            return

        # Prepare parameters
        if self.viewpoints_tabs.currentIndex() == 0:
            raster_resolution = self.vp_existing_resolution_spin.value()
        else:
            # In generated-grid mode output raster cell size follows viewpoint spacing.
            raster_resolution = self.vp_spacing_spin.value()

        params = {
            'dsm_path': dsm_path,
            'viewpoints_path': viewpoints_path,
            'viewpoints_layer_name': viewpoints_layer_name,
            'generated_points': generated_points,
            'output_path': output_path,
            'max_distance': self.max_distance_spin.value(),
            'observer_height': self.observer_height_spin.value(),
            'raster_resolution': raster_resolution,
            'use_chm': self.use_chm_check.isChecked(),
            'chm_path': chm_path if self.use_chm_check.isChecked() else None,
        }

        # Disable controls
        self.run_button.setEnabled(False)
        self.cancel_process_button.setEnabled(True)
        self.progress_bar.setVisible(True)
        self.status_text.setVisible(True)
        self.status_text.clear()

        # Start worker
        self.worker = ViewshedWorker(params)
        self.worker.progress.connect(self.on_progress)
        self.worker.finished.connect(self.on_finished)
        self.worker.start()

    def cancel_calculation(self):
        """Request cancellation of an active viewshed calculation."""
        if self.worker is not None and self.worker.isRunning():
            self.worker.stop()
            self.cancel_process_button.setEnabled(False)
            self.status_text.setVisible(True)
            self.status_text.append("Cancellation requested. Waiting for current step to finish...")

    def on_progress(self, value, message):
        """Handle progress update."""
        self.progress_bar.setValue(value)
        self.status_text.append(message)

    def _resolve_raster_source(self, line_edit, layer_combo):
        """Resolve raster source path from project layer selection or manual path."""
        layer = layer_combo.currentLayer()
        if layer is not None:
            source = layer.source() or ""
            if source:
                return source.split('|', 1)[0]
        return line_edit.text().strip()

    def _resolve_vector_source(self, line_edit, layer_combo):
        """Resolve vector source path and layer name from selection or manual path."""
        layer = layer_combo.currentLayer()
        if layer is not None:
            source = layer.source() or ""
            if source:
                source_parts = source.split('|')
                base_path = source_parts[0]
                layer_name = None

                for part in source_parts[1:]:
                    if part.lower().startswith('layername='):
                        layer_name = part.split('=', 1)[1]
                        break

                if layer_name is None:
                    layer_name = layer.name()

                return base_path, layer_name

        return line_edit.text().strip(), None

    def on_finished(self, success, message):
        """Handle calculation finished."""
        self.run_button.setEnabled(True)
        self.cancel_process_button.setEnabled(False)
        self.progress_bar.setVisible(False)

        self.worker = None

        if message == "Cancelled by user":
            QMessageBox.information(self, "Cancelled", message)
            self.status_text.append("\nCalculation cancelled by user.")
            return

        if success:
            QMessageBox.information(self, "Success", message)
            self.status_text.append("\n✓ Calculation complete!")

            # Load output raster to QGIS
            output_path = self.output_line.text()
            loaded_layer = self.iface.addRasterLayer(output_path, Path(output_path).stem)
            if not loaded_layer:
                fallback_layer = QgsRasterLayer(output_path, Path(output_path).stem)
                if fallback_layer.isValid():
                    QgsProject.instance().addMapLayer(fallback_layer)
                    loaded_layer = fallback_layer
            if loaded_layer:
                self.status_text.append("Output layer added to QGIS project.")
            else:
                self.status_text.append("Could not automatically load output layer in QGIS.")
        else:
            QMessageBox.critical(self, "Error", message)
            self.status_text.append(f"\n✗ {message}")
