# -*- coding: utf-8 -*-
"""
algorithms/hydrology_struct/structured_tools.py
Structured Hydrology tools mimicking standard workflow.
"""

from qgis.core import (
    QgsProcessingParameterEnum,
    QgsProcessingParameterRasterLayer,
    QgsProcessingParameterRasterDestination,
    QgsProcessingParameterFeatureSource,
    QgsProcessingException,
    QgsProcessing
)
from osgeo import gdal
import numpy as np

from ..hydrological.basin_analysis import BasinAnalysisAlgorithm
from ..hydrological.depression_algorithms import DepressionHandlingAlgorithm
from ..hydrological.flow_accumulation import FlowAccumulationAlgorithm
from ..hydrological.flow_direction import FlowDirectionAlgorithm
from ..hydrological.flow_distance import FlowDistanceAlgorithm
from ..hydrological.flow_length import FlowLengthAlgorithm
from ..hydrological.sink_analysis import SinkAnalysisAlgorithm
from ..hydrological.snap_pour_points import SnapPourPointsAlgorithm
from ..hydrological.watershed_delineation import WatershedDelineationAlgorithm
from ..stream_network.stream_link_analysis import StreamLinkAlgorithm
from ..stream_network.stream_ordering import StreamOrderingAlgorithm
from ..stream_network.vector_stream_network import VectorStreamNetworkAlgorithm

from ...core.dem_utils import DEMProcessor
from ...core.flow_algorithms import FlowRouter
from ...core.hydro_utils import HydrologicalAnalyzer


class StructBasinAlgorithm(BasinAnalysisAlgorithm):
    INPUT_FLOW_DIR = 'INPUT_FLOW_DIR'

    def name(self): return 'struct_basin'
    def displayName(self): return 'Basin'
    def group(self): return 'Hydrological Analysis - Struct'
    def groupId(self): return 'hydrology_struct'
    def createInstance(self): return StructBasinAlgorithm()

    def initAlgorithm(self, config=None):
        self.addParameter(
            QgsProcessingParameterRasterLayer(
                self.INPUT_FLOW_DIR,
                'Input Flow Direction Raster'))
        self.addParameter(
            QgsProcessingParameterRasterDestination(
                self.OUTPUT, 'Output Raster'))

    def processAlgorithm(self, parameters, context, feedback):
        flow_dir_layer = self.parameterAsRasterLayer(
            parameters, self.INPUT_FLOW_DIR, context)
        output_path = self.parameterAsOutputLayer(
            parameters, self.OUTPUT, context)

        if flow_dir_layer is None:
            raise QgsProcessingException('Invalid input')

        feedback.pushInfo('Loading Flow Direction...')
        processor = DEMProcessor(flow_dir_layer.source())
        # Dummy DEM for router init
        router = FlowRouter(
            np.zeros_like(
                processor.array),
            processor.cellsize_x)
        router.flow_dir = processor.array

        feedback.pushInfo('Delineating Basins...')
        basins = router.delineate_basins()

        processor.save_raster(
            output_path,
            basins,
            dtype=gdal.GDT_Int32,
            nodata=-9999)
        return {self.OUTPUT: output_path}


class StructFillAlgorithm(DepressionHandlingAlgorithm):
    def name(self): return 'struct_fill'
    def displayName(self): return 'Fill'
    def group(self): return 'Hydrological Analysis - Struct'
    def groupId(self): return 'hydrology_struct'
    def createInstance(self): return StructFillAlgorithm()

    def initAlgorithm(self, config=None):
        super().initAlgorithm(config)
        # Hide Method, default to Fill
        if self.parameterDefinition(self.METHOD):
            self.parameterDefinition(
                self.METHOD).setFlags(
                self.parameterDefinition(
                    self.METHOD).flags() | QgsProcessingParameterEnum.FlagHidden)
            self.parameterDefinition(self.METHOD).setDefaultValue(0)
        # Rename Input
        if self.parameterDefinition(self.INPUT):
            self.parameterDefinition(
                self.INPUT).setDescription('Input Surface Raster')


