import json
import sqlite3
from typing import Optional

from qgis.core import QgsApplication, QgsFeature, QgsVectorLayer
from qgis.PyQt.QtCore import QMetaType, QSize, QVariant, pyqtSignal
from qgis.PyQt.QtGui import QStandardItem, QStandardItemModel
from qgis.PyQt.QtWidgets import QCheckBox, QComboBox, QDialog, QGridLayout, QGroupBox, QHBoxLayout, QLabel, QLineEdit, QMessageBox, QPushButton, QRadioButton, QTabWidget, QTextEdit, QTreeView, QVBoxLayout, QWidget

from ..compat import ALIGN_TOP, CHECKED, DLG_ACCEPTED, ITEM_FLAG_EDITABLE, ITEM_FLAG_SELECTABLE, MESSAGE_LEVEL_CRITICAL, MESSAGE_LEVEL_SUCCESS, MESSAGE_LEVEL_WARNING, USER_ROLE
from ..gp.feature_class_functions import layer_path_parser
from ..gp.import_feature_class import ImportFeatureClass, ImportFieldMap
from ..gp.import_temp_layer import ImportMapLayer
from ..gp.order_by_line_task import OrderByLineTask
from ..gp.sample_frame_task import SampleFrameTask
from ..model.db_item import DBItem, DBItemModel
from ..model.db_item_spatial import DBItemSpatial
from ..model.pour_point import PourPoint
from ..model.project import Project
from ..model.sample_frame import SampleFrame, insert_sample_frame
from ..model.scratch_vector import ScratchVector
from ..QRiS.settings import Settings
from .frm_new_attribute import FrmNewAttribute
from .utilities import add_standard_form_buttons, validate_name
from .widgets.metadata import MetadataWidget
from .widgets.stats_widget import StatsWidget

# Text constants
flow_path = "Flow Path"  # level path
attributes_name = "Categories"
default_flow_path_name = ""
flow_path_field_name = "flow_path"
flows_into_field_name = "flows_into"
label_field_name = "display_label"


