# -*- coding: utf-8 -*-
"""📦 Subnational Aggregation Processor module.

This module contains functionality for subnational aggregation processor.
"""

import os
import shutil
import traceback
from typing import Optional

from qgis import processing
from qgis.core import QgsCoordinateReferenceSystem, QgsTask, QgsVectorLayer

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


class SubnationalAggregationProcessingTask(QgsTask):
    """
    A QgsTask subclass for calculating GeoE3 x Population Score and or GeoE3 Score per aggregation area.

    It iterates over subnational boundaries, calculates the majoriy GeoE3 Score or GeoE3 x Population Score
    into a geopackage with the same polygons as original subnational boundaries but new attributes for
    the majority scores and applies a QML style.

    The subnational boundaries will NOT be split bu the study area polygons, so
    containment is not guaranteed between study area polygons and the subnational boundary polygons.
    Because of this, we will use the VRT combined outputs for the GeoE3 Score and GeoE3 x Population Score
    inputs.

    It will write 4 columns to the output gpkg:

    fid - generic incrementing id for each subnational area
    name - subnational area name
    geoe3_score - majority GeoE3 Score for that subnational area
    geoe3_pop_score - majority geoe3 x population score for that subnational area

    The GeoE3 Score can be one of 5 classes:

    | Range  | Description               | Color      |
    |--------|---------------------------|------------|
    | 0 - 1  | Very Low Enablement       | ![#FF0000](#) `#FF0000` |
    | 1 - 2  | Low Enablement            | ![#FFA500](#) `#FFA500` |
    | 2 - 3  | Moderately Enabling       | ![#FFFF00](#) `#FFFF00` |
    | 3 - 4  | Enabling                  | ![#90EE90](#) `#90EE90` |
    | 4 - 5  | Highly Enabling           | ![#0000FF](#) `#0000FF` |

    The GeoE3 x Population Score can be one of 15 classes:

    | Color      | Description                                 |
    |------------|---------------------------------------------|
    | ![#FF0000](#) `#FF0000` | Very low enablement, low population       |
    | ![#FF0000](#) `#FF0000` | Very low enablement, medium population    |
    | ![#FF0000](#) `#FF0000` | Very low enablement, high population      |
    | ![#FFA500](#) `#FFA500` | Low enablement, low population            |
    | ![#FFA500](#) `#FFA500` | Low enablement, medium population         |
    | ![#FFA500](#) `#FFA500` | Low enablement, high population           |
    | ![#FFFF00](#) `#FFFF00` | Moderately enabling, low population       |
    | ![#FFFF00](#) `#FFFF00` | Moderately enabling, medium population    |
    | ![#FFFF00](#) `#FFFF00` | Moderately enabling, high population      |
    | ![#90EE90](#) `#90EE90` | Enabling, low population                  |
    | ![#90EE90](#) `#90EE90` | Enabling, medium population               |
    | ![#90EE90](#) `#90EE90` | Enabling, high population                 |
    | ![#0000FF](#) `#0000FF` | Highly enabling, low population           |
    | ![#0000FF](#) `#0000FF` | Highly enabling, medium population        |
    | ![#0000FF](#) `#0000FF` | Highly enabling, high population          |

    See the geoe3_by_population_score_processor.py module for more details on how this is computed.

    📒 The majority score for each subnational area is calculated by counting the
       number of pixels in each GeoE3 Score and GeoE3 x Population Score class as per the
       above tables. Then majority pixel count is then allocated to the aggregate area.
       In the unlikely event of there being two or more classes with an equal pixel count,
       the highest enablement and population class is assigned.

    See zonalstatisticsfb algorithm for more details on how the majority score is calculated.
    https://qgis.org/pyqgis/3.34/analysis/QgsZonalStatistics.html#qgis.analysis.QgsZonalStatistics.Majority


    Args:
        study_area_gpkg_path (str): Path to the study area geopackage. Used to determine the CRS.
        working_directory (str): Parent directory to save the output agregated data. Outputs will
            be saved in a subdirectory called "subnational_aggregates".
        target_crs (Optional[QgsCoordinateReferenceSystem]): CRS for the output rasters.
        force_clear (bool): Flag to force clearing of all outputs before processing.

    Note: item.attribute("aggregation_layer_source") must be set for the analysis item.
    """

    def __init__(
        self,
        item: JsonTreeItem,
        study_area_gpkg_path: str,
        working_directory: str,
        target_crs: Optional[QgsCoordinateReferenceSystem] = None,
        force_clear: bool = False,
    ):
        super().__init__("Subnational Aggregation Processor", QgsTask.CanCancel)
        self.item = item  # should be an item with role=analysis

        self.study_area_gpkg_path = study_area_gpkg_path

        self.aggregation_areas_path = item.attribute("aggregation_layer_source")

        self.aggregation_layer: QgsVectorLayer = QgsVectorLayer(
            self.aggregation_areas_path,
            "aggregation_areas",
            "ogr",
        )
        if not self.aggregation_layer.isValid():
            raise Exception(f"Invalid aggregation areas layer:\n{self.aggregation_areas_path}")  # noqa E231

        self.output_dir = os.path.join(working_directory, "subnational_aggregation")
        os.makedirs(self.output_dir, exist_ok=True)

        # This folder should already exist from the aggregation analysis and population raster processing
        self.geoe3_score_folder = os.path.join(working_directory, "geoe3_score")

        if not os.path.exists(self.geoe3_score_folder):
            raise Exception(
                f"GeoE3 folder not found.\n{self.geoe3_score_folder}\nPlease run GeoE3 raster processing first."
            )
        # This folder is optional  - user needs to have configured population layer in analysis properties dialog
        self.geoe3_by_population_folder = os.path.join(working_directory, "geoe3_by_population_score")

        if not os.path.exists(self.geoe3_by_population_folder):
            log_message(
                f"GeoE3 folder not found.\n{self.geoe3_by_population_folder}\nPlease run GeoE3 raster processing first."
            )
            self.geoe3_by_population_folder = None
        # These two folders are optional - will only be set if user configured mask for
        # job opportunities in analysis properties dialog
        self.geoe3_score_ghsl_masked_folder = os.path.join(working_directory, "geoe3_score_ghsl_masked")

        if not os.path.exists(self.geoe3_score_ghsl_masked_folder):
            log_message(
                f"GeoE3 GHSL masked folder not found.\n{self.geoe3_score_ghsl_masked_folder}\nPlease run GeoE3 raster processing first."
            )
            self.geoe3_score_ghsl_masked_folder = None

        self.geoe3_score_by_population_ghsl_masked_folder = os.path.join(
            working_directory, "geoe3_score_by_population_ghsl_masked"
        )

        if not os.path.exists(self.geoe3_score_by_population_ghsl_masked_folder):
            log_message(
                f"GeoE3 by population GHSL masked folder not found.\n{self.geoe3_score_by_population_ghsl_masked_folder}\nPlease run GeoE3 raster processing first."
            )
            self.geoe3_score_by_population_ghsl_masked_folder = None

        self.force_clear = force_clear
        if self.force_clear and os.path.exists(self.output_dir):
            for file in os.listdir(self.output_dir):
                os.remove(os.path.join(self.output_dir, file))

        self.target_crs = target_crs
        if not self.target_crs:
            layer: QgsVectorLayer = QgsVectorLayer(
                f"{self.study_area_gpkg_path}|layername=study_area_clip_polygons",
                "study_area_clip_polygons",
                "ogr",
            )
            self.target_crs = layer.crs()
            log_message(f"Target CRS not set. Using CRS from study area clip polygon: {self.target_crs.authid()}")
            log_message(f"{self.study_area_gpkg_path}|ayername=study_area_clip_polygon")
            del layer

        log_message("Initialized GeoE3 Subnational Area Aggregation Processing Task")

    def run(self) -> bool:
        """
        Executes the GeoE3 Subnational Area Aggregation Processing Task calculation task.

        Returns:
            bool: True if the task completed successfully, False otherwise.
        """
        try:
            cleaned_aggregation_layer = self.fix_geometries()
            if self.geoe3_score_folder:
                log_message("Calculating GeoE3 Score Aggregation")
                layer_name = "geoe3_score_subnational_aggregation"
                raster_layer_path = os.path.join(self.geoe3_score_folder, "GeoE3_Score_combined.vrt")
                output = self.aggregate(
                    cleaned_aggregation_layer,
                    raster_layer_path,
                    layer_name,
                )
                self.apply_qml_style(
                    source_qml=resources_path("resources", "qml", "geoe3_score_vector.qml"),
                    qml_path=os.path.join(self.output_dir, f"{layer_name}.qml"),
                )
                self.item.setAttribute(f"{layer_name}", f"{output}|layername={layer_name}")
            if self.geoe3_by_population_folder:
                log_message("Calculating GeoE3 x Population Score Aggregation")
                layer_name = "geoe3_by_population_subnational_aggregation"
                raster_layer_path = os.path.join(self.geoe3_by_population_folder, "geoe3_by_population_score.vrt")
                output = self.aggregate(
                    cleaned_aggregation_layer,
                    raster_layer_path,
                    layer_name,
                )
                self.apply_qml_style(
                    source_qml=resources_path("resources", "qml", "geoe3_by_population_vector_score.qml"),
                    qml_path=os.path.join(self.output_dir, f"{layer_name}.qml"),
                )
                self.item.setAttribute(f"{layer_name}", f"{output}|layername={layer_name}")
            if self.geoe3_score_ghsl_masked_folder:
                log_message("Calculating GeoE3 Score GHSL Masked Aggregation")
                layer_name = "geoe3_score_ghsl_masked_subnational_aggregation"
                raster_layer_path = os.path.join(
                    self.geoe3_score_ghsl_masked_folder,
                    "geoe3_score_ghsl_masked.vrt",
                )
                output = self.aggregate(
                    cleaned_aggregation_layer,
                    raster_layer_path,
                    layer_name,
                )
                self.apply_qml_style(
                    source_qml=resources_path("resources", "qml", "geoe3_score_vector.qml"),
                    qml_path=os.path.join(self.output_dir, f"{layer_name}.qml"),
                )
                self.item.setAttribute(f"{layer_name}", f"{output}|layername={layer_name}")
            if self.geoe3_score_by_population_ghsl_masked_folder:
                log_message("Calculating GeoE3 Score by Population GHSL Masked Aggregation")
                layer_name = "geoe3_score_by_population_ghsl_masked_subnational_aggregation"
                raster_layer_path = os.path.join(
                    self.geoe3_score_by_population_ghsl_masked_folder,
                    "geoe3_score_by_population_ghsl_masked.vrt",
                )
                output = self.aggregate(
                    cleaned_aggregation_layer,
                    raster_layer_path,
                    layer_name,
                )
                self.apply_qml_style(
                    source_qml=resources_path("resources", "qml", "geoe3_by_population_vector_score.qml"),
                    qml_path=os.path.join(self.output_dir, f"{layer_name}.qml"),
                )
                self.item.setAttribute(f"{layer_name}", f"{output}|layername={layer_name}")
            log_message(self.item.attributesAsMarkdown())
            return True
        except Exception as e:
            log_message(f"Task failed: {e}")
            log_message(traceback.format_exc())
            return False

    def fix_geometries(self) -> QgsVectorLayer:
        """
        Fix geometries in the aggregation layer.

        Returns:
            QgsVectorLayer: The layer with fixed geometries and retained fid field.
        """

        params = {
            "INPUT": self.aggregation_layer,
            "METHOD": 1,  # Structur method
            "OUTPUT": "TEMPORARY_OUTPUT",
        }
        output = processing.run("native:fixgeometries", params)["OUTPUT"]
        # Also drop all columns
        params = {"INPUT": output, "FIELDS": ["fid"], "OUTPUT": "TEMPORARY_OUTPUT"}
        output = processing.run("native:retainfields", params)["OUTPUT"]

        return output

    def aggregate(self, aggregation_layer: QgsVectorLayer, raster_layer_path: str, name: str) -> str:
        """
        Use aggregation vector to calculate the majority GeoE3 Score and GeoE3 x Population Score for each valid polygon.

        Args:
            aggregation_layer (QgsVectorLayer): The aggregation layer with fixed geometries.
            raster_layer_path (str): Path to the raster layer to aggregate.
            name (str): Name for the output layer.

        Returns:
            str: Path to the output geopackage file.
        """
        output = os.path.join(self.output_dir, f"{name}.gpkg")
        params = {
            "INPUT": aggregation_layer,
            "INPUT_RASTER": raster_layer_path,
            "RASTER_BAND": 1,
            "COLUMN_PREFIX": "_",
            "STATISTICS": [9],  # Majority
            "OUTPUT": output,
        }
        processing.run("native:zonalstatisticsfb", params)
        return output

    def apply_qml_style(self, source_qml: str, qml_path: str) -> None:
        """
        Apply QML style by copying from source to destination.

        Args:
            source_qml (str): Path to the source QML file.
            qml_path (str): Path where the QML file should be copied to.
        """

        log_message(f"Copying QML style from {source_qml} to {qml_path}")
        # Apply QML Style
        if os.path.exists(source_qml):
            shutil.copy(source_qml, qml_path)
        else:
            log_message("QML style file not found. Skipping QML copy.")

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

        Args:
            result (bool): The result of the task execution.
        """
        if result:
            log_message("Subnational aggregate calculation completed successfully.")
        else:
            log_message("Subnational aggregate calculation failed.")
