# -*- coding: utf-8 -*-
"""📦 Opportunities Mask Processor module.

This module contains functionality for opportunities mask processor.
"""

import os
import traceback
from typing import Optional
from urllib.parse import unquote

from qgis import processing
from qgis.core import (
    Qgis,
    QgsCoordinateReferenceSystem,
    QgsCoordinateTransform,
    QgsFeedback,
    QgsGeometry,
    QgsProcessingContext,
    QgsProcessingFeedback,
    QgsProject,
    QgsRasterLayer,
    QgsRectangle,
    QgsTask,
    QgsVectorLayer,
)

from geest.core import JsonTreeItem
from geest.utilities import log_message, resources_path

from .area_iterator import AreaIterator
from .ghsl_downloader import GHSLDownloader
from .ghsl_processor import GHSLProcessor
from .utilities import (
    check_and_reproject_layer,
    combine_rasters_to_vrt,
    geometry_to_memory_layer,
    subset_vector_layer,
)


class OpportunitiesMaskProcessor(QgsTask):
    """
    A QgsTask subclass for generating job opportunity mask layers.

    It iterates over bounding boxes and study areas, selects the intersecting features
    (if inputs are points or polygons) or in the case of a raster maske, clips the raster
    data to match the study area bbox.

    Args:
        item (JSONTreeItem): Analysis item containing the needed parameters.
        study_area_gpkg_path (str): Path to the GeoPackage containing study area masks.
        output_dir (str): Directory to save the output rasters.
        cell_size_m (float): Cell size for the output rasters.
        context (Optional[QgsProcessingContext]): QGIS processing context.
        feedback (Optional[QgsFeedback]): QGIS feedback object.
        force_clear (bool): Flag to force clearing of all outputs before processing.

    Concrete implementation of a geoe3 insight for masking by job opportunities.

    It will create a raster layer where all cells outside the masked areas (defined
    by the input polygons layer) are set to a no data value.

    This is used when you want to represent the GeoE3 Score and GeoE3 x Population Score
    only in areas where there are job opportunities / job creation initiatives.

    The input layer should be a polygon layer with the job opportunities. Its attributes
    are completely ignored, only the geometry is used to create a mask.

    The output raster will have the same extent and cell size as the study area.

    The output raster will have either have values of 1 in the mask and 0 outside.

    The output raster will be a vrt which is a composite of all the individual area rasters.

    The output raster will be saved in the working directory under a subfolder called 'opportunity_masks'.

    Preconditions:

    This workflow expects that the user has configured the root analysis node dialog with
    the polygon mask settings configured.

    Input can be any of:

    * a point layer (with a buffer distance)
    * a polygon layer (attributes are ignored)
    * a raster layer (any non-null pixel will be set to 1)

    """

    def __init__(
        self,
        item: JsonTreeItem,
        study_area_gpkg_path: str,
        working_directory: str,
        cell_size_m: float,
        context: Optional[QgsProcessingContext] = None,
        feedback: Optional[QgsFeedback] = None,
        force_clear: bool = False,
    ):
        super().__init__("Opportunities Mask Processor", QgsTask.CanCancel)
        self.study_area_gpkg_path = study_area_gpkg_path
        self.cell_size_m = cell_size_m
        self.working_directory = working_directory
        layer: QgsVectorLayer = QgsVectorLayer(
            f"{self.study_area_gpkg_path}|layername=study_area_clip_polygons",
            "study_area_clip_polygons",
            "ogr",
        )
        self.target_crs = layer.crs()
        del layer
        self.context = context
        self.feedback = feedback
        self.clipped_rasters = []
        self.item = item
        self.features_layer = None
        self.raster_layer = None
        self.mask_mode = self.item.attribute("mask_mode", None)  # if set,  will be "point", "polygon" or "raster"
        if not self.mask_mode:
            raise Exception("Mask mode not set in the analysis.")

        self.workflow_name = f"opportunities_{self.mask_mode}_mask"
        # In normal workflows this comes from the item, but this workflow is a bit different
        # so we set it manually.
        self.layer_id = "opportunities_mask"
        if self.mask_mode == "point":
            self.buffer_distance_m = self.item.attribute("buffer_distance_m", 1000)
        if self.mask_mode in ["point", "polygon"]:
            # There are two ways a user can specify the polygon mask layer
            # either as a shapefile path added in a line edit or as a layer source
            # using a QgsMapLayerComboBox. We prioritize the shapefile path, so check that first.
            layer_source = self.item.attribute(f"{self.mask_mode}_mask_shapefile", None)
            if layer_source:
                layer_source = unquote(layer_source)
            provider_type = "ogr"
            if not layer_source:
                # Fall back to the QgsMapLayerComboBox source
                layer_source = self.item.attribute(f"{self.mask_mode}_mask_layer_source", None)
                provider_type = self.item.attribute(f"{self.mask_mode}_mask_layer_provider_type", "ogr")
            if not layer_source:
                log_message(
                    f"{self.mask_mode}_mask_shapefile not found",
                    tag="GeoE3",
                    level=Qgis.Critical,
                )
                raise Exception(f"{self.mask_mode}_mask_shapefile not found")
            self.features_layer = QgsVectorLayer(layer_source, self.mask_mode, provider_type)
            if not self.features_layer.isValid():
                log_message(f"{self.mask_mode}_mask_shapefile not valid", level=Qgis.Critical)
                log_message(f"Layer Source: {layer_source}", level=Qgis.Critical)
                raise Exception(f"{self.mask_mode}_mask_shapefile not valid")

            # Check the geometries and reproject if necessary
            self.features_layer = check_and_reproject_layer(self.features_layer, self.target_crs)

        elif self.mask_mode == "raster":
            # Check the input raster is ok. The raster itself does not need to be a mask
            # (i.e. with 1 and nodata values) - we will take care of that in this class.
            # Then the _process_raster_for_area method is where we turn it into a mask
            log_message("Loading source raster mask layer")
            # First try the one defined in the line edit
            self.raster_layer = QgsRasterLayer(unquote(self.item.attribute("raster_mask_raster")), "Raster Mask", "ogr")
            if not self.raster_layer.isValid():
                # Then fall back to the QgsMapLayerComboBox source
                self.raster_layer = QgsRasterLayer(
                    self.item.attribute("raster_mask_layer_source"),
                    "Raster Mask",
                    self.item.attribute("raster_mask_layer_provider_type"),
                )
            if not self.raster_layer.isValid():
                log_message("No valid raster layer provided for mask", level=Qgis.Critical)
                log_message(f"Raster Source: {self.raster_layer.source()}", level=Qgis.Critical)
                raise Exception("No valid raster layer provided for mask")
        elif self.mask_mode == "ghsl":
            # Use the global human settlements layer defined in the project settings
            log_message("Loading global human settlements layer for mask")
            layer_source = os.path.join(working_directory, "study_area", "ghsl_settlements_layer.parquet")
            provider_type = "ogr"

            # Check if file exists and is not empty (0 bytes)
            needs_download = False
            if not os.path.exists(layer_source):
                log_message(f"GHSL parquet file not found: {layer_source}", level=Qgis.Warning)
                needs_download = True
            elif os.path.getsize(layer_source) == 0:
                log_message(f"GHSL parquet file is 0 bytes: {layer_source}", level=Qgis.Warning)
                needs_download = True

            if needs_download:
                log_message("Attempting to download and process GHSL data...")
                if not self._download_ghsl_data(working_directory):
                    log_message(
                        "Failed to download GHSL data for mask",
                        tag="GeoE3",
                        level=Qgis.Critical,
                    )
                    raise Exception("Failed to download GHSL data for mask")

            # Try loading the layer again after potential download
            if not os.path.exists(layer_source):
                log_message(
                    f"{self.mask_mode} parquet file not found after download attempt",
                    tag="GeoE3",
                    level=Qgis.Critical,
                )
                raise Exception(f"{self.mask_mode} parquet file not found")

            if os.path.getsize(layer_source) == 0:
                log_message(
                    f"{self.mask_mode} parquet file is 0 bytes after download attempt",
                    tag="GeoE3",
                    level=Qgis.Critical,
                )
                raise Exception(f"{self.mask_mode} parquet file is 0 bytes")

            self.features_layer = QgsVectorLayer(layer_source, self.mask_mode, provider_type)
            if not self.features_layer.isValid():
                log_message(f"{self.mask_mode} parquet file not valid", level=Qgis.Critical)
                log_message(f"Layer Source: {layer_source}", level=Qgis.Critical)
                raise Exception(f"{self.mask_mode} parquet file not valid")
            # Check the crs of the layer, if it is not in the target crs, reproject it
            if self.features_layer.crs() != self.target_crs:
                log_message("Reprojecting features layer to match target crs")
                # Note that the layer returned will be a memory layer
                # which may be inefficiecy here.
                log_message(f"Layer CRS: {self.features_layer.crs().authid()}")
                log_message(f"Target CRS: {self.target_crs.authid()}")
                self.features_layer = check_and_reproject_layer(self.features_layer, self.target_crs)

        # Workflow directory is the subdir under working_directory
        self.workflow_directory = os.path.join(working_directory, "opportunity_masks")
        os.makedirs(self.workflow_directory, exist_ok=True)

        self.output_filename = "opportunities_mask"
        # And customise which key we will write the result file to:
        self.result_file_key = "opportunities_mask_result_file"
        self.result_key = "opportunities_mask_result"

        # TODO make user configurable
        self.force_clear = False
        if self.force_clear and os.path.exists(self.workflow_directory):
            for file in os.listdir(self.workflow_directory):
                os.remove(os.path.join(self.workflow_directory, file))

        self.mask_list = []

        log_message("---------------------------------------------")
        log_message("Initialized GeoE3 Opportunities Mask Workflow")
        log_message("---------------------------------------------")
        log_message(f"Item: {self.item.name}")
        log_message(f"Study area GeoPackage path: {self.study_area_gpkg_path}")
        log_message(f"Working_directory: {self.working_directory}")
        log_message(f"Workflow directory: {self.workflow_directory}")
        log_message(f"Mask mode: {self.mask_mode}")
        layer_source = None
        if self.mask_mode in ["point", "polygon", "ghsl"] and self.features_layer:
            layer_source = self.features_layer.source()
        elif self.raster_layer:
            layer_source = self.raster_layer.source()
        log_message(f"Layer path: {layer_source if layer_source else 'Not set'}")
        log_message(f"Cell size: {self.cell_size_m}")
        log_message(f"CRS: {self.target_crs.authid() if self.target_crs else 'None'}")
        log_message(f"Force clear: {self.force_clear}")
        log_message("---------------------------------------------")

    def run(self) -> bool:
        """
        Executes the task to process mask for each area.

        Returns:
            bool: True if the task completed successfully, False otherwise.
        """
        try:
            area_iterator = AreaIterator(self.study_area_gpkg_path)
            for index, (current_area, clip_area, current_bbox, progress) in enumerate(area_iterator):
                if self.feedback and self.feedback.isCanceled():
                    return False
                if self.mask_mode == "raster":
                    area_raster = self._subset_raster_layer(current_bbox, index)
                    mask_layer = self._process_raster_for_area(
                        current_area, clip_area, current_bbox, area_raster, index
                    )
                else:
                    vector_layer = subset_vector_layer(
                        self.workflow_directory,
                        self.features_layer,
                        current_area,
                        str(index),
                    )
                    mask_layer = self._process_features_for_area(
                        current_area, clip_area, current_bbox, vector_layer, index
                    )
                if mask_layer:
                    self.mask_list.append(mask_layer)

            vrt_filepath = os.path.join(
                self.workflow_directory,
                f"{self.output_filename}_combined.vrt",
            )
            source_qml = resources_path("resources", "qml", "mask.qml")
            vrt_filepath = combine_rasters_to_vrt(self.mask_list, self.target_crs, vrt_filepath, source_qml)
            self.item.setAttribute(self.result_file_key, vrt_filepath)
            self.item.setAttribute(self.result_key, "Opportunities Mask Created OK")
        except Exception as e:
            log_message(f"Task failed: {e}")
            log_message(traceback.format_exc())
            self.item.setAttribute(self.result_key, str(e))
            return False

    def finished(self, result: bool) -> None:
        """
        Called when the task completes.

        Args:
            result (bool): The result of the task execution.
        """
        if result:
            log_message("Opportunities mask processing completed successfully.")
        else:
            log_message("Opportunities mask processing failed.")

    def _download_ghsl_data(self, working_directory: str) -> bool:
        """Download and process GHSL data for the study area if not already present.

        This method downloads GHSL tiles that intersect the study area,
        processes them, and saves the results to ghsl_settlements_layer.parquet.

        Args:
            working_directory: The project working directory.

        Returns:
            True if GHSL data was successfully downloaded/processed, False otherwise.
        """
        try:
            log_message("Downloading GHSL data for opportunities mask...")

            # Get study area extent for GHSL download
            study_area_dir = os.path.join(working_directory, "study_area")
            study_area_path = os.path.join(study_area_dir, "study_area.gpkg")

            if not os.path.exists(study_area_path):
                log_message(f"Study area not found: {study_area_path}", level=Qgis.Critical)
                return False

            # Load study area to get extent
            study_layer = QgsVectorLayer(study_area_path, "study_area", "ogr")
            if not study_layer.isValid():
                log_message("Study area layer is invalid", level=Qgis.Critical)
                return False

            extent = study_layer.extent()
            src_crs = study_layer.crs()

            # Transform extent to Mollweide (ESRI:54009) for GHSL
            dst_crs = QgsCoordinateReferenceSystem("ESRI:54009")
            transform = QgsCoordinateTransform(src_crs, dst_crs, QgsProject.instance())
            extent_mollweide = transform.transformBoundingBox(extent)

            log_message(f"Study area extent in Mollweide: {extent_mollweide.toString()}")

            # Download GHSL tiles
            downloader = GHSLDownloader(
                extents=extent_mollweide,
                output_path=study_area_dir,
                filename="ghsl_temp",
                use_cache=True,
                delete_existing=False,
                feedback=self.feedback,
            )

            tiles = downloader.tiles_intersecting_bbox()
            if not tiles:
                log_message("No GHSL tiles intersect study area", level=Qgis.Warning)
                return False

            log_message(f"Downloading {len(tiles)} GHSL tiles...")
            tile_paths = []
            for tile_id in tiles:
                paths = downloader.download_and_unpack_tile(tile_id)
                # Filter to only .tif files
                tif_paths = [p for p in paths if p.endswith(".tif")]
                tile_paths.extend(tif_paths)

            if not tile_paths:
                log_message("No GHSL tiles downloaded", level=Qgis.Warning)
                return False

            log_message(f"Downloaded {len(tile_paths)} GHSL .tif files")

            # Process tiles
            log_message("Processing GHSL tiles...")
            processor = GHSLProcessor(input_raster_paths=tile_paths)

            # Reclassify
            reclassified = processor.reclassify_rasters(suffix="reclass")

            # Polygonize
            polygonized = processor.polygonize_rasters(reclassified)

            # Combine vectors
            output_parquet = os.path.join(study_area_dir, "ghsl_settlements_layer.parquet")
            result = processor.combine_vectors(polygonized, output_parquet, extent_mollweide)

            if result and os.path.exists(output_parquet):
                file_size = os.path.getsize(output_parquet)
                if file_size > 0:
                    log_message(f"GHSL data downloaded and processed: {output_parquet} ({file_size} bytes)")
                    return True
                else:
                    log_message("GHSL output file is 0 bytes", level=Qgis.Warning)
                    return False
            else:
                log_message("Failed to combine GHSL vectors", level=Qgis.Warning)
                return False

        except Exception as e:
            log_message(f"Error downloading GHSL data: {str(e)}", level=Qgis.Warning)
            log_message(traceback.format_exc(), level=Qgis.Warning)
            return False

    def _process_features_for_area(
        self,
        current_area: QgsGeometry,
        clip_area: QgsGeometry,
        current_bbox: QgsGeometry,
        area_features: QgsVectorLayer,
        index: int,
    ) -> str:
        """
        Executes the actual workflow logic for a single area.

        Args:
            current_area (QgsGeometry): Current polygon from our study area.
            clip_area (QgsGeometry): Polygon to clip the features by.
            current_bbox (QgsGeometry): Bounding box of the current area.
            area_features (QgsVectorLayer): A vector layer of features to analyse that includes only features in the study area.
            index (int): Iteration / number of area being processed.

        Returns:
            str: A raster layer file path if processing completes successfully.
        """
        log_message(f"{self.workflow_name}  Processing Started for area {index}")
        log_message(f"Mask mode: {self.mask_mode}")
        if self.mask_mode == "point":
            log_message("Buffering job opportunity points")
            buffered_points_layer = self._buffer_features(area_features, index)
            log_message("Clipping features to the current area's clip area")
            clipped_layer = self._clip_features(buffered_points_layer, clip_area, index)
            log_message(f"Clipped features saved to {clipped_layer.source()}")
            log_message("Generating mask layer")
            mask_layer = self.generate_mask_layer(clipped_layer, current_bbox, index)
            log_message(f"Mask layer saved to {mask_layer}")
            return mask_layer
        elif self.mask_mode in ["polygon", "ghsl"]:
            # Step 1: clip the selected features to the current area's clip area
            log_message("Clipping features to the current area's clip area")
            clipped_layer = self._clip_features(area_features, clip_area, index)
            log_message(f"Clipped features saved to {clipped_layer.source()}")
            log_message("Generating mask layer")
            mask_layer = self.generate_mask_layer(clipped_layer, current_bbox, index)
            log_message(f"Mask layer saved to {mask_layer}")
            return mask_layer
        else:  # Raster
            pass

    def _buffer_features(self, layer: QgsVectorLayer, index: int) -> QgsVectorLayer:
        """
        Buffer the input features by the buffer_distance m.

        Args:
            layer (QgsVectorLayer): The input feature layer.
            index (int): The index of the current area.

        Returns:
            QgsVectorLayer: The buffered features layer.
        """
        output_name = f"opportunites_points_buffered_{index}"

        output_path = os.path.join(self.workflow_directory, f"{output_name}.shp")
        buffered_layer = processing.run(
            "native:buffer",
            {
                "INPUT": layer,
                "DISTANCE": self.buffer_distance_m,
                "SEGMENTS": 15,
                "DISSOLVE": True,
                "OUTPUT": output_path,
            },
        )["OUTPUT"]

        buffered_layer = QgsVectorLayer(output_path, output_name, "ogr")
        return buffered_layer

    def _clip_features(self, layer: QgsVectorLayer, clip_area: QgsGeometry, index: int) -> QgsVectorLayer:
        """
        Clip the input features by the clip area.

        Args:
            layer (QgsVectorLayer): The input feature layer.
            clip_area (QgsGeometry): The geometry to clip the features by.
            index (int): The index of the current area.

        Returns:
            QgsVectorLayer: The mask features layer clipped to the clip area.
        """
        output_name = f"opportunites_polygons_clipped_{index}"
        clip_layer = geometry_to_memory_layer(clip_area, self.target_crs, "clip_area")
        output_path = os.path.join(self.workflow_directory, f"{output_name}.shp")
        params = {"INPUT": layer, "OVERLAY": clip_layer, "OUTPUT": output_path}
        output = processing.run("native:clip", params)["OUTPUT"]
        del output
        clipped_layer = QgsVectorLayer(output_path, output_name, "ogr")
        return clipped_layer

    def generate_mask_layer(self, clipped_layer: QgsVectorLayer, current_bbox: QgsGeometry, index: int) -> str:
        """Generate the mask layer.

        This will be used to create a mask by rasterizing the input polygon layer.

        Args:
            clipped_layer (QgsVectorLayer): The clipped vector mask layer.
            current_bbox (QgsGeometry): The bounding box of the current area.
            index (int): The index of the current area.

        Returns:
            str: Path to the mask raster layer generated from the input clipped polygon layer.
        """

        rasterized_polygons_path = os.path.join(self.workflow_directory, f"opportunites_mask_{index}.tif")

        params = {
            "INPUT": clipped_layer,
            "FIELD": None,
            "BURN": 1,
            "USE_Z": False,
            "UNITS": 1,
            "WIDTH": self.cell_size_m,
            "HEIGHT": self.cell_size_m,
            "EXTENT": current_bbox.boundingBox(),
            "NODATA": 0,
            "OPTIONS": "",
            "DATA_TYPE": 0,  # byte
            "INIT": None,
            "INVERT": False,
            "EXTRA": "-co NBITS=1 -at",  # -at is for all touched cells
            "OUTPUT": rasterized_polygons_path,
        }

        output = processing.run("gdal:rasterize", params)["OUTPUT"]
        del output
        return rasterized_polygons_path

    def _subset_raster_layer(self, bbox: QgsGeometry, index: int) -> str:
        """
        Reproject and clip the raster to the bounding box of the current area.

        Overloaded version of the same method in the base class because that one
        fills the raster replacing nodata with 0 which is not what we want for this workflow.

        Args:
            bbox (QgsGeometry): The bounding box of the current area.
            index (int): The index of the current area.

        Returns:
            str: The path to the reprojected and clipped raster.
        """
        # Convert the bbox to QgsRectangle
        bbox: QgsRectangle = bbox.boundingBox()

        reprojected_raster_path = os.path.join(
            self.workflow_directory,
            f"{self.layer_id}_clipped_and_reprojected_{index}.tif",
        )

        params = {
            "INPUT": self.raster_layer,
            "TARGET_CRS": self.target_crs,
            "RESAMPLING": 0,
            "TARGET_RESOLUTION": self.cell_size_m,
            "NODATA": -9999,
            "OUTPUT": reprojected_raster_path,
            "TARGET_EXTENT": f"{bbox.xMinimum()},{bbox.xMaximum()},{bbox.yMinimum()},{bbox.yMaximum()} [{self.target_crs.authid()}]",  # noqa E231
        }

        processing.run("gdal:warpreproject", params, feedback=QgsProcessingFeedback())
        return reprojected_raster_path

    def _process_raster_for_area(
        self,
        current_area: QgsGeometry,
        clip_area: QgsGeometry,
        current_bbox: QgsGeometry,
        area_raster: str,
        index: int,
    ) -> str:
        """
        Executes the actual workflow logic for a single area using a raster.

        In this case we will convert it to a mask raster where any non-null pixel
        is given a value of 1 and all other pixels are set to nodata.

        Args:
            current_area (QgsGeometry): Current polygon from our study area.
            clip_area (QgsGeometry): Polygon to clip the raster to.
            current_bbox (QgsGeometry): Bounding box of the current area.
            area_raster (str): A raster layer of features to analyse that includes only bbox pixels in the study area.
            index (int): Index of the current area.

        Returns:
            str: Path to the reclassified raster.
        """
        log_message(f"{self.workflow_name}  Processing Raster Started for area {index}")
        opportunities_mask_path = os.path.join(self.workflow_directory, f"opportunities_mask_{index}.tif")
        # 📒 NOTE: Explanation for the formula logic:
        #
        # (A!=A) identifies NoData cells (since NaN != NaN is True for NoData cells) and sets them to 0.
        # (A==A) identifies valid cells and sets them to 1.
        # This ensures that all valid cells are set to 1 and NoData cells remain as NoData.
        #
        params = {
            "INPUT_A": area_raster,
            "BAND_A": 1,
            "FORMULA": "(A!=A)*0+(A==A)*1",
            "NO_DATA": None,
            "EXTENT_OPT": 0,
            "PROJWIN": None,
            "RTYPE": 0,
            "OPTIONS": "",
            "EXTRA": "",
            "OUTPUT": opportunities_mask_path,
        }

        processing.run("gdal:rastercalculator", params)
        return opportunities_mask_path