class FrmSampleFrame(QDialog):
    complete = pyqtSignal(bool)

    def __init__(self, parent, qris_project: Project, import_source_path: Optional[str] = None, sample_frame: SampleFrame = None, create_sample_frame: bool = False):

        self.qris_project = qris_project
        self.sample_frame = sample_frame
        self.import_source_path = import_source_path
        self.create_sample_frame = create_sample_frame

        super().__init__(parent)
        self.setWindowTitle("Create New (Empty) Sample Frame")

        # Figure out what tabs we need
        self.tab_inputs = None
        self.tab_attributes = None  # just in case
        self.attribute_filter = None

        if self.import_source_path is not None:
            if isinstance(self.import_source_path, QgsVectorLayer):
                self.tab_inputs = SampleFrameInputs(self, self.qris_project, self.import_source_path)
                self.tab_attributes = SampleFrameAttributesAddFields(self, self.import_source_path)
                self.setWindowTitle("Import New Sample Frame From Temporary Layer")
            else:
                self.tab_inputs = SampleFrameInputs(self, self.qris_project, self.import_source_path)
                self.tab_attributes = SampleFrameAttributesAddFields(self, self.import_source_path)
                self.setWindowTitle("Import New Sample Frame From Feature Class")
        else:
            self.tab_attributes = SampleFrameAttributes(self, self.sample_frame)

        if create_sample_frame is True:
            self.tab_inputs = SampleFrameInputsCreate(self, self.qris_project)
            self.setWindowTitle("Create New Sample Frame from Existing Layers")

        self.tab_properties = SampleFrameProperties(self, self.sample_frame)
        self.tab_description = SampleFrameDescription(self, self.sample_frame)

        metadata_json = json.dumps(sample_frame.user_metadata) if sample_frame is not None else None
        self.metadata_widget = MetadataWidget(self, metadata_json)
        self.stats_widget = StatsWidget(db_item=sample_frame, db_path=qris_project.project_file, parent=self)

        self.setupUi()
        if sample_frame is not None:
            self.setWindowTitle("Sample Frame Properties")
            self.txtName.setText(self.sample_frame.name)
            self.chkAddToMap.setCheckState(CHECKED)
            self.chkAddToMap.setVisible(False)
            self.chkStartEditSession.setVisible(False)
        else:
            self.txtName.setFocus()

    def set_inputs(self, cross_sections=None, polygon=None):

        if cross_sections is not None:
            self.tab_inputs.cboCrossSections.setCurrentIndex(self.tab_inputs.cross_sections_model.getItemIndex(cross_sections))
        if polygon is not None:
            self.tab_inputs.cboFramePolygon.setCurrentIndex(self.tab_inputs.polygons_model.getItemIndex(polygon))

    def promote_to_sample_frame(self, db_item: DBItem):

        self.txtName.setText(db_item.name)
        self.setWindowTitle(f"Promote {db_item.name} to Sample Frame")

        db_path = self.qris_project.project_file
        id_field = None
        if isinstance(db_item, PourPoint):
            layer_name = "catchments"
            id_field = "pour_point_id"
        elif isinstance(db_item, ScratchVector):
            layer_name = db_item.fc_name
            db_path = db_item.gpkg_path
        else:
            layer_name = db_item.db_table_name
            id_field = db_item.id_column_name
        self.import_source_path = f"{db_path}|layername={layer_name}"
        self.attribute_filter = f"{id_field} = {db_item.id}" if id_field is not None else None

        self.tab_inputs = SampleFrameInputs(self, self.qris_project, self.import_source_path)
        self.tabs.addTab(self.tab_inputs, "Inputs")
        self.basepath, self.layer_name, self.layer_id = layer_path_parser(self.import_source_path)

        metadata_json = json.dumps(db_item.metadata) if db_item.metadata is not None else None
        self.metadata_widget.load_json(metadata_json)

    def accept(self):

        # validate name
        if not validate_name(self, self.txtName):
            return

        # validate name is unique
        with sqlite3.connect(self.qris_project.project_file) as conn:
            curs = conn.cursor()
            curs.execute("SELECT name FROM sample_frames WHERE name = ? and sample_frame_type_id = 1 and not id = ?", [self.txtName.text(), self.sample_frame.id if self.sample_frame is not None else 0])
            row = curs.fetchone()
            if row is not None:
                QMessageBox.warning(self, "Duplicate Name", f"A sample frame with the name '{self.txtName.text()}' already exists. Please choose a unique name.")
                return

        metadata_json = self.metadata_widget.get_json()
        metadata = json.loads(metadata_json) if metadata_json is not None else None
        fields = self.tab_attributes.get_fields() if self.tab_attributes is not None else None

        out_metadata = {}
        if fields is not None:
            out_metadata["fields"] = fields
        if metadata is not None:
            out_metadata["metadata"] = metadata
        # add other items to the metadata
        out_metadata["default_flow_path_name"] = self.tab_properties.txtDefaultFlowPathName.text()

        try:
            if self.sample_frame is not None:
                self.sample_frame.update(self.qris_project.project_file, self.txtName.text(), self.tab_description.txtDescription.toPlainText(), out_metadata)
                self.qris_project.project_changed.emit()
            else:
                self.sample_frame = insert_sample_frame(self.qris_project.project_file, self.txtName.text(), self.tab_description.txtDescription.toPlainText(), out_metadata)
                self.qris_project.add_db_item(self.sample_frame)
        except Exception as ex:
            if "unique" in str(ex).lower():
                QMessageBox.warning(self, "Duplicate Name", f"A sample frame with the name '{self.txtName.text()}' already exists. Please choose a unique name.")
                self.txtName.setFocus()
            else:
                QMessageBox.warning(self, "Error Saving sample frame", str(ex))
            return

        if self.import_source_path is not None:
            try:
                # set the ok and cancel buttons to disabled
                self.buttonBox.setEnabled(False)

                sample_frame_path = f"{self.qris_project.project_file}|layername=sample_frame_features"
                out_field_map = self.get_field_map()
                if self.tab_inputs is not None:
                    if isinstance(self.tab_inputs, SampleFrameInputs):
                        if self.tab_inputs.cboDisplayLabel.currentIndex() > 0:
                            out_field_map.append(ImportFieldMap(self.tab_inputs.cboDisplayLabel.currentText(), "display_label", direct_copy=True))
                        if self.tab_inputs.cboFlowPathField.currentIndex() > 0:
                            out_field_map.append(ImportFieldMap(self.tab_inputs.cboFlowPathField.currentText(), "flow_path", direct_copy=True))
                        if self.tab_inputs.cboTopologyField.currentIndex() > 0:
                            out_field_map.append(ImportFieldMap(self.tab_inputs.cboTopologyField.currentText(), "source_topology"))
                            if self.tab_inputs.rdoFlowsInto.isChecked():
                                # Also capture the source feature's own fid so we can build the
                                # source_fid → dest_fid lookup in post-processing
                                out_field_map.append(ImportFieldMap("fid", "source_fid"))

                clip_mask = None
                if self.tab_inputs is not None:
                    clip_item = self.tab_inputs.cboClipToAOI.currentData(USER_ROLE)
                    if clip_item is not None:
                        if clip_item.id > 0:
                            clip_mask = ("sample_frame_features", "sample_frame_id", clip_item.id)

                attributes = {}
                attributes["sample_frame_id"] = self.sample_frame.id

                if isinstance(self.import_source_path, QgsVectorLayer):
                    import_mask_task = ImportMapLayer(self.import_source_path, sample_frame_path, attributes, out_field_map, clip_mask, self.attribute_filter, self.qris_project.project_file)
                else:
                    import_mask_task = ImportFeatureClass(self.import_source_path, sample_frame_path, attributes, out_field_map, clip_mask, attribute_filter=self.attribute_filter, proj_gpkg=self.qris_project.project_file)
                # # DEBUG
                # result = import_mask_task.run()
                # self.on_import_complete(result)
                # PRODUCTION
                import_mask_task.import_complete.connect(self.on_import_complete)
                QgsApplication.taskManager().addTask(import_mask_task)
                return
            except Exception as ex:
                try:
                    self.sample_frame.delete(self.qris_project.project_file)
                except Exception as ex_delete:
                    QMessageBox.warning(self, "Error Importing Sample Frame Features", str(ex_delete))
                    # enable the buttons
                    self.buttonBox.setEnabled(True)
                Settings().msg_bar("Error Importing Sample Frame Features", str(ex), MESSAGE_LEVEL_CRITICAL)
                return

        if self.create_sample_frame is True:
            try:
                db_item_polygon = self.tab_inputs.cboFramePolygon.currentData(USER_ROLE)
                if isinstance(db_item_polygon, DBItemSpatial):
                    polygon_layer = QgsVectorLayer(f"{self.qris_project.project_file}|layername={db_item_polygon.fc_name}")
                    polygon_layer.setSubsetString(f"{db_item_polygon.fc_id_column_name} = {db_item_polygon.id}")
                else:
                    polygon_layer = QgsVectorLayer(f"{db_item_polygon.gpkg_path}|layername={db_item_polygon.fc_name}")
                cross_sections = self.tab_inputs.cboCrossSections.currentData(USER_ROLE)
                cross_sections_layer = QgsVectorLayer(f"{self.qris_project.project_file}|layername=cross_section_features")
                cross_sections_layer.setSubsetString(f"cross_section_id = {cross_sections.id}")
                self.centerline_geom = None
                centerline_profile = self.tab_inputs.cboCenterline.currentData(USER_ROLE)
                if centerline_profile is not None:
                    cl_layer = QgsVectorLayer(f"{self.qris_project.project_file}|layername={centerline_profile.fc_name}")
                    cl_layer.setSubsetString(f"profile_id = {centerline_profile.id}")
                    cl_feat = QgsFeature()
                    if cl_layer.getFeatures().nextFeature(cl_feat):
                        self.centerline_geom = cl_feat.geometry()
                out_path = f"{self.qris_project.project_file}|layername=sample_frame_features"
                task = SampleFrameTask(polygon_layer, cross_sections_layer, out_path, self.sample_frame.id)
                task.sample_frame_complete.connect(self.on_split_complete)
                QgsApplication.taskManager().addTask(task)
                return
            except Exception as ex:
                try:
                    self.sample_frame.delete(self.qris_project.project_file)
                    Settings().log(f"Error Importing Sample Frame Features: {ex!s}", MESSAGE_LEVEL_CRITICAL)
                except Exception as ex:
                    Settings().log(f"Error Deleting sample frame: {ex!s}", MESSAGE_LEVEL_CRITICAL)
                return

        self.complete.emit(True)
        super().accept()
        return

    def get_field_map(self):

        field_map = []
        # if self.tab_inputs is not None:
        #     for field in self.tab_inputs.load_fields():
        #         field_map.append(ImportFieldMap(field, field, True))
        if self.tab_attributes is not None:
            for field_dict in self.tab_attributes.get_fields():
                field = field_dict["label"]
                field_map.append(ImportFieldMap(field, field, parent="attributes"))
        return field_map

    def on_split_complete(self, result: bool):
        if not result:
            self.on_import_complete(False)
            return
        if self.centerline_geom is not None:
            order_task = OrderByLineTask(
                f"{self.qris_project.project_file}|layername=sample_frame_features", self.centerline_geom, filter_expression=f"sample_frame_id = {self.sample_frame.id}", label_field="display_label", chain_field="flows_into"
            )
            order_task.order_complete.connect(self.on_import_complete)
            QgsApplication.taskManager().addTask(order_task)
        else:
            self.on_import_complete(True)

    def on_import_complete(self, result):
        if result is True:
            if self.tab_inputs is not None and isinstance(self.tab_inputs, SampleFrameInputs):
                if self.tab_inputs.cboTopologyField.currentIndex() > 0:
                    self._resolve_topology()
            Settings().msg_bar("Sample Frame Imported", f'Sample Frame "{self.txtName.text()}" has been created successfully.', MESSAGE_LEVEL_SUCCESS)
            self.complete.emit(True)
            super().accept()
        else:
            Settings().msg_bar("Error Importing Sample Frame Features", f'Sample Frame "{self.txtName.text()}" could not be imported.', MESSAGE_LEVEL_CRITICAL)
            try:
                self.sample_frame.delete(self.qris_project.project_file)
            except Exception as ex:
                Settings().msg_bar("Error Deleting sample frame", str(ex), MESSAGE_LEVEL_CRITICAL)
                super().accept()
            return

    def _resolve_topology(self):
        """Post-process: resolve source_topology values stored in metadata to flows_into fids.

        Sequence mode   : source_topology is an integer sequence number. Features are sorted
                          by that value and wired into a downstream chain via flows_into.
        Flows Into mode : source_topology is the source fid of the downstream feature.
                          The chain order is derived by graph traversal, then wired the same way.
        """
        is_sequence = self.tab_inputs.rdoSequence.isChecked()
        try:
            with sqlite3.connect(self.qris_project.project_file) as conn:
                curs = conn.cursor()
                curs.execute("SELECT fid, metadata FROM sample_frame_features WHERE sample_frame_id = ?", [self.sample_frame.id])
                rows = curs.fetchall()

                # Parse metadata for all features
                feature_meta = {}
                for fid, metadata_json in rows:
                    meta = json.loads(metadata_json) if metadata_json else {}
                    feature_meta[fid] = meta

                if is_sequence:
                    # Order directly from integer sequence values
                    ordered = sorted(((fid, int(meta["source_topology"])) for fid, meta in feature_meta.items() if meta.get("source_topology") is not None), key=lambda x: x[1])
                else:
                    # Derive order by traversing the source topology chain
                    # Build source_fid → dest_fid lookup
                    source_fid_to_dest = {str(meta["source_fid"]): fid for fid, meta in feature_meta.items() if meta.get("source_fid") is not None}
                    # Build source_fid → downstream source_fid map
                    src_to_next = {str(meta["source_fid"]): str(meta["source_topology"]) for meta in feature_meta.values() if meta.get("source_fid") is not None and meta.get("source_topology") is not None}
                    # Find the root: the source_fid not referenced as a downstream by anyone
                    all_source_fids = set(source_fid_to_dest.keys())
                    all_downstream = set(src_to_next.values())
                    roots = all_source_fids - all_downstream
                    if len(roots) != 1:
                        Settings().msg_bar("Topology Warning", "Topology chain is not linear — flows_into not set.", MESSAGE_LEVEL_WARNING)
                        return

                    # Traverse the chain to get ordered dest fids
                    ordered = []
                    current = next(iter(roots))
                    visited = set()
                    while current is not None:
                        if current in visited:
                            Settings().msg_bar("Topology Warning", "Cycle detected in topology chain — flows_into not set.", MESSAGE_LEVEL_WARNING)
                            return
                        visited.add(current)
                        dest_fid = source_fid_to_dest.get(current)
                        if dest_fid is not None:
                            ordered.append((dest_fid, len(ordered) + 1))
                        current = src_to_next.get(current)

                if not ordered:
                    return

                # Wire up flows_into chain — same for both modes
                for i, (fid, _) in enumerate(ordered):
                    next_fid = ordered[i + 1][0] if i < len(ordered) - 1 else None
                    curs.execute("UPDATE sample_frame_features SET flows_into = ? WHERE fid = ?", [next_fid, fid])
                conn.commit()
        except Exception as ex:
            Settings().msg_bar("Error resolving topology", str(ex), MESSAGE_LEVEL_WARNING)

    def setupUi(self):

        self.resize(500, 400)
        self.setMinimumSize(QSize(500, 400))

        self.vert = QVBoxLayout(self)
        self.setLayout(self.vert)

        self.horiz_name = QHBoxLayout()
        self.vert.addLayout(self.horiz_name)

        self.lblName = QLabel("Name")
        self.horiz_name.addWidget(self.lblName)
        self.txtName = QLineEdit()
        self.txtName.setToolTip("The name of the sample frame")
        self.horiz_name.addWidget(self.txtName)

        self.tabs = QTabWidget(self)
        self.vert.addWidget(self.tabs)

        if self.tab_inputs is not None:
            self.tabs.addTab(self.tab_inputs, "Inputs")
        self.tabs.addTab(self.tab_properties, "Basic Properties")
        self.tabs.addTab(self.tab_description, "Description")
        self.tabs.addTab(self.tab_attributes, attributes_name)
        self.stats_widget.add_stats_tab(self.tabs)
        self.tabs.addTab(self.metadata_widget, "Metadata")

        self.chkAddToMap = QCheckBox("Add to Map")
        self.chkAddToMap.setChecked(True)

        self.chkStartEditSession = QCheckBox("Start Edit Session")
        self.chkStartEditSession.setChecked(False)
        self.chkAddToMap.stateChanged.connect(lambda state: (self.chkStartEditSession.setEnabled(state == CHECKED), self.chkStartEditSession.setChecked(False) if state != CHECKED else None))

        add_to_map_row = QHBoxLayout()
        add_to_map_row.addWidget(self.chkAddToMap)
        add_to_map_row.addWidget(self.chkStartEditSession)
        add_to_map_row.addStretch()
        self.vert.addLayout(add_to_map_row)

        self.buttonBox = add_standard_form_buttons(self, "inputs/sample-frames")
        self.vert.addLayout(self.buttonBox)


