Asset-Frameworker/gui/unified_view_model.py

# gui/unified_view_model.py
from PySide6.QtCore import QAbstractItemModel, QModelIndex, Qt
from pathlib import Path # Added for file_name extraction
from rule_structure import SourceRule, AssetRule, FileRule # Removed AssetType, ItemType import

class UnifiedViewModel(QAbstractItemModel):
    """
    A QAbstractItemModel for displaying and editing the hierarchical structure
    of SourceRule -> AssetRule -> FileRule.
    """
    Columns = [
        "Name", "Supplier Override", "Asset-Type Override",
        "Target Asset Name Override", "Item-Type Override",
        "Status", "Output Path"
    ]

    COL_NAME = 0
    COL_SUPPLIER = 1
    COL_ASSET_TYPE = 2
    COL_TARGET_ASSET = 3
    COL_ITEM_TYPE = 4
    COL_STATUS = 5
    COL_OUTPUT_PATH = 6

    def __init__(self, parent=None):
        super().__init__(parent)
        self._source_rules = [] # Now stores a list of SourceRule objects

    def load_data(self, source_rules_list: list): # Accepts a list
        """Loads or reloads the model with a list of SourceRule objects."""
        self.beginResetModel()
        self._source_rules = source_rules_list if source_rules_list else [] # Assign the new list
        # Ensure back-references for parent lookup are set on the NEW items
        for source_rule in self._source_rules:
            for asset_rule in source_rule.assets:
                asset_rule.parent_source = source_rule # Set parent SourceRule
                for file_rule in asset_rule.files:
                    file_rule.parent_asset = asset_rule # Set parent AssetRule
        self.endResetModel()

    def clear_data(self):
        """Clears the model data."""
        self.beginResetModel()
        self._source_rules = [] # Clear the list
        self.endResetModel()

    def get_all_source_rules(self) -> list:
        """Returns the internal list of SourceRule objects."""
        return self._source_rules
    def rowCount(self, parent: QModelIndex = QModelIndex()) -> int:
        """Returns the number of rows under the given parent."""
        if not parent.isValid():
            # Parent is the invisible root. Children are the SourceRules.
            return len(self._source_rules)

        parent_item = parent.internalPointer()

        if isinstance(parent_item, SourceRule):
             # Parent is a SourceRule. Children are AssetRules.
             return len(parent_item.assets)
        elif isinstance(parent_item, AssetRule):
             # Parent is an AssetRule. Children are FileRules.
            return len(parent_item.files)
        elif isinstance(parent_item, FileRule):
            return 0 # FileRules have no children

        return 0 # Should not happen for valid items


    def columnCount(self, parent: QModelIndex = QModelIndex()) -> int:
        """Returns the number of columns."""
        return len(self.Columns)

    def parent(self, index: QModelIndex) -> QModelIndex:
        """Returns the parent of the model item with the given index."""
        if not index.isValid():
            return QModelIndex()

        child_item = index.internalPointer()
        if child_item is None:
             return QModelIndex()

        # Determine the parent based on the item type
        if isinstance(child_item, SourceRule):
             # Parent is the invisible root
             return QModelIndex()
        elif isinstance(child_item, AssetRule):
             # Parent is a SourceRule. Find its row in the _source_rules list.
             parent_item = getattr(child_item, 'parent_source', None)
             if parent_item and parent_item in self._source_rules:
                 try:
                     parent_row = self._source_rules.index(parent_item)
                     return self.createIndex(parent_row, 0, parent_item)
                 except ValueError:
                      return QModelIndex() # Should not happen if parent_source is correct
             else:
                 return QModelIndex() # Parent SourceRule not found or reference missing

        elif isinstance(child_item, FileRule):
            # Parent is an AssetRule. Find its row within its parent SourceRule.
            parent_item = getattr(child_item, 'parent_asset', None) # Get parent AssetRule
            if parent_item:
                 grandparent_item = getattr(parent_item, 'parent_source', None) # Get the SourceRule
                 if grandparent_item:
                     try:
                         parent_row = grandparent_item.assets.index(parent_item)
                         # We need the index of the grandparent (SourceRule) to create the parent index
                         grandparent_row = self._source_rules.index(grandparent_item)
                         return self.createIndex(parent_row, 0, parent_item) # Create index for the AssetRule parent
                     except ValueError:
                         return QModelIndex() # Parent AssetRule or Grandparent SourceRule not found in respective lists
                 else:
                      return QModelIndex() # Grandparent (SourceRule) reference missing
            else:
                 return QModelIndex() # Parent AssetRule reference missing

        return QModelIndex() # Should not be reached


    def index(self, row: int, column: int, parent: QModelIndex = QModelIndex()) -> QModelIndex:
        """Returns the index of the item in the model specified by the given row, column and parent index."""
        if not self.hasIndex(row, column, parent):
             return QModelIndex()

        parent_item = None
        if not parent.isValid():
            # Parent is invisible root. Children are SourceRules.
            if row < len(self._source_rules):
                 child_item = self._source_rules[row]
                 return self.createIndex(row, column, child_item)
            else:
                 return QModelIndex() # Row out of bounds for top-level items
        else:
            # Parent is a valid index, get its item
            parent_item = parent.internalPointer()

        child_item = None
        if isinstance(parent_item, SourceRule):
            # Parent is SourceRule. Children are AssetRules.
            if row < len(parent_item.assets):
                child_item = parent_item.assets[row]
                # Ensure parent reference is set
                if not hasattr(child_item, 'parent_source'):
                     child_item.parent_source = parent_item
        elif isinstance(parent_item, AssetRule):
            # Parent is AssetRule. Children are FileRules.
            if row < len(parent_item.files):
                child_item = parent_item.files[row]
                # Ensure parent reference is set
                if not hasattr(child_item, 'parent_asset'):
                    child_item.parent_asset = parent_item

        if child_item:
            # Create index for the child item under the parent
            return self.createIndex(row, column, child_item)
        else:
            # Invalid row or parent type has no children (FileRule)
            return QModelIndex()

    def data(self, index: QModelIndex, role: int = Qt.DisplayRole):
        """Returns the data stored under the given role for the item referred to by the index."""
        if not index.isValid(): # Check only index validity, data list might be empty but valid
            return None

        item = index.internalPointer()
        column = index.column()

        # --- Handle different item types ---
        if isinstance(item, SourceRule): # This might only be relevant if SourceRule is displayed
            if role == Qt.DisplayRole:
                if column == 0: return item.input_path
                # Use supplier_override if set, otherwise empty string
                if column == self.COL_SUPPLIER: return item.supplier_override if item.supplier_override is not None else ""
                # Other columns return None or "" for SourceRule
            elif role == Qt.EditRole:
                 # Return supplier_override for editing
                 if column == self.COL_SUPPLIER: return item.supplier_override if item.supplier_override is not None else ""
            return None # Default for SourceRule for other roles/columns

        elif isinstance(item, AssetRule):
            if role == Qt.DisplayRole:
                if column == self.COL_NAME: return item.asset_name
                # Use asset_type_override if set, otherwise fall back to predicted asset_type
                if column == self.COL_ASSET_TYPE:
                    display_value = item.asset_type_override if item.asset_type_override is not None else item.asset_type
                    return display_value if display_value else ""
                # Placeholder columns
                if column == self.COL_STATUS: return "" # Status (Not handled yet)
                if column == self.COL_OUTPUT_PATH: return "" # Output Path (Not handled yet)
            elif role == Qt.EditRole:
                 # Return asset_type_override for editing (delegate expects string or None)
                 if column == self.COL_ASSET_TYPE:
                     return item.asset_type_override # Return string or None
            return None # Default for AssetRule


        elif isinstance(item, FileRule):
            if role == Qt.DisplayRole:
                if column == self.COL_NAME: return Path(item.file_path).name # Display only filename
                # Use target_asset_name_override if set, otherwise empty string
                if column == self.COL_TARGET_ASSET:
                    return item.target_asset_name_override if item.target_asset_name_override is not None else ""
                # Use item_type_override if set, otherwise empty string (assuming predicted isn't stored directly)
                if column == self.COL_ITEM_TYPE:
                    # Assuming item_type_override stores the string name of the ItemType enum
                    return item.item_type_override if item.item_type_override else ""
                if column == self.COL_STATUS: return "" # Status (Not handled yet)
                if column == self.COL_OUTPUT_PATH: return "" # Output Path (Not handled yet)
            elif role == Qt.EditRole:
                # Return target_asset_name_override for editing
                if column == self.COL_TARGET_ASSET: return item.target_asset_name_override if item.target_asset_name_override is not None else ""
                # Return item_type_override for editing (delegate expects string or None)
                if column == self.COL_ITEM_TYPE: return item.item_type_override # Return string or None
            return None # Default for FileRule

        return None # Should not be reached if item is one of the known types

    def setData(self, index: QModelIndex, value, role: int = Qt.EditRole) -> bool:
        """Sets the role data for the item at index to value."""
        if not index.isValid() or role != Qt.EditRole: # Check only index and role
            return False

        item = index.internalPointer()
        if item is None: # Extra check for safety
             return False
        column = index.column()
        changed = False

        # --- Handle different item types ---
        if isinstance(item, SourceRule): # If SourceRule is editable
            if column == self.COL_SUPPLIER:
                # Ensure value is string or None
                new_value = str(value).strip() if value is not None else None
                if new_value == "": new_value = None # Treat empty string as None
                # Update supplier_override
                if item.supplier_override != new_value:
                    item.supplier_override = new_value
                    changed = True

        elif isinstance(item, AssetRule):
            if column == self.COL_ASSET_TYPE:
                # Delegate provides string value (e.g., "Surface", "Model") or None
                new_value = str(value) if value is not None else None
                if new_value == "": new_value = None # Treat empty string as None
                # Update asset_type_override
                if item.asset_type_override != new_value:
                    item.asset_type_override = new_value
                    changed = True

        elif isinstance(item, FileRule):
            if column == self.COL_TARGET_ASSET: # Target Asset Name Override
                # Ensure value is string or None
                new_value = str(value).strip() if value is not None else None
                if new_value == "": new_value = None # Treat empty string as None
                # Update target_asset_name_override
                if item.target_asset_name_override != new_value:
                    item.target_asset_name_override = new_value
                    changed = True
            elif column == self.COL_ITEM_TYPE: # Item-Type Override
                 # Delegate provides string value (e.g., "MAP_COL") or None
                 new_value = str(value) if value is not None else None
                 if new_value == "": new_value = None # Treat empty string as None
                 # Update item_type_override
                 if item.item_type_override != new_value:
                    item.item_type_override = new_value
                    changed = True


        if changed:
            # Emit dataChanged for the specific index and affected roles
            self.dataChanged.emit(index, index, [Qt.DisplayRole, Qt.EditRole])
            return True

        return False

    def flags(self, index: QModelIndex) -> Qt.ItemFlags:
        """Returns the item flags for the given index."""
        if not index.isValid():
             return Qt.NoItemFlags # No flags for invalid index

        # Start with default flags for a valid item
        default_flags = Qt.ItemIsEnabled | Qt.ItemIsSelectable

        item = index.internalPointer()
        column = index.column()

        can_edit = False
        # Determine editability based on item type and column
        if isinstance(item, SourceRule): # If SourceRule is displayed/editable
            if column == 1: can_edit = True
        elif isinstance(item, AssetRule):
            if column == 2: can_edit = True
        elif isinstance(item, FileRule):
            if column == 3: can_edit = True
            if column == 4: can_edit = True

        if can_edit:
            return default_flags | Qt.ItemIsEditable
        else:
            return default_flags

    def headerData(self, section: int, orientation: Qt.Orientation, role: int = Qt.DisplayRole):
        """Returns the data for the given role and section in the header."""
        if orientation == Qt.Horizontal and role == Qt.DisplayRole:
            if 0 <= section < len(self.Columns):
                return self.Columns[section]
        # Optionally handle Vertical header (row numbers)
        # if orientation == Qt.Vertical and role == Qt.DisplayRole:
        #     return str(section + 1)
        return None

    # Helper to get item from index
    def getItem(self, index: QModelIndex):
        """Safely returns the item associated with the index."""
        if index.isValid():
            item = index.internalPointer()
            if item: # Ensure internal pointer is not None
                 return item
        return None # Return None for invalid index or None pointer