# -*- coding: utf-8 -*-

"""
/***************************************************************************
 LandCoverDownloader
                                 A QGIS plugin
 Downloads Sentinel-2 Land Cover Data from Esri's Living Atlas
 Generated by Plugin Builder: http://g-sherman.github.io/Qgis-Plugin-Builder/
                              -------------------
        begin                : 2026-01-11
        copyright            : (C) 2026 by CustomCartographix
        email                : john@customcartographix.com
 ***************************************************************************/

/***************************************************************************
 *                                                                         *
 *   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 script contains the necessary functions for the land cover       *
 *   downloader processing algorithm to work.                              *
 *                                                                         *
 ***************************************************************************/
"""

# Import necessary modules
from os import path

from PyQt5.QtCore import QUrl
from PyQt5.QtNetwork import QNetworkRequest, QNetworkReply
from qgis.core import (QgsNetworkAccessManager,
                       QgsProcessingException,
                       QgsProcessingLayerPostProcessorInterface,
                       QgsProject,
                       QgsVectorFileWriter,
                       QgsVectorLayer)
from qgis import processing


# Bundled QML styles applied to the algorithm outputs.
#   * Land cover raster → Esri Living Atlas paletted style.
#   * AOI polygon        → red outline, no fill.
LAND_COVER_STYLE_PATH = path.join(
    path.dirname(__file__), 'resources', 'land_cover_style.qml')
AOI_OUTLINE_STYLE_PATH = path.join(
    path.dirname(__file__), 'resources', 'aoi_outline_style.qml')


class _StyleFromQmlPostProcessor(QgsProcessingLayerPostProcessorInterface):
    """
    Base Processing post-processor that loads a bundled QML style onto a layer
    once QGIS has added it to the project.

    Subclasses set ``STYLE_PATH`` (an absolute path to a ``.qml`` file) and
    optionally override ``STYLE_LABEL`` for user-facing messages.

    QGIS keeps only a weak reference to post-processors, so instances must be
    held elsewhere. Each subclass caches its own instance on ``_instance``
    (attribute lookup on ``cls`` shadows the base's ``None``).
    """

    STYLE_PATH = None
    STYLE_LABEL = 'style'
    _instance = None

    def postProcessLayer(self, layer, context, feedback):
        if layer is None or not layer.isValid() or self.STYLE_PATH is None:
            return
        result = layer.loadNamedStyle(self.STYLE_PATH)
        # loadNamedStyle returns (message, success) on QGIS 3.x
        if isinstance(result, tuple) and len(result) == 2 and not result[1]:
            feedback.pushInfo(
                'Could not apply ' + self.STYLE_LABEL + ': ' + str(result[0]))
        layer.triggerRepaint()

    @classmethod
    def create(cls):
        cls._instance = cls()
        return cls._instance


class SetLandCoverStylePostProcessor(_StyleFromQmlPostProcessor):
    """Applies the Esri Living Atlas paletted raster style to the output raster."""
    STYLE_PATH = LAND_COVER_STYLE_PATH
    STYLE_LABEL = 'land cover style'


class SetAoiOutlineStylePostProcessor(_StyleFromQmlPostProcessor):
    """Applies the red-outline / no-fill polygon style to the generated AOI."""
    STYLE_PATH = AOI_OUTLINE_STYLE_PATH
    STYLE_LABEL = 'AOI outline style'


def loadUtmLayer():
    """
    Load the packaged World UTM Grid shapefile as a QgsVectorLayer.

    Uses an absolute path derived from this file's location so callers don't
    have to chdir() into the plugin directory.
    """
    utm_path = path.join(path.dirname(__file__), 'resources', 'World_UTM_Grid.shp')
    return QgsVectorLayer(utm_path, 'World_UTM_Grid')