class SampleFrameInputs(QWidget):
    def __init__(self, parent, qris_project, import_feature_class=None):
        super().__init__(parent)

        self.qris_project = qris_project
        self.import_feature_class = import_feature_class
        name = self.import_feature_class.name() if isinstance(self.import_feature_class, QgsVectorLayer) else self.import_feature_class

        self.setupUi()
        self.txtImport.setText(name)
        self.load_fields()
        self.load_clip_to_aoi()

    def load_fields(self):

        # clear fields
        self.cboDisplayLabel.clear()
        self.cboFlowPathField.clear()
        self.cboTopologyField.clear()

        # add None option
        self.cboDisplayLabel.addItem("None")
        self.cboFlowPathField.addItem("None")
        self.cboTopologyField.addItem("None")

        # use QgsVectorLayer to get fields
        if self.import_feature_class is not None:
            if isinstance(self.import_feature_class, QgsVectorLayer):
                layer = self.import_feature_class
            else:
                layer = QgsVectorLayer(self.import_feature_class, "temp", "ogr")
            fields = layer.fields()
            for field in fields:
                self.cboDisplayLabel.addItem(field.name())
                self.cboFlowPathField.addItem(field.name())
                if field.type() in [QMetaType.Type.Int, QMetaType.Type.Double]:  # integer
                    self.cboTopologyField.addItem(field.name())

    def load_clip_to_aoi(self):
        # get the aoi and valley bottoms from the project
        self.cboClipToAOI.clear()

        clip_items = {**self.qris_project.valley_bottoms, **self.qris_project.aois}
        no_aoi = DBItem("None", 0, "None - Retain full dataset extent")
        clip_items[0] = no_aoi
        aoi_model = DBItemModel(clip_items)
        self.cboClipToAOI.setModel(aoi_model)

    def setupUi(self):

        self.vert = QVBoxLayout(self)
        self.setLayout(self.vert)

        self.grid = QGridLayout(self)
        self.vert.addLayout(self.grid)

        self.lblImport = QLabel("Import Features From")
        self.grid.addWidget(self.lblImport, 0, 0)
        self.txtImport = QLineEdit()
        self.txtImport.setEnabled(False)
        self.grid.addWidget(self.txtImport, 0, 1)
        # self.cmdImport = QPushButton('...')
        # self.cmdImport.clicked.connect(self.cmdImport_clicked)
        # self.grid.addWidget(self.cmdImport, 0, 2)

        self.lblClipToAOI = QLabel("Clip to AOI / Valley Bottom")
        self.grid.addWidget(self.lblClipToAOI, 1, 0)
        self.cboClipToAOI = QComboBox()
        self.cboClipToAOI.setToolTip("Optionally clip the sample frame to the selected AOI or valley bottom")
        self.grid.addWidget(self.cboClipToAOI, 1, 1, 1, 2)

        self.grid.setRowMinimumHeight(2, 8)

        # Optional fields group box
        self.grpOptional = QGroupBox("Specify Fields to Import (Optional)")
        grp_grid = QGridLayout()
        self.grpOptional.setLayout(grp_grid)
        self.grid.addWidget(self.grpOptional, 3, 0, 1, 3)

        self.lblName = QLabel("Sample Frame Labels")
        grp_grid.addWidget(self.lblName, 0, 0)
        self.cboDisplayLabel = QComboBox()
        self.cboDisplayLabel.setToolTip("Field to use for display labels.")
        grp_grid.addWidget(self.cboDisplayLabel, 0, 1, 1, 2)

        self.lblFlowPathField = QLabel(f"{flow_path}")
        grp_grid.addWidget(self.lblFlowPathField, 1, 0)
        self.cboFlowPathField = QComboBox()
        self.cboFlowPathField.setToolTip("Field to specify unique flow path names.")
        grp_grid.addWidget(self.cboFlowPathField, 1, 1, 1, 2)

        self.lblTopologyField = QLabel("Topology")
        grp_grid.addWidget(self.lblTopologyField, 2, 0)
        self.cboTopologyField = QComboBox()
        self.cboTopologyField.setToolTip("Field to use for topology inference. Value should ref.")
        grp_grid.addWidget(self.cboTopologyField, 2, 1, 1, 2)

        self.rdoSequence = QRadioButton("Sequence")
        self.rdoFlowsInto = QRadioButton("Flows Into Reference")
        self.rdoSequence.setChecked(True)
        topology_type_layout = QHBoxLayout()
        topology_type_layout.addWidget(self.rdoSequence)
        topology_type_layout.addWidget(self.rdoFlowsInto)
        topology_type_layout.addStretch()
        topology_type_widget = QWidget()
        topology_type_widget.setLayout(topology_type_layout)
        grp_grid.addWidget(QLabel("Topology Represents"), 3, 0)
        grp_grid.addWidget(topology_type_widget, 3, 1, 1, 2)
        topology_type_widget.setEnabled(False)
        self.cboTopologyField.currentIndexChanged.connect(lambda idx: topology_type_widget.setEnabled(idx > 0))

        self.vert.addStretch()