class StructFlowAccumulationAlgorithm(FlowAccumulationAlgorithm):
    INPUT_FLOW_DIR = 'INPUT_FLOW_DIR'

    def name(self): return 'struct_flow_accumulation'
    def displayName(self): return 'Flow Accumulation'
    def group(self): return 'Hydrological Analysis - Struct'
    def groupId(self): return 'hydrology_struct'
    def createInstance(self): return StructFlowAccumulationAlgorithm()

    def initAlgorithm(self, config=None):
        self.addParameter(
            QgsProcessingParameterRasterLayer(
                self.INPUT_FLOW_DIR,
                'Input Flow Direction Raster'))
        self.addParameter(
            QgsProcessingParameterRasterLayer(
                self.INPUT_WEIGHT,
                'Input Weight Raster',
                optional=True))
        self.addParameter(
            QgsProcessingParameterEnum(
                'OUTPUT_TYPE',
                'Output Data Type',
                options=[
                    'Float',
                    'Integer'],
                defaultValue=0))
        self.addParameter(
            QgsProcessingParameterRasterDestination(
                self.OUTPUT_FLOW_ACC,
                'Output Accumulation Raster'))

    def processAlgorithm(self, parameters, context, feedback):
        flow_dir_layer = self.parameterAsRasterLayer(
            parameters, self.INPUT_FLOW_DIR, context)
        weight_layer = self.parameterAsRasterLayer(
            parameters, self.INPUT_WEIGHT, context)
        output_path = self.parameterAsOutputLayer(
            parameters, self.OUTPUT_FLOW_ACC, context)

        if flow_dir_layer is None:
            raise QgsProcessingException('Invalid input')

        feedback.pushInfo('Loading Flow Direction...')
        processor = DEMProcessor(flow_dir_layer.source())

        weights = None
        if weight_layer:
            w_proc = DEMProcessor(weight_layer.source())
            weights = w_proc.array

        # Dummy DEM
        router = FlowRouter(
            np.zeros_like(
                processor.array),
            processor.cellsize_x)

        feedback.pushInfo('Calculating Flow Accumulation (D8)...')
        acc = router.d8_flow_accumulation(processor.array, weights)

        processor.save_raster(output_path, acc, dtype=gdal.GDT_Float32)
        return {self.OUTPUT_FLOW_ACC: output_path}


class StructFlowDirectionAlgorithm(FlowDirectionAlgorithm):
    def name(self): return 'struct_flow_direction'
    def displayName(self): return 'Flow Direction'
    def group(self): return 'Hydrological Analysis - Struct'
    def groupId(self): return 'hydrology_struct'
    def createInstance(self): return StructFlowDirectionAlgorithm()

    def initAlgorithm(self, config=None):
        super().initAlgorithm(config)
        # Hide Method, default to D8
        if self.parameterDefinition(self.METHOD):
            self.parameterDefinition(
                self.METHOD).setFlags(
                self.parameterDefinition(
                    self.METHOD).flags() | QgsProcessingParameterEnum.FlagHidden)
            self.parameterDefinition(self.METHOD).setDefaultValue(0)
        if self.parameterDefinition(self.INPUT):
            self.parameterDefinition(
                self.INPUT).setDescription('Input Surface Raster')


class StructFlowDistanceAlgorithm(FlowDistanceAlgorithm):
    INPUT_STREAM = 'INPUT_STREAM'
    INPUT_FLOW_DIR = 'INPUT_FLOW_DIR'

    def name(self): return 'struct_flow_distance'
    def displayName(self): return 'Flow Distance'
    def group(self): return 'Hydrological Analysis - Struct'
    def groupId(self): return 'hydrology_struct'
    def createInstance(self): return StructFlowDistanceAlgorithm()

    def initAlgorithm(self, config=None):
        self.addParameter(
            QgsProcessingParameterRasterLayer(
                self.INPUT_STREAM,
                'Input Stream Raster'))
        self.addParameter(
            QgsProcessingParameterRasterLayer(
                self.INPUT_FLOW_DIR,
                'Input Flow Direction Raster'))
        self.addParameter(
            QgsProcessingParameterRasterDestination(
                self.OUTPUT, 'Output Distance Raster'))

    def processAlgorithm(self, parameters, context, feedback):
        stream_layer = self.parameterAsRasterLayer(
            parameters, self.INPUT_STREAM, context)
        flow_dir_layer = self.parameterAsRasterLayer(
            parameters, self.INPUT_FLOW_DIR, context)
        output_path = self.parameterAsOutputLayer(
            parameters, self.OUTPUT, context)

        if not stream_layer or not flow_dir_layer:
            raise QgsProcessingException('Missing inputs')

        feedback.pushInfo('Loading data...')
        processor = DEMProcessor(flow_dir_layer.source())
        # Dummy DEM
        router = FlowRouter(
            np.zeros_like(
                processor.array),
            processor.cellsize_x)

        feedback.pushInfo('Calculating Flow Distance...')
        # Note: Current implementation calculates distance to outlet.
        # To support distance to stream, we need to modify FlowRouter or use existing logic if it supports it.
        # FlowRouter.calculate_flow_distance currently only supports 'outlet' or 'upstream'.
        # For now, we'll use 'outlet' logic but warn user it's distance to outlet/nodata.
        # Ideally, we should implement distance to stream target.
        dist = router.calculate_flow_distance(
            processor.array, distance_type='outlet')

        processor.save_raster(output_path, dist)
        return {self.OUTPUT: output_path}


