# -*- coding: utf-8 -*-
"""
/***************************************************************************
 PandaPipesForQGISDockWidget
                                 A QGIS plugin
 An interface between the open source fluid modelling tool, Panda Pipes and QGIS.
 Generated by Plugin Builder: http://g-sherman.github.io/Qgis-Plugin-Builder/
                             -------------------
        begin                : 2025-05-26
        git sha              : $Format:%H$
        copyright            : (C) 2025 by Thomas Raj
        email                : lemungineer@gmail.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.                                   *
 *                                                                         *
 ***************************************************************************/
"""

# Standard library
import os

# Third-party libraries
import pandas as pd
import numpy as np
import pandapipes as pp
import processing

# PyQt / Qt imports
from PyQt5.QtCore import QVariant, QCoreApplication
from PyQt5.QtWidgets import (
    QHeaderView,
    QTreeWidgetItem,
    QDialog,
    QMessageBox,
    QListWidget
)

# QGIS PyQt imports
from qgis.PyQt import QtGui, QtWidgets, uic
from qgis.PyQt.QtCore import pyqtSignal, QAbstractTableModel, Qt

# QGIS core imports
from qgis.core import (
    QgsVectorLayer,
    QgsField,
    QgsFields,
    QgsFeature,
    QgsGeometry,
    QgsPointXY,
    QgsProject,
    QgsEditorWidgetSetup,
    QgsDefaultValue,
    QgsProcessingFeatureSourceDefinition,
    QgsFeatureRequest,
    QgsMapLayerProxyModel,
    QgsMapLayer
)

# QGIS GUI imports
from qgis.gui import QgsFileWidget

# QGIS utilities
from qgis.utils import iface

# Local imports
from .logic.power_component_fields import *
from .logic.function_helpers import *

FORM_CLASS, _ = uic.loadUiType(os.path.join(
    os.path.dirname(__file__), 'power_network_component_dialog_base.ui'))