class SampleFrameInputsCreate(QWidget):
    def __init__(self, parent, project: Project = None):
        super().__init__(parent)

        self.qris_project = project

        self.setupUi()
        self.load_layers()

    def load_layers(self):

        self.cboFramePolygon.clear()
        self.cboCenterline.clear()
        self.cboCrossSections.clear()

        # Set Cross Sections, set init if exists
        cross_sections = {id: xsection for id, xsection in self.qris_project.cross_sections.items()}
        self.cross_sections_model = DBItemModel(cross_sections)
        self.cboCrossSections.setModel(self.cross_sections_model)
        # if cross_sections_init is not None:
        #     self.cboCrossSections.setCurrentIndex(self.cross_sections_model.getItemIndex(cross_sections_init))

        # Masks (filtered to just AOI)
        aois = {f"aoi_{id}": aoi for id, aoi in self.qris_project.aois.items()}
        valley_bottoms = {f"valley_bottom_{id}": valley_bottom for id, valley_bottom in self.qris_project.valley_bottoms.items()}
        # context = {f'context_{id}': layer for id, layer in self.qris_project.scratch_vectors.items() if QgsVectorLayer(f'{layer.gpkg_path}|layername={layer.fc_name}').geometryType() == QgsWkbTypes.PolygonGeometry}
        self.polygons = {**valley_bottoms, **aois}
        self.polygons_model = DBItemModel(self.polygons)
        self.cboFramePolygon.setModel(self.polygons_model)
        # if polygon_init is not None:
        #     index = self.polygons_model.getItemIndex(polygon_init)
        #     if index is not None:
        #         self.cboFramePolygon.setCurrentIndex(index)

        # Centerlines
        centerlines = {id: centerline for id, centerline in self.qris_project.profiles.items()}
        self.centerlines_model = DBItemModel(centerlines)
        self.cboCenterline.setModel(self.centerlines_model)

    def setupUi(self):

        self.vert = QVBoxLayout(self)
        self.setLayout(self.vert)

        self.grid = QGridLayout(self)
        self.vert.addLayout(self.grid)

        self.lblPolygon = QLabel("Polygon Layer")
        self.grid.addWidget(self.lblPolygon, 0, 0)
        self.cboFramePolygon = QComboBox()
        self.grid.addWidget(self.cboFramePolygon, 0, 1, 1, 2)

        self.lblCenterline = QLabel("Centerline")
        self.grid.addWidget(self.lblCenterline, 1, 0)
        self.cboCenterline = QComboBox()
        self.grid.addWidget(self.cboCenterline, 1, 1, 1, 2)

        self.lblCrossSections = QLabel("Cross Sections")
        self.grid.addWidget(self.lblCrossSections, 2, 0)
        self.cboCrossSections = QComboBox()
        self.grid.addWidget(self.cboCrossSections, 2, 1, 1, 2)

        self.vert.addStretch()