def downloadLandCoverRaster(i, utm_zones, year, scratch_folder):
    """
    Download a single land cover tile for the UTM zone at ``utm_zones[i]``
    for the given ``year`` and write it to ``scratch_folder``.

    Returns the filename of the saved land cover raster. Raises
    ``QgsProcessingException`` on network failure so the calling algorithm can
    surface a clear error via the Processing feedback panel instead of
    continuing with a corrupt file.
    """

    # Download data based on whichever UTM zone the aoi is within
    # Get UTM zone and create output filename
    zone = utm_zones[i]
    zone_id = str(zone[0]) + zone[1]
    output_filename = path.join(scratch_folder, zone_id + ".tif")

    # Get date range for URL
    if year == "2024":
        date_range = "20240101-20241231"
    else:
        year_int = int(year)
        next_year = year_int + 1
        date_range = year + "0101-" + str(next_year) + "0101"

    # Download data from Living Atlas
    url = ("https://lulctimeseries.blob.core.windows.net/lulctimeseriesv003/lc" + year + "/" + zone_id +
           "_" + date_range + ".tif")
    qurl = QUrl(url)
    request = QNetworkRequest(qurl)
    network_access_manager = QgsNetworkAccessManager()
    response = network_access_manager.blockingGet(request)

    # Surface network errors instead of writing garbage to disk
    if response.error() != QNetworkReply.NoError:
        raise QgsProcessingException(
            "Failed to download land cover tile " + zone_id + " for " + year +
            " (" + response.errorString() + "). URL: " + url)

    content = response.content()

    # Write to scratch folder using a context manager so the writer is closed
    # deterministically (fixes ResourceWarning: unclosed file).
    with open(output_filename, "wb") as f:
        f.write(content)
    temp_lc = open(output_filename, )
    temp_lc.close()

    # Return filename of downloaded land cover raster
    return output_filename


def _selectUtmZonesForAoi(aoi_layer, utm_layer, scratch_folder):
    """
    Select UTM grid features intersecting ``aoi_layer``, write the selection to
    ``scratch_folder`` as an ESRI Shapefile, then return the zones as a list of
    ``[zone_number, zone_letter]`` pairs.
    """

    processing.run("native:selectbylocation",
                   {'INPUT': utm_layer, 'PREDICATE': [0], 'INTERSECT': aoi_layer,
                    'METHOD': 0})

    selected_utm_path = path.join(scratch_folder, 'selected_utm.shp')

    # Use writeAsVectorFormatV3 (V1 was deprecated in QGIS 3.10 and V2 in 3.20).
    save_options = QgsVectorFileWriter.SaveVectorOptions()
    save_options.driverName = 'ESRI Shapefile'
    save_options.fileEncoding = 'utf-8'
    save_options.onlySelectedFeatures = True
    QgsVectorFileWriter.writeAsVectorFormatV3(
        utm_layer,
        selected_utm_path,
        QgsProject.instance().transformContext(),
        save_options)

    temp_selection = QgsVectorLayer(selected_utm_path, 'temp_selection')
    return [[feature[1], feature[2]] for feature in temp_selection.getFeatures()]


def runLandCoverPipeline(aoi_layer, year_string, scratch_folder,
                         output_lc_raster, feedback):
    """
    Shared download / clip / mosaic pipeline used by all three algorithms.

    Given an AOI layer already in the model coordinate system (EPSG:3857),
    determine which Living Atlas tiles it intersects, download each tile into
    ``scratch_folder``, then clip and mosaic them into ``output_lc_raster``.

    Honors cancellation via ``feedback.isCanceled()`` between phases.
    """

    if feedback.isCanceled():
        return
    feedback.setProgressText('Downloading Land Cover Data...')

    utm_layer = loadUtmLayer()
    utm_zones = _selectUtmZonesForAoi(aoi_layer, utm_layer, scratch_folder)

    lc_return = []
    for i in range(len(utm_zones)):
        if feedback.isCanceled():
            return
        feedback.setProgressText(
            'Downloading LULC Data ' + str(i + 1) + '/' + str(len(utm_zones)) + '...')
        lc_return.append(downloadLandCoverRaster(
            i, utm_zones, year_string, scratch_folder))

    if feedback.isCanceled():
        return

    feedback.setProgressText('Clipping and Merging Land Cover Data...')
    mosaicAndClipRasters(lc_return, aoi_layer, scratch_folder,
                         output_lc_raster, feedback)