class PowerComponentDialog(QDialog, FORM_CLASS):
    def __init__(self, parent=None):
        super(PowerComponentDialog, self).__init__(parent)
        self.setupUi(self)  # Load the UI components
    
        # Combo Boxes

        components = [
                        "Bus","Bus DC","Line","Line DC","Switch",
                        "Load","Load DC","Asymmetric Load","Motor",
                        "Static Generator","Asymmetric Static Generator","Generator","External Grid","Source DC",
                        "Transformer","Three Winding Transformer","Shunt","Impedance","Ward","Extended Ward",
                        "HV DC Link / DC Line","Measurement","Storage","Static Var Compensator (SVC)",
                        "Thyristor-Controlled Series Capacitor (TCSC)",
                        "Static Synchronous Compensator (SSC)",
                        "Voltage Source Converter (VSC)", "Stacked Voltage Source Converter (VSC Stacked)"
                      ]

        self.cmbComponentType.addItems(components)

        self.update_geometry()
        self.cmbComponentType.currentIndexChanged.connect(self.update_geometry)

        # QGIS Combo Boxes
        self.mlcbBaseLayer.setFilters(QgsMapLayerProxyModel.VectorLayer)
        self.mlcbJunctionLayer.setFilters(QgsMapLayerProxyModel.PointLayer)

        self.update_adapt_geom_types()
        self.cmbComponentType.currentIndexChanged.connect(self.update_adapt_geom_types)

        self.update_tree_fields_to_keep()

        # Dynamic Combo Boxes
        self.mlcbBaseLayer.layerChanged.connect(lambda layer, fc=self.fcbLayerIDField: fc.setLayer(layer if layer and layer.type() == QgsMapLayer.VectorLayer else None))

        # Immediately apply the logic using the current layer
        layer = self.mlcbBaseLayer.currentLayer()
        if layer and layer.type() == QgsMapLayer.VectorLayer:
            self.fcbLayerIDField.setLayer(layer)
        else:
            self.fcbLayerIDField.setLayer(None)

        # Tree Widgets
        self.update_tree_fields()

        self.twFields.setAlternatingRowColors(True)
        self.twFieldsToKeep.setAlternatingRowColors(True)

        # Updating Tree Widget
        self.cmbComponentType.currentIndexChanged.connect(self.update_tree_fields)

        # Updating Tree Widget
        current_layer = self.mlcbBaseLayer.currentLayer()

        # Buttons
        self.pbGenerateComponent.clicked.connect(self.on_create_component_clicked)
        self.pbAdaptExistingLayer.clicked.connect(self.on_adapt_component_clicked)

        if current_layer:
            self.update_tree_fields_to_keep

        if current_layer != None:
            self.mlcbBaseLayer.currentIndexChanged.connect(self.update_tree_fields_to_keep)

        # Buttons
        self.pbFieldsToggle.clicked.connect(self.toggle_selection_fields)
        self.pbKeepFieldsToggle.clicked.connect(self.toggle_selection_keep_fields)

    # Non-User Input Functions

    def update_tree_fields(self):
        self.twFields.clear()

        component_type = self.cmbComponentType.currentText()
        component_fields = power_component_fields_dict[component_type]

        for field in component_fields:
            item = QTreeWidgetItem([field])
            item.setCheckState(0, Qt.Unchecked)  # Column 0 will have a checkbox
            self.twFields.addTopLevelItem(item)

    def update_geometry(self):
        component_type = self.cmbComponentType.currentText()
        available_geometries = power_component_geometries_dict[component_type]

        self.cmbLayerGeometryType.clear()
        self.cmbLayerGeometryType.addItems(available_geometries)

    def update_tree_fields_to_keep(self):
        layer_name = self.mlcbBaseLayer.currentLayer()
        print('CHANGE')
        print(layer_name)

        if layer_name == None:
            return
        
        if type(layer_name) == QgsVectorLayer:
            layer = layer_name
        else:
            layer = QgsProject.instance().mapLayersByName(layer_name)[0]

        self.twFieldsToKeep.clear()

        if not layer or not layer.fields():
            return

        for field in layer.fields():
            item = QTreeWidgetItem([field.name(), field.typeName()])
            item.setCheckState(0, Qt.Unchecked)  # Column 0 will have a checkbox
            self.twFieldsToKeep.addTopLevelItem(item)

    # User input functions

    def toggle_selection_fields(self):
        root = self.twFields.invisibleRootItem()
                
        # Helper: recursively check if any item is checked
        def any_checked(item):
            if item.checkState(0) == Qt.Checked:
                return True
            for i in range(item.childCount()):
                if any_checked(item.child(i)):
                    return True
            return False
        
        # Helper: recursively set check state
        def set_check_state(item, state):
            item.setCheckState(0, state)
            for i in range(item.childCount()):
                set_check_state(item.child(i), state)

        # Decide if we are checking or unchecking
        should_uncheck = any(any_checked(root.child(i)) for i in range(root.childCount()))

        new_state = Qt.Unchecked if should_uncheck else Qt.Checked

        # Apply new state to all items
        for i in range(root.childCount()):
            set_check_state(root.child(i), new_state)

    def toggle_selection_keep_fields(self):
        root = self.twFieldsToKeep.invisibleRootItem()
                
        # Helper: recursively check if any item is checked
        def any_checked(item):
            if item.checkState(0) == Qt.Checked:
                return True
            for i in range(item.childCount()):
                if any_checked(item.child(i)):
                    return True
            return False
        
        # Helper: recursively set check state
        def set_check_state(item, state):
            item.setCheckState(0, state)
            for i in range(item.childCount()):
                set_check_state(item.child(i), state)

        # Decide if we are checking or unchecking
        should_uncheck = any(any_checked(root.child(i)) for i in range(root.childCount()))

        new_state = Qt.Unchecked if should_uncheck else Qt.Checked

        # Apply new state to all items
        for i in range(root.childCount()):
            set_check_state(root.child(i), new_state)

    def update_adapt_geom_types(self):
        component_type = self.cmbComponentType.currentText()

        if component_type in ["Line", "Line DC", "Transformer", "Three Winding Transformer", "HV DC Link / DC Line", "Impedance", "Thyristor-Controlled Series Capacitor (TCSC)", "Switch"]:
            self.mlcbBaseLayer.setFilters((QgsMapLayerProxyModel.LineLayer))
        else:
            self.mlcbBaseLayer.setFilters((QgsMapLayerProxyModel.PointLayer))

    def on_create_component_clicked(self):
        # Pulling list of checked fields
        fields = []

        for i in range(self.twFields.topLevelItemCount()):
            item = self.twFields.topLevelItem(i)
            if item.checkState(0) == Qt.Checked:
                field_name = item.text(0)  # Assuming the field name is in column 0
                fields.append(field_name)

        component = self.cmbComponentType.currentText()
        geom_type = self.cmbLayerGeometryType.currentText()

        # Creating Layer
        project = QgsProject.instance()
        project_crs = project.crs().authid()

        bus_layer_name = self.mlcbJunctionLayer.currentText() # Junction Layer

        if 'Bus' not in component:
            if type(bus_layer_name) == QgsVectorLayer:
                bus_layer_name = bus_layer_name
            else:
                bus_layer_name = QgsProject.instance().mapLayersByName(bus_layer_name)[0]

        def create_layer(geom_type,crs,name):
            return QgsVectorLayer(f"{geom_type}?crs={crs}",name,"memory")

        def add_fields(layer, field_list):
            provider = layer.dataProvider()
            provider.addAttributes(field_list)
            layer.updateFields()

        filtered_fields = []
        for fname in fields:
            # Find the QgsField with this name in field_types
            match = next((f for f in power_field_types if f.name() == fname), None)
            if match:
                filtered_fields.append(match)

        layer = create_layer(geom_type, project_crs, f'{component} Layer')
        add_fields(layer,filtered_fields)
        
        if component == 'Pipe' and 'length_km' in fields:
            idx = layer.fields().indexOf("length_km")
            layer.setDefaultValueDefinition(idx, QgsDefaultValue("round($length / 1000, 3)", True))

        layer.setName(f'{component} Layer')
        project.addMapLayer(layer)

        junction_fields = ["junction","controlled_junction", "flow_junction", "from_junction", "return_junction", "to_junction"]

        common_values = set(fields) & set(junction_fields)

        for value in common_values:
            setup_value_relation(layer=layer,field_name=value,junction_layer=junction_layer)


    def on_adapt_component_clicked(self):
        component_type = self.cmbComponentType.currentText()
        component_name_field = self.fcbLayerIDField.currentText()

        if type(self.mlcbBaseLayer.currentText()) == QgsVectorLayer:
            layer = self.mlcbBaseLayer.currentText()
        else:
            layer = QgsProject.instance().mapLayersByName(self.mlcbBaseLayer.currentText())[0]

        try:
            if type(self.mlcbJunctionLayer.currentText()) == QgsVectorLayer:
                junction_layer = self.mlcbJunctionLayer.currentText()
            else:
                junction_layer = QgsProject.instance().mapLayersByName(self.mlcbJunctionLayer.currentText())[0]
        except:
            pass

        try:
            fields.remove('name')
        except:
            pass

        # Pulling list of checked fields
        fields = []

        for i in range(self.twFields.topLevelItemCount()):
            item = self.twFields.topLevelItem(i)
            if item.checkState(0) == Qt.Checked:
                field_name = item.text(0)  # Assuming the field name is in column 0
                fields.append(field_name)

        # Pulling list of checked fields
        fields_to_keep = []

        for i in range(self.twFieldsToKeep.topLevelItemCount()):
            item = self.twFieldsToKeep.topLevelItem(i)
            if item.checkState(0) == Qt.Checked:
                field_name = item.text(0)  # Assuming the field name is in column 0
                fields_to_keep.append(field_name)

        if component_type == 'Junction':
            print('CHECK 1')
            layer = create_junction_layer_from_existing(layer=layer,fields_to_keep=fields_to_keep)

            fields_to_keep.extend(['name'])

        elif component_type == 'Pipe':
            print('CHECK 2')
            layer = create_pipe_layer_from_existing(layer=layer,fields_to_keep=fields_to_keep,
                                                    junction_layer=junction_layer,
                                                    component_name_field=component_name_field)
            
            fields_to_keep.extend(['name','from_junction','to_junction'])
            try:
                fields.remove('from_junction')
            except:
                pass

            try:
                fields.remove('to_junction')
            except:
                pass
            
        elif 'junction' in fields:
            print('CHECK 3')
            layer = create_component_with_junctions_from_existing(layer=layer,fields_to_keep=fields_to_keep,
                                                                  junction_layer=junction_layer,component_name_field=component_name_field)

            try:
                fields.remove('name')
            except:
                pass

            try:
                fields.remove('junction')
            except:
                pass
        else:
            print('CHECK 4')
            fields_to_keep.extend(['name'])

            layer = create_component_from_existing(layer=layer,component_name_field=component_name_field,
                                                   fields_to_keep=fields_to_keep)

        filtered_fields = []
        for fname in fields:
            # Find the QgsField with this name in field_types
            match = next((f for f in field_types if f.name() == fname), None)
            if match:
                filtered_fields.append(match)

        # 3. Add the fields to the memory layer
        provider = layer.dataProvider()
        provider.addAttributes(filtered_fields)

        # 4. Update the layer’s fields so QGIS recognizes them
        layer.updateFields()

        junction_fields = ["junction","controlled_junction", "flow_junction", "from_junction", "return_junction", "to_junction"]

        common_values = set(fields) & set(junction_fields)

        for value in common_values:
            setup_value_relation(layer=layer,field_name=value,junction_layer=junction_layer)

        layer.setName(f'{component_type} Layer')
        QgsProject.instance().addMapLayer(layer)
        

    

