# -*- coding: utf-8 -*-
"""
create_river_line.py

〰 Create River Line from DEM

Extracts a river/channel network from a DEM using:
    Fill sinks XXL (Wang & Liu)
    → Flow accumulation (top-down, D8)
    → Threshold to binary channel raster
    → Convert to integer (CELL type)
    → r.thin (GRASS) — skeletonise to 1-pixel centreline
    → r.to.vect (GRASS) — vectorise to line

The output is intentionally "raw" — the user should inspect and
clean it before feeding it into the Generate REM tool.
"""

import os
import tempfile

from qgis.PyQt.QtCore import QCoreApplication
from qgis.core import (
    QgsProcessing,
    QgsProcessingAlgorithm,
    QgsProcessingException,
    QgsProcessingParameterRasterLayer,
    QgsProcessingParameterNumber,
    QgsProcessingParameterVectorDestination,
    QgsProcessingMultiStepFeedback,
    QgsVectorLayer,
)
import processing


class CreateRiverLineAlgorithm(QgsProcessingAlgorithm):

    INPUT_DEM            = "INPUT_DEM"
    FLOW_ACCUM_THRESHOLD = "FLOW_ACCUM_THRESHOLD"
    OUTPUT_RIVER_LINE    = "OUTPUT_RIVER_LINE"

    def tr(self, string):
        return QCoreApplication.translate("CreateRiverLineAlgorithm", string)

    def createInstance(self):
        return CreateRiverLineAlgorithm()

    def name(self):
        return "createriverline"

    def displayName(self):
        return self.tr("〰 Create River Line from DEM")

    def group(self):
        return self.tr("REM Tools")

    def groupId(self):
        return "remtools"

    def shortHelpString(self):
        return self.tr("""
        <h3>〰 Create River Line from DEM</h3>
        <p>Automatically extracts a river/channel network from a DEM using
        flow accumulation analysis.</p>

        <h4>Pipeline</h4>
        <ol>
          <li>Fill sinks XXL (Wang &amp; Liu) — prepares DEM for flow analysis</li>
          <li>Flow accumulation (D8 top-down)</li>
          <li>Threshold — keep only cells with ≥ N upstream cells</li>
          <li>Convert to integer raster (required by r.thin)</li>
          <li>r.thin — skeletonise to 1-pixel centreline</li>
          <li>r.to.vect — vectorise to line layer</li>
        </ol>

        <h4>Output</h4>
        <p>A raw river line layer. Inspect it, delete spurious branches,
        merge gaps, and simplify before using it in
        <b>🌄 Generate REM</b>.</p>

        <h4>Tips</h4>
        <ul>
          <li>Start with threshold = 10000 for a 30m DEM</li>
          <li>Higher threshold = fewer, larger rivers</li>
          <li>Lower threshold = denser network (noisier)</li>
          <li>Use QGIS editing tools to clean the output line</li>
        </ul>
        """)

    def initAlgorithm(self, config=None):

        self.addParameter(QgsProcessingParameterRasterLayer(
            self.INPUT_DEM,
            self.tr("Input DEM (projected, metre units)"),
        ))

        self.addParameter(QgsProcessingParameterNumber(
            self.FLOW_ACCUM_THRESHOLD,
            self.tr("Flow accumulation threshold (cells) — higher = fewer rivers"),
            type=QgsProcessingParameterNumber.Integer,
            defaultValue=10000, minValue=10,
        ))

        self.addParameter(QgsProcessingParameterVectorDestination(
            self.OUTPUT_RIVER_LINE,
            self.tr("Output river line (inspect and clean before using in REM)"),
            type=QgsProcessing.TypeVectorLine,
        ))

    def processAlgorithm(self, parameters, context, feedback):

        fb = QgsProcessingMultiStepFeedback(6, feedback)

        dem_layer = self.parameterAsRasterLayer(parameters, self.INPUT_DEM, context)
        if not dem_layer or not dem_layer.isValid():
            raise QgsProcessingException("Invalid DEM layer.")

        threshold = self.parameterAsInt(parameters, self.FLOW_ACCUM_THRESHOLD, context)
        dem_path  = dem_layer.source()

        if dem_layer.crs().isGeographic():
            fb.reportError(
                "WARNING: DEM appears to use a geographic CRS (degrees). "
                "A projected CRS (metres) is strongly recommended.", False
            )

        # ---- Step 1: Fill sinks XXL ------------------------------------
        fb.setCurrentStep(1)
        fb.pushInfo("Step 1/6 – Filling sinks (XXL Wang & Liu) …")

        filled = processing.run(
            "sagang:fillsinksxxlwangliu",
            {
                "ELEV":     dem_path,
                "MINSLOPE": 0.01,
                "FILLED":   "TEMPORARY_OUTPUT",
            },
            context=context, feedback=fb,
        )

        if fb.isCanceled():
            return {}

        # ---- Step 2: Flow accumulation ---------------------------------
        fb.setCurrentStep(2)
        fb.pushInfo("Step 2/6 – Computing flow accumulation (D8) …")

        try:
            flow = processing.run(
                "sagang:flowaccumulationtopdown",
                {
                    "ELEVATION":    filled["FILLED"],
                    "FLOW":         "TEMPORARY_OUTPUT",
                    "METHOD":       0,          # D8
                    "STEP":         1,
                    "FLOW_UNIT":    0,
                    "LINEAR_DO":    True,
                    "LINEAR_MIN":   threshold,
                    "CONVERGENCE":  1.1,
                    "MFD_CONTOUR":  False,
                    "NO_NEGATIVES": True,
                    "VAL_MEAN":     "TEMPORARY_OUTPUT",
                    "ACCU_TARGET":  filled["FILLED"],
                    "FLOW_LENGTH":  "TEMPORARY_OUTPUT",
                    "WEIGHT_LOSS":  "TEMPORARY_OUTPUT",
                },
                context=context, feedback=fb,
            )
        except Exception as e:
            raise QgsProcessingException(f"Flow accumulation failed: {str(e)}")

        if fb.isCanceled():
            return {}

        # ---- Step 3: Threshold -----------------------------------------
        fb.setCurrentStep(3)
        fb.pushInfo(f"Step 3/6 – Thresholding at {threshold} cells …")

        flow_name = os.path.splitext(os.path.basename(flow["FLOW"]))[0]
        channel = processing.run(
            "qgis:rastercalculator",
            {
                "EXPRESSION": f'"{flow_name}@1" >= {threshold}',
                "LAYERS":     [flow["FLOW"]],
                "CELLSIZE":   0,
                "EXTENT":     None,
                "CRS":        None,
                "OUTPUT":     "TEMPORARY_OUTPUT",
            },
            context=context, feedback=fb,
        )

        if fb.isCanceled():
            return {}

        # ---- Step 4: Convert to integer (required by r.thin) -----------
        fb.setCurrentStep(4)
        fb.pushInfo("Step 4/6 – Converting to integer raster …")

        channel_int = processing.run(
            "gdal:rastercalculator",
            {
                "INPUT_A": channel["OUTPUT"],
                "BAND_A":  1,
                "FORMULA": "A.astype(int)",
                "NO_DATA": 0,
                "RTYPE":   1,   # Int16
                "OUTPUT":  "TEMPORARY_OUTPUT",
            },
            context=context, feedback=fb,
        )

        if fb.isCanceled():
            return {}

        # ---- Step 5 & 6: Thin + vectorise ------------------------------
        fb.setCurrentStep(5)
        fb.pushInfo("Step 5/6 – Thinning to centreline (r.thin) …")

        output_path = self.parameterAsOutputLayer(
            parameters, self.OUTPUT_RIVER_LINE, context
        )

        try:
            thinned = processing.run(
                "grass7:r.thin",
                {
                    "input":      channel_int["OUTPUT"],
                    "iterations": 200,
                    "output":     "TEMPORARY_OUTPUT",
                    "GRASS_REGION_CELLSIZE_PARAMETER": 0,
                },
                context=context, feedback=fb,
            )

            fb.setCurrentStep(6)
            fb.pushInfo("Step 6/6 – Vectorising centreline (r.to.vect) …")

            vect = processing.run(
                "grass7:r.to.vect",
                {
                    "input":  thinned["output"],
                    "type":   0,        # line
                    "column": "value",
                    "output": output_path,
                    "GRASS_REGION_CELLSIZE_PARAMETER": 0,
                    "GRASS_OUTPUT_TYPE_PARAMETER":     0,
                },
                context=context, feedback=fb,
            )
            line_path = vect["output"]

        except Exception as e:
            fb.reportError(
                f"GRASS not available ({e}) — using GDAL polygonize fallback. "
                "Output may be less clean.", False
            )

            poly = processing.run(
                "gdal:polygonize",
                {
                    "INPUT":               channel_int["OUTPUT"],
                    "BAND":                1,
                    "FIELD":               "DN",
                    "EIGHT_CONNECTEDNESS": False,
                    "OUTPUT":              "TEMPORARY_OUTPUT",
                },
                context=context, feedback=fb,
            )
            filt = processing.run(
                "qgis:extractbyattribute",
                {
                    "INPUT":    poly["OUTPUT"],
                    "FIELD":    "DN",
                    "OPERATOR": 0,
                    "VALUE":    "1",
                    "OUTPUT":   "TEMPORARY_OUTPUT",
                },
                context=context, feedback=fb,
            )
            lines = processing.run(
                "qgis:polygonstolines",
                {"INPUT": filt["OUTPUT"], "OUTPUT": output_path},
                context=context, feedback=fb,
            )
            line_path = lines["OUTPUT"]

        # Report result
        result_lyr = QgsVectorLayer(line_path, "river_line", "ogr")
        n = result_lyr.featureCount()
        fb.pushInfo(f"✓ Extracted {n} river line feature(s).")

        if n > 500:
            fb.reportError(
                f"WARNING: {n} features extracted — the network is quite dense. "
                "Consider increasing the flow accumulation threshold and re-running, "
                "or manually delete unwanted branches before generating the REM.",
                fatalError=False,
            )

        return {self.OUTPUT_RIVER_LINE: line_path}