class SampleFrameProperties(QWidget):
    def __init__(self, parent, sample_frame: SampleFrame = None):
        super().__init__(parent)

        self.sample_frame = sample_frame

        self.setupUi()

        self.txtLabelField.setText(label_field_name)
        self.txtFlowField.setText(flow_path_field_name)

        if self.sample_frame is not None:
            self.txtDefaultFlowPathName.setText(self.sample_frame.default_flow_path_name)
            self.txtDefaultFlowPathName.setEnabled(False)
        else:
            self.txtDefaultFlowPathName.setText(default_flow_path_name)

    def setupUi(self):

        self.vert = QVBoxLayout(self)
        self.grid = QGridLayout()
        self.vert.addLayout(self.grid)

        self.lblLabelField = QLabel("Labels Field Name")
        self.grid.addWidget(self.lblLabelField, 1, 0)
        self.lblLabelField.setVisible(False)
        self.txtLabelField = QLineEdit()
        self.txtLabelField.setEnabled(False)
        self.grid.addWidget(self.txtLabelField, 1, 1)
        self.txtLabelField.setVisible(False)

        self.lblFlowField = QLabel("Flow Into Field Name")
        self.grid.addWidget(self.lblFlowField, 2, 0)
        self.lblFlowField.setVisible(False)
        self.txtFlowField = QLineEdit()
        self.txtFlowField.setEnabled(False)
        self.grid.addWidget(self.txtFlowField, 2, 1)
        self.txtFlowField.setVisible(False)

        self.lblDefaultFlowPathName = QLabel(f"Default {flow_path} Name")
        self.grid.addWidget(self.lblDefaultFlowPathName, 3, 0)
        self.txtDefaultFlowPathName = QLineEdit()
        self.txtDefaultFlowPathName.setPlaceholderText("Not specified")
        self.grid.addWidget(self.txtDefaultFlowPathName, 3, 1)

        self.vert.addStretch()