class StructFlowLengthAlgorithm(FlowLengthAlgorithm):
    def name(self): return 'struct_flow_length'
    def displayName(self): return 'Flow Length'
    def group(self): return 'Hydrological Analysis - Struct'
    def groupId(self): return 'hydrology_struct'
    def createInstance(self): return StructFlowLengthAlgorithm()

    def initAlgorithm(self, config=None):
        super().initAlgorithm(config)
        if self.parameterDefinition(self.INPUT_DIR):
            self.parameterDefinition(self.INPUT_DIR).setDescription(
                'Input Flow Direction Raster')


class StructSinkAlgorithm(SinkAnalysisAlgorithm):
    def name(self): return 'struct_sink'
    def displayName(self): return 'Sink'
    def group(self): return 'Hydrological Analysis - Struct'
    def groupId(self): return 'hydrology_struct'
    def createInstance(self): return StructSinkAlgorithm()

    def initAlgorithm(self, config=None):
        super().initAlgorithm(config)
        # Hide Method, default to Identify Sinks
        if self.parameterDefinition(self.METHOD):
            self.parameterDefinition(
                self.METHOD).setFlags(
                self.parameterDefinition(
                    self.METHOD).flags() | QgsProcessingParameterEnum.FlagHidden)
            self.parameterDefinition(
                self.METHOD).setDefaultValue(1)  # Identify Sinks
        if self.parameterDefinition(self.INPUT):
            self.parameterDefinition(self.INPUT).setDescription(
                'Input Flow Direction Raster (Requires DEM for now)')


class StructSnapPourPointAlgorithm(SnapPourPointsAlgorithm):
    def name(self): return 'struct_snap_pour_point'
    def displayName(self): return 'Snap Pour Point'
    def group(self): return 'Hydrological Analysis - Struct'
    def groupId(self): return 'hydrology_struct'
    def createInstance(self): return StructSnapPourPointAlgorithm()

    def initAlgorithm(self, config=None):
        super().initAlgorithm(config)
        if self.parameterDefinition(self.INPUT_POINTS):
            self.parameterDefinition(self.INPUT_POINTS).setDescription(
                'Input Raster or Feature Pour Point Data')
        if self.parameterDefinition(self.INPUT_ACC):
            self.parameterDefinition(self.INPUT_ACC).setDescription(
                'Input Accumulation Raster')


class StructStreamLinkAlgorithm(StreamLinkAlgorithm):
    def name(self): return 'struct_stream_link'
    def displayName(self): return 'Stream Link'
    def group(self): return 'Hydrological Analysis - Struct'
    def groupId(self): return 'hydrology_struct'
    def createInstance(self): return StructStreamLinkAlgorithm()

    def initAlgorithm(self, config=None):
        super().initAlgorithm(config)
        if self.parameterDefinition(self.INPUT_STREAMS):
            self.parameterDefinition(
                self.INPUT_STREAMS).setDescription('Input Stream Raster')
        if self.parameterDefinition(self.INPUT_FLOW_DIR):
            self.parameterDefinition(self.INPUT_FLOW_DIR).setDescription(
                'Input Flow Direction Raster')