def _writeAoiCutline(aoi_layer, scratch_folder):
    """
    Write ``aoi_layer`` to an ESRI Shapefile in ``scratch_folder`` for use as
    a ``gdalwarp`` cutline. Needed because gdal.Warp requires a data-source
    path — a memory layer or already-loaded QgsVectorLayer won't do.
    Returns the cutline shapefile path.
    """
    cutline_path = path.join(scratch_folder, 'aoi_cutline.shp')
    save_options = QgsVectorFileWriter.SaveVectorOptions()
    save_options.driverName = 'ESRI Shapefile'
    save_options.fileEncoding = 'utf-8'
    QgsVectorFileWriter.writeAsVectorFormatV3(
        aoi_layer,
        cutline_path,
        QgsProject.instance().transformContext(),
        save_options)
    return cutline_path


def _makeWarpProgressCallback(feedback):
    """
    Build a GDAL warp progress callback that forwards to a Processing
    ``QgsProcessingFeedback`` and honors cancellation. Returning 0 from a
    GDAL callback aborts the warp.
    """
    def _cb(complete, message, unknown):
        if feedback is None:
            return 1
        if feedback.isCanceled():
            return 0
        feedback.setProgress(int(complete * 100))
        return 1
    return _cb


def mosaicAndClipRasters(input_lc_filenames, aoi_layer, scratch_folder,
                         output_lc_raster, feedback=None):
    """
    Mosaic, reproject, and clip all downloaded land cover tiles into a single
    output raster using ``osgeo.gdal.Warp`` directly.

    Why gdal.Warp instead of gdal:cliprasterbymasklayer + gdal:merge:

    - **Memory.** ``gdal_merge.py`` (which backs ``gdal:merge``) allocates the
      full output as an int64 numpy array. On a multi-UTM-zone AOI that can
      run into tens of GiB and abort with ``ArrayMemoryError`` (see
      https://github.com/CustomCartographix/LandCoverDownloader/issues/1).
      ``gdalwarp`` streams block-by-block, so peak memory is a handful of MB
      regardless of output extent.
    - **Multi-CRS input.** Living Atlas tiles are stored in their native UTM
      projections, which differ across zones. ``gdal.Warp`` accepts multiple
      inputs in different CRSs and reprojects each to ``dstSRS`` in a single
      pass. ``gdal:buildvirtualraster`` cannot.
    - **Color table.** ``gdal.Warp`` preserves the source palette on Byte
      rasters, so the paletted classes come through even without the QML
      style being applied on top.
    - **Fewer steps, fewer temp files.** No per-tile clipped intermediate.

    Categorical class values require nearest-neighbour resampling; bilinear or
    cubic would invent nonsense classes (e.g. 2.7).
    """
    # gdal is a QGIS dependency; deferring the import keeps plugin startup
    # robust if the bindings are ever unavailable at load time.
    from osgeo import gdal

    cutline_path = _writeAoiCutline(aoi_layer, scratch_folder)

    warp_options = gdal.WarpOptions(
        format='GTiff',
        dstSRS=aoi_layer.crs().authid(),
        cutlineDSName=cutline_path,
        cropToCutline=True,
        resampleAlg='near',       # categorical data — do NOT change
        multithread=False,
        creationOptions=['COMPRESS=LZW', 'TILED=YES', 'PREDICTOR=2'],
        callback=_makeWarpProgressCallback(feedback),
    )

    result = gdal.Warp(output_lc_raster, input_lc_filenames, options=warp_options)
    if result is None:
        raise QgsProcessingException(
            'gdal.Warp failed to mosaic/clip land cover tiles into ' +
            str(output_lc_raster))
    # Close the dataset so the file is flushed to disk before QGIS opens it
    result = None