class SampleFrameDescription(QWidget):
    def __init__(self, parent, sample_frame: SampleFrame = None):
        super().__init__(parent)

        self.setupUi()

        if sample_frame is not None and sample_frame.description:
            self.txtDescription.setPlainText(sample_frame.description)

    def setupUi(self):

        self.vert = QVBoxLayout(self)
        self.setLayout(self.vert)

        self.txtDescription = QTextEdit()
        self.vert.addWidget(self.txtDescription)


class SampleFrameAttributes(QWidget):
    def __init__(self, parent, sample_frame: SampleFrame = None):
        super().__init__(parent)

        self.sample_frame = sample_frame
        self.model = EditableTreeModel()

        self.setupUi()

        if sample_frame is not None:
            self.load_attributes()

    def load_attributes(self):

        if not self.sample_frame.fields:
            return
        # add to model
        for field in self.sample_frame.fields:
            name = field if isinstance(field, str) else field["label"]
            values = list() if isinstance(field, str) else field["values"]
            item = QStandardItem(name)
            item.setData(values, USER_ROLE)
            self.model.appendRow(item)
            # add attributes as children
            for attribute in values:
                child_item = QStandardItem(str(attribute))
                item.appendRow(child_item)

    def selection_changed(self, selected, deselected):

        if self.sample_frame is None:
            self.cmdEdit.setEnabled(len(selected.indexes()) > 0)
            self.cmdDelete.setEnabled(len(selected.indexes()) > 0)

    def cmdAdd_clicked(self):

        # get list of field names in the model
        existing_fields = [self.model.item(i).text() for i in range(self.model.rowCount())]

        frm = FrmNewAttribute(self, self.sample_frame, existing_fields=existing_fields)
        frm.exec()

        if frm.result() == DLG_ACCEPTED:
            if frm.name is not None:
                self.add_attribute(frm.name, frm.attributes)

    def cmdEdit_clicked(self):

        # get list of field names in the model
        existing_fields = [self.model.item(i).text() for i in range(self.model.rowCount())]

        # get selected item
        index = self.treeView.currentIndex()
        item = self.model.itemFromIndex(index)
        field_name = item.text()

        # get attributes These are the child nodes of this item
        attributes = [item.child(i).text() for i in range(item.rowCount())]

        frm = FrmNewAttribute(self, field_name, attributes, existing_fields)
        frm.exec()

        if frm.result() == DLG_ACCEPTED:
            if frm.name is not None:
                # remove existing item
                self.model.removeRow(index.row())

                # add new item
                self.add_attribute(frm.name, frm.attributes)

    def cmdDelete_clicked(self):

        # get selected item
        index = self.treeView.currentIndex()
        self.model.itemFromIndex(index)

        # remove existing item
        self.model.removeRow(index.row())

    def add_attribute(self, field_name: str, attributes: list):

        # add to model
        item = QStandardItem(field_name)
        item.setData(attributes, USER_ROLE)
        self.model.appendRow(item)

        # add attributes as children
        for attribute in attributes:
            child_item = QStandardItem(attribute)
            item.appendRow(child_item)

        # add to sample frame
        # self.sample_frame.add_attribute(field_name, attributes)

    def get_fields(self):

        fields = []
        for row in range(self.model.rowCount()):
            item = self.model.item(row)
            field_name = item.text()
            field_id = field_name.lower().replace(" ", "_")
            attributes = [item.child(i).text() for i in range(item.rowCount())]
            fields.append({"id": field_id, "label": field_name, "type": "list", "values": attributes})
        return fields

    def setupUi(self):

        self.horiz = QHBoxLayout()
        self.setLayout(self.horiz)

        self.treeView = QTreeView(self)
        self.treeView.header().hide()
        self.treeView.setModel(self.model)
        self.treeView.selectionModel().selectionChanged.connect(self.selection_changed)
        self.tree_state = {}

        self.cmdAdd = QPushButton("Add")
        if self.sample_frame is not None:
            self.cmdAdd.setEnabled(False)
        self.cmdAdd.clicked.connect(self.cmdAdd_clicked)

        self.cmdEdit = QPushButton("Edit")
        self.cmdEdit.setEnabled(False)
        self.cmdEdit.clicked.connect(self.cmdEdit_clicked)

        self.cmdDelete = QPushButton("Delete")
        self.cmdDelete.setEnabled(False)
        self.cmdDelete.clicked.connect(self.cmdDelete_clicked)

        # add buttons to layout on left side, then table on right side
        self.vert = QVBoxLayout()
        self.vert.addWidget(self.cmdAdd)
        self.vert.addWidget(self.cmdEdit)
        self.vert.addWidget(self.cmdDelete)
        self.vert.addStretch()

        self.horiz.addLayout(self.vert)
        self.horiz.addWidget(self.treeView)