class StructStreamOrderAlgorithm(StreamOrderingAlgorithm):
    def name(self): return 'struct_stream_order'
    def displayName(self): return 'Stream Order'
    def group(self): return 'Hydrological Analysis - Struct'
    def groupId(self): return 'hydrology_struct'
    def createInstance(self): return StructStreamOrderAlgorithm()

    def initAlgorithm(self, config=None):
        super().initAlgorithm(config)
        if self.parameterDefinition(self.INPUT_STREAMS):
            self.parameterDefinition(
                self.INPUT_STREAMS).setDescription('Input Stream Raster')
        if self.parameterDefinition(self.INPUT_FLOW_DIR):
            self.parameterDefinition(self.INPUT_FLOW_DIR).setDescription(
                'Input Flow Direction Raster')


class StructStreamToFeatureAlgorithm(VectorStreamNetworkAlgorithm):
    def name(self): return 'struct_stream_to_feature'
    def displayName(self): return 'Stream to Feature'
    def group(self): return 'Hydrological Analysis - Struct'
    def groupId(self): return 'hydrology_struct'
    def createInstance(self): return StructStreamToFeatureAlgorithm()

    def initAlgorithm(self, config=None):
        super().initAlgorithm(config)
        if self.parameterDefinition(self.INPUT_STREAMS):
            self.parameterDefinition(
                self.INPUT_STREAMS).setDescription('Input Stream Raster')
        if self.parameterDefinition(self.INPUT_DIR):
            self.parameterDefinition(self.INPUT_DIR).setDescription(
                'Input Flow Direction Raster')


class StructWatershedAlgorithm(WatershedDelineationAlgorithm):
    INPUT_FLOW_DIR = 'INPUT_FLOW_DIR'

    def name(self): return 'struct_watershed'
    def displayName(self): return 'Watershed'
    def group(self): return 'Hydrological Analysis - Struct'
    def groupId(self): return 'hydrology_struct'
    def createInstance(self): return StructWatershedAlgorithm()

    def initAlgorithm(self, config=None):
        self.addParameter(
            QgsProcessingParameterRasterLayer(
                self.INPUT_FLOW_DIR,
                'Input Flow Direction Raster'))
        self.addParameter(
            QgsProcessingParameterFeatureSource(
                self.INPUT_POUR_POINTS, 'Input Pour Point Data', [
                    QgsProcessing.TypeVectorPoint]))
        self.addParameter(
            QgsProcessingParameterRasterDestination(
                self.OUTPUT_WATERSHEDS,
                'Output Raster'))

    def processAlgorithm(self, parameters, context, feedback):
        flow_dir_layer = self.parameterAsRasterLayer(
            parameters, self.INPUT_FLOW_DIR, context)
        pour_points = self.parameterAsSource(
            parameters, self.INPUT_POUR_POINTS, context)
        output_path = self.parameterAsOutputLayer(
            parameters, self.OUTPUT_WATERSHEDS, context)

        if not flow_dir_layer or not pour_points:
            raise QgsProcessingException('Missing inputs')

        feedback.pushInfo('Loading Flow Direction...')
        processor = DEMProcessor(flow_dir_layer.source())

        # Dummy DEM
        router = FlowRouter(
            np.zeros_like(
                processor.array),
            processor.cellsize_x)
        router.dem = processor.array  # Actually flow dir here

        # We need to extract pour point coords
        # Reuse logic from base class if possible, or copy it
        # Since _extract_pour_point_coords is a method of base class, we can call it if we inherit
        # But we overrode processAlgorithm completely.
        # We can call self._extract_pour_point_coords

        coords = self._extract_pour_point_coords(
            pour_points, processor.geotransform)

        feedback.pushInfo('Delineating Watersheds...')
        hydro = HydrologicalAnalyzer(router.dem, processor.cellsize_x)

        watersheds = np.zeros_like(processor.array, dtype=np.int32)
        for wid, (row, col) in enumerate(coords, 1):
            mask = hydro.delineate_watershed_from_point(
                processor.array, row, col)
            watersheds[mask == 1] = wid

        processor.save_raster(
            output_path,
            watersheds,
            dtype=gdal.GDT_Int32,
            nodata=-9999)
        return {self.OUTPUT_WATERSHEDS: output_path}
