# Source code for marche_a_lombre.mns_downloader

    # -*- coding: utf-8 -*-
    """
    Part of MarcheALOmbre QGIS Plugin
    Copyright (C) 2025 Yolanda Seifert
    Licensed under GPL v2+
    """
    import os
    import time
    import math
    import tempfile
    import shutil
    import xml.etree.ElementTree as ET

    from qgis.core import (
        QgsNetworkAccessManager,
        QgsRectangle,
        QgsCoordinateReferenceSystem,
        QgsCoordinateTransform
    )
    from qgis.PyQt.QtCore import QUrl, QCoreApplication, QEventLoop
    from qgis.PyQt.QtNetwork import QNetworkRequest, QNetworkReply
    from osgeo import gdal, osr

    from .geo_definitions import MANUAL_DEFS

    class MNSDownloader:

        BASE_URL = "https://data.geopf.fr/wms-r"
        CAPABILITIES_URL = f"{BASE_URL}?SERVICE=WMS&VERSION=1.3.0&REQUEST=GetCapabilities"
        TILE_SIZE_PX = 4000

        def __init__(self, crs, transform_context, feedback=None):
            """
            Initializes the downloader

            Args:
                crs (str): The epsg string of the target CRS
                transform_context (QgsCoordinateTransformContext): Context for coordinate transforms
                feedback (QgsProcessingFeedback, optional): Feedback object for logging
            """
            self.manager = QgsNetworkAccessManager.instance()
            self.crs = crs
            self.transform_context = transform_context
            self.feedback = feedback
            self._capabilities_xml_cache = None

        def log(self, message):
            if self.feedback:
                self.feedback.pushInfo(message)
            else:
                print(message)

        def _fetch_capabilities(self):
            """
            Fetches WMS Capabilities to find layers dynamically

            Returns:
                xml.etree.ElementTree.Element: root element of the parsed XML

            Raises:
                Exception: If the network request fails
            """
            if self._capabilities_xml_cache is not None:
                return self._capabilities_xml_cache

            self.log("Fetching WMS Capabilities...")
            request = QNetworkRequest(QUrl(self.CAPABILITIES_URL))
            request.setAttribute(QNetworkRequest.FollowRedirectsAttribute, True)

            reply = self.manager.get(request)

            loop = QEventLoop()
            reply.finished.connect(loop.quit)
            loop.exec_()

            if reply.error() != QNetworkReply.NoError:
                raise Exception(f"Capabilities failed: {reply.errorString()}")

            content = reply.readAll()
            self._capabilities_xml_cache = ET.fromstring(content)
            return self._capabilities_xml_cache

        def get_layer_candidates(self, wgs84_point, is_mns=True):
            """
            Parses capabilities to find the best layer for the location

            Args:
                wgs84_point (QgsPointXY): point to query in WGS84 coordinates
                is_mns (bool, optional): True -> searches for Surface Models (MNS)
                                         False -> searches for Terrain Models (MNT)

            Returns:
                list[dict]: A list of layer candidates containing 'name' and 'score'.
            """
            try:
                root = self._fetch_capabilities()
            except Exception as e:
                self.log(f"Error fetching capabilities: {e}")
                return []

            ns = {'wms': 'http://www.opengis.net/wms'}
            candidates = []

            if is_mns:
                search_type = "MNS"
                fallback_lidar_global = "IGNF_LIDAR-HD_MNS_ELEVATION.ELEVATIONGRIDCOVERAGE.WGS84G"
                fallback_highres = "ELEVATION.ELEVATIONGRIDCOVERAGE.HIGHRES.MNS"
            else:
                search_type = "MNT"
                fallback_lidar_global = "IGNF_LIDAR-HD_MNT_ELEVATION.ELEVATIONGRIDCOVERAGE.WGS84G"
                fallback_highres = "ELEVATION.ELEVATIONGRIDCOVERAGE.HIGHRES"

            for layer in root.findall('.//wms:Layer', ns):
                name_elem = layer.find('wms:Name', ns)
                if name_elem is None: continue
                name = name_elem.text

                if "SHADOW" in name: continue

                geo_bbox = layer.find('wms:EX_GeographicBoundingBox', ns)
                if geo_bbox is None: continue

                try:
                    w = float(geo_bbox.find('wms:westBoundLongitude', ns).text)
                    e = float(geo_bbox.find('wms:eastBoundLongitude', ns).text)
                    s = float(geo_bbox.find('wms:southBoundLatitude', ns).text)
                    n = float(geo_bbox.find('wms:northBoundLatitude', ns).text)

                    # Check if point is inside the layer
                    px, py = wgs84_point.x(), wgs84_point.y()
                    if not (w <= px <= e and s <= py <= n):
                        continue

                    score = 0
                    if search_type in name and "LIDAR-HD" in name and "WGS84G" not in name:
                        score = 1000
                    elif name == fallback_lidar_global:
                        score = 500
                    elif name == fallback_highres:
                        score = 100

                    if score == 0: continue

                    candidates.append({'name': name, 'score': score})
                except:
                    continue

            return sorted(candidates, key=lambda x: x['score'], reverse=True)

        def validate_raster_content(self, file_path):
            """
            Validates the downloaded file by reading it

            Args:
                file_path (str): Path to the file to check

            Returns:
                tuple[bool, str]: (isValid, statusMessage)
            """
            try:
                if os.path.getsize(file_path) < 1000:
                    return False, "File too small (< 1KB)"

                with open(file_path, 'rb') as f:
                    header = f.read(512)
                    if b"ServiceException" in header or b"<?xml" in header:
                        return False, "Contains WMS XML Error"

                gdal.PushErrorHandler('CPLQuietErrorHandler')
                ds = gdal.Open(file_path)

                if not ds:
                    gdal.PopErrorHandler()
                    return False, "GDAL could not open file"

                band = ds.GetRasterBand(1)
                try:
                    data = band.ReadAsArray()
                except Exception as e:
                    ds = None
                    gdal.PopErrorHandler()
                    return False, f"File Truncated/Corrupt: {str(e)}"

                if data is None:
                    ds = None
                    gdal.PopErrorHandler()
                    return False, "ReadAsArray returned None"

                mn = data.min()
                mx = data.max()
                ds = None
                gdal.PopErrorHandler()

                if mn <= -9000 and mx <= -9000:
                    return False, f"All NoData (Min:{mn} Max:{mx})"

                if mn == mx:
                     return False, f"Flat Data (Val:{mn})"

                return True, "Valid"

            except Exception as e:
                return False, f"Validation Exception: {str(e)}"

        def download_dual_quality_mns(self, trail_extent, high_res_path, low_res_path, trail_lat, input_crs, high_res=0.5, low_res=15.0):
            """Download two MNS, one high quality around the trail and one low resolution with a greater extent for longer shadows

            Args:
                trail_extent (QgsRectangle): Extent around the hiking trail
                high_res_path (str): Path to High-Res MNS
                low_res_path (str): Path to Low-Res MNS
                trail_lat (float): latitude of trail_extent center
                input_crs (str): Coordinate Reference System
                high_res (float, optional): High resolution. Defaults to 0.5.
                low_res (float, optional): Low resolution. Defaults to 30.0.

            Returns:
                bool: True if download successful
            """
            # High-Res
            self.read_tif(trail_extent, high_res, high_res_path, input_crs=input_crs)

            # Low-Res (greater extent)
            self.log("Downloading large Low Resolution MNS for obstacles at greater distance (e.g. mountains)")
            buffer_dist = 22000.0 # altitude difference of 2000m with solar elevation of 5° casts 22km shadow
            buffer_n = buffer_dist
            buffer_s = buffer_dist
            if trail_lat > 23.4: # north/south buffer not necessary for high/low lat
                buffer_n = 0
            if trail_lat < -23.4:
                buffer_s = 0

            horizon_extent = QgsRectangle(
                trail_extent.xMinimum() - buffer_dist,
                trail_extent.yMinimum() - buffer_s,
                trail_extent.xMaximum() + buffer_dist,
                trail_extent.yMaximum() + buffer_n
            )
            return self.read_tif(horizon_extent, low_res, low_res_path, input_crs=input_crs)

        def read_tif(self, extent, resolution, output_path, input_crs, is_mns=True):
            """
            Dowloads the MNS/MNT data for a specific extent and resolution

            Args:
                extent (QgsRectangle): The area to download
                resolution (float): Pixel resolution in meters
                output_path (str): File path to save the GeoTIFF
                input_crs (str): epsg code of the CRS
                is_mns (bool, optional): True -> MNS (Surface), False -> MNT (Terrain). Defaults to True.

            Returns:
                bool: True if successful, False otherwise
            """
            source_ref = QgsCoordinateReferenceSystem(input_crs)
            wgs84_ref = QgsCoordinateReferenceSystem("EPSG:4326")

            tr_to_wgs84 = QgsCoordinateTransform(source_ref, wgs84_ref, self.transform_context)
            tr_to_wgs84.setBallparkTransformsAreAppropriate(True)

            center_input = extent.center()
            center_wgs84 = tr_to_wgs84.transform(center_input)

            auth_id = source_ref.authid()
            is_identity = (abs(center_input.x() - center_wgs84.x()) < 0.1) and (auth_id != wgs84_ref.authid())

            if is_identity: # transformation failed
                if auth_id in MANUAL_DEFS:
                    self.log(f"Switching to Manual Definition for Coordinate Transformation.")
                    # CRS from manual definition
                    source_ref = QgsCoordinateReferenceSystem.fromProj4(MANUAL_DEFS[auth_id])
                    # Redo transform
                    tr_to_wgs84 = QgsCoordinateTransform(source_ref, wgs84_ref, self.transform_context)
                    tr_to_wgs84.setBallparkTransformsAreAppropriate(True)
                    center_wgs84 = tr_to_wgs84.transform(center_input)
                else:
                    self.log("Warning: Transform returned identity and no manual definition available.")

            # Search available layers using WGS84 center
            candidates = self.get_layer_candidates(center_wgs84, is_mns)

            if not candidates:
                # Default based on detected CRS to avoid "Metropole" layer in DOM-TOM
                if "2154" in self.crs:
                    default = "IGNF_LIDAR-HD_MNS_ELEVATION.ELEVATIONGRIDCOVERAGE.LAMB93"
                else:
                    # Use the generic WGS84G layer for DOM-TOM if specific one isn't found
                    default = "IGNF_LIDAR-HD_MNS_ELEVATION.ELEVATIONGRIDCOVERAGE.WGS84G"

                self.log(f"No candidates found. Using default: {default}")
                candidates = [{'name': default}]

            # Calculate size in pixels based on the projected extent
            width = int(extent.width() / resolution)
            height = int(extent.height() / resolution)

            for i, cand in enumerate(candidates):
                layer_name = cand['name']
                self.log(f"Attempting layer {layer_name}...")

                success = False
                # Check if the request is small enough for a single download
                if width <= self.TILE_SIZE_PX and height <= self.TILE_SIZE_PX:
                    success = self._download_single_tile(extent, width, height, output_path, layer_name)
                else:
                    # If too big switch to tiled download
                    self.log(f"Large request ({width}x{height}). Switching to tiled download...")
                    success = self._download_tiled(extent, resolution, width, height, output_path, layer_name)

                if success:
                    is_valid, msg = self.validate_raster_content(output_path)
                    if is_valid:
                        self.log(f"Success: {msg}")
                        return True
                    else:
                        self.log(f"Layer {layer_name} INVALID: {msg}")

            return False

        def _download_tiled(self, extent, resolution, total_w, total_h, output_path, layer_name):
            """
            Splits the extent into chunks of max 4000px, downloads them, and merges them
            """
            cols = math.ceil(total_w / self.TILE_SIZE_PX)
            rows = math.ceil(total_h / self.TILE_SIZE_PX)

            temp_dir = tempfile.mkdtemp()
            tile_files = []

            current_y = extent.yMaximum()

            try:
                for i in range(rows):
                    current_x = extent.xMinimum()

                    # Calculate height of this row (last row might be smaller)
                    row_height_px = min(self.TILE_SIZE_PX, total_h - (i * self.TILE_SIZE_PX))
                    row_height_m = row_height_px * resolution

                    for j in range(cols):
                        if self.feedback and self.feedback.isCanceled():
                            return False

                        # Calculate width of this col
                        col_width_px = min(self.TILE_SIZE_PX, total_w - (j * self.TILE_SIZE_PX))
                        col_width_m = col_width_px * resolution

                        # Define tile extent
                        tile_extent = QgsRectangle(
                            current_x,
                            current_y - row_height_m,
                            current_x + col_width_m,
                            current_y
                        )

                        tile_path = os.path.join(temp_dir, f"tile_{i}_{j}.tif")

                        self.log(f"Downloading tile {i+1},{j+1} / {rows},{cols} ({col_width_px}x{row_height_px})...")

                        success = self._download_single_tile(
                            tile_extent, col_width_px, row_height_px, tile_path, layer_name
                        )

                        if not success:
                            raise Exception("Tile download failed")

                        tile_files.append(tile_path)
                        current_x += col_width_m

                    current_y -= row_height_m

                # Merge tiles using GDAL VRT
                self.log("Merging tiles...")
                vrt_options = gdal.BuildVRTOptions(resampleAlg='nearest')
                vrt_path = os.path.join(temp_dir, "merged.vrt")
                vrt = gdal.BuildVRT(vrt_path, tile_files, options=vrt_options)
                vrt = None # Flush to disk

                # Translate VRT to final TIFF
                translate_options = gdal.TranslateOptions(format='GTiff', creationOptions=['COMPRESS=DEFLATE', 'TILED=YES'])
                gdal.Translate(output_path, vrt_path, options=translate_options)

                return True

            except Exception as e:
                self.log(f"Error during tiled download: {e}")
                return False
            finally:
                shutil.rmtree(temp_dir, ignore_errors=True)

        def _download_single_tile(self, extent, width, height, output_path, layer_name):
            """
            Requests a single tile from IGN Wep Map Service
            """
            params = [
                f"SERVICE=WMS",
                f"VERSION=1.3.0",
                f"REQUEST=GetMap",
                f"LAYERS={layer_name}",
                f"STYLES=normal",
                f"FORMAT=image/tiff",
                f"CRS={self.crs}",
                f"BBOX={extent.xMinimum()},{extent.yMinimum()},{extent.xMaximum()},{extent.yMaximum()}",
                f"WIDTH={int(width)}",
                f"HEIGHT={int(height)}",
                f"TRANSPARENT=false"
            ]

            full_url_str = f"{self.BASE_URL}?" + "&".join(params)
            # self.log(f"Requesting URL: {full_url_str}")

            request = QNetworkRequest(QUrl(full_url_str))
            request.setAttribute(QNetworkRequest.FollowRedirectsAttribute, True)
            reply = self.manager.get(request)

            while reply.isRunning():
                if self.feedback and self.feedback.isCanceled():
                    reply.abort()
                    self.feedback.reportError("Download canceled by user.")
                    return False
                QCoreApplication.processEvents()

            # Check HTTP Status
            if reply.error() != QNetworkReply.NoError:
                if self.feedback:
                    if reply.error() != QNetworkReply.OperationCanceledError:
                        self.feedback.reportError(f"HTTP Error: {reply.errorString()}")
                self.log(f"QNetworkReply Error: {reply.errorString()}")
                return False

            content = reply.readAll()

            # Write Content to Disk
            if not content or len(content) < 100:
                if self.feedback:
                    self.feedback.reportError(f"Download failed (File too small). Server returned: {bytes(content)}")
                return False

            try:
                with open(output_path, 'wb') as f:
                    f.write(content)

                # Validation Step
                is_valid, status_msg = self.validate_raster_content(output_path)
                if not is_valid:
                    self.log(f"Downloaded file validation failed: {status_msg}")
                    return False

                self._embed_georeferencing(output_path, extent, width, height)
                # self.log(f"Saved to {output_path}")
                return True

            except Exception as e:
                if self.feedback:
                    self.feedback.reportError(f"GDAL Error: {e}")
                return False

        def _embed_georeferencing(self, tif_path, extent, width, height):
            """
            Opens the existing TIFF using GDAL and injects spatial metadata (GeoTransform and Projection)
            """
            ds = gdal.Open(tif_path, 1)
            if ds is None:
                raise Exception("Could not open file with GDAL.")

            nodata_val = -9999.0
            band = ds.GetRasterBand(1)
            data = band.ReadAsArray()
            data[data < -1000] = nodata_val
            band.WriteArray(data)
            if band:
                band.SetNoDataValue(nodata_val)

            x_res = (extent.xMaximum() - extent.xMinimum()) / width
            y_res = (extent.yMaximum() - extent.yMinimum()) / height

            # GeoTransform list format:
            # [0] Top-Left X Coordinate
            # [1] W-E Pixel Resolution
            # [2] Rotation
            # [3] Top-Left Y Coordinate
            # [4] Rotation
            # [5] N-S Pixel Resolution (negative for north-up)
            geotransform = [
                extent.xMinimum(), x_res, 0,
                extent.yMaximum(), 0, -y_res
            ]
            ds.SetGeoTransform(geotransform)

            srs = osr.SpatialReference()
            srs.SetFromUserInput(self.crs)

            ds.SetProjection(srs.ExportToWkt())
            # Close the file
            band.FlushCache()
            ds = None