class SampleFrameAttributesAddFields(QWidget):
    def __init__(self, parent, import_feature_class=None):
        super().__init__(parent)

        self.import_feature_class = import_feature_class
        self.fields = self.load_fields()
        self.chk_fields = []

        self.setupUi()

    def load_fields(self):

        # use QgsVectorLayer to get fields
        if self.import_feature_class is not None:
            if isinstance(self.import_feature_class, QgsVectorLayer):
                self.layer: QgsVectorLayer = self.import_feature_class
            else:
                self.layer = QgsVectorLayer(self.import_feature_class, "temp", "ogr")
            return self.layer.fields().names()

    def get_fields(self):

        fields = [chk.text() for chk in self.chk_fields if chk.isChecked()]
        outfields = []
        for field in fields:
            field_id = field.replace(" ", "_").lower()
            values = []
            for feature in self.layer.getFeatures():
                value = feature[field]
                if isinstance(value, QVariant):
                    value = None
                if value is not None and value not in values:
                    values.append(str(value))
            outfields.append({"id": field_id, "label": field, "type": "list", "values": values})

        return outfields

    def setupUi(self):

        self.vert = QVBoxLayout(self)
        self.setLayout(self.vert)

        # List of checkboxes for each field in the feature class
        self.lbl_fields = QLabel(f"Import Fields for Sample Frame {attributes_name}")
        self.vert.addWidget(self.lbl_fields)

        self.box = QGroupBox()
        self.box_layout = QVBoxLayout()
        self.box_layout.setAlignment(ALIGN_TOP)
        self.box.setLayout(self.box_layout)

        for field in self.fields:
            chk = QCheckBox(field)
            chk.setChecked(True)
            self.chk_fields.append(chk)
            self.box_layout.addWidget(chk)

        self.vert.addWidget(self.box)

        self.vert.addStretch()


class EditableTreeModel(QStandardItemModel):
    def flags(self, index):
        default_flags = super().flags(index)

        if index.parent().isValid():  # If the item has a parent, it's not a top-level item
            return default_flags & ~ITEM_FLAG_SELECTABLE  # Remove the selectable flag
        else:
            return default_flags | ITEM_FLAG_EDITABLE  # Add the editable flag for top-level items
