#!/usr/bin/env python3
from __future__ import annotations

import logging, argparse, os, sys, enum, warnings
from typing import Any, Optional, Tuple, Type
from pathlib import Path
import numpy as np
from PySide6.QtWidgets import (
    QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
    QPushButton, QLabel, QLineEdit, QFileDialog, QTableWidget,
    QTableWidgetItem, QComboBox, QMessageBox, QTabWidget,
    QTextEdit, QFormLayout,
    QHeaderView, QDialog, QDialogButtonBox
)
from PySide6.QtCore import Qt
from PySide6.QtGui import QIcon

from gguf_connector.reader import GGUFReader, GGUFValueType, ReaderField
from gguf_connector.writer import GGUFWriter
from gguf_connector.const import TokenType, RopeScalingType, PoolingType, GGMLQuantizationType, Keys

logger = logging.getLogger("pig-reader")
icon_path = os.path.join(os.path.dirname(__file__), "logo.png")

# Map of key names to enum types for automatic enum interpretation
KEY_TO_ENUM_TYPE = {
    Keys.Tokenizer.TOKEN_TYPE: TokenType,
    Keys.Rope.SCALING_TYPE: RopeScalingType,
    Keys.LLM.POOLING_TYPE: PoolingType,
    Keys.General.FILE_TYPE: GGMLQuantizationType,
}

# Define the tokenizer keys that should be edited together
TOKENIZER_LINKED_KEYS = [
    Keys.Tokenizer.LIST,
    Keys.Tokenizer.TOKEN_TYPE,
    Keys.Tokenizer.SCORES
]

class TokenizerEditorDialog(QDialog):
    def __init__(self, tokens, token_types, scores, parent=None):
        super().__init__(parent)
        self.setWindowTitle("Edit Tokenizer Data")
        self.resize(900, 600)

        self.tokens = tokens.copy() if tokens else []
        self.token_types = token_types.copy() if token_types else []
        self.scores = scores.copy() if scores else []

        # Ensure all arrays have the same length
        max_len = max(len(self.tokens), len(self.token_types), len(self.scores))
        if len(self.tokens) < max_len:
            self.tokens.extend([""] * (max_len - len(self.tokens)))
        if len(self.token_types) < max_len:
            self.token_types.extend([0] * (max_len - len(self.token_types)))
        if len(self.scores) < max_len:
            self.scores.extend([0.0] * (max_len - len(self.scores)))

        layout = QVBoxLayout(self)

        # Add filter controls
        filter_layout = QHBoxLayout()
        filter_layout.addWidget(QLabel("Filter:"))
        self.filter_edit = QLineEdit()
        self.filter_edit.setPlaceholderText("Type to filter tokens...")
        self.filter_edit.textChanged.connect(self.apply_filter)
        filter_layout.addWidget(self.filter_edit)

        # Add page controls
        self.page_size = 100  # Show 100 items per page
        self.current_page = 0
        self.total_pages = max(1, (len(self.tokens) + self.page_size - 1) // self.page_size)

        self.page_label = QLabel(f"Page 1 of {self.total_pages}")
        filter_layout.addWidget(self.page_label)

        prev_page = QPushButton("Previous")
        prev_page.clicked.connect(self.previous_page)
        filter_layout.addWidget(prev_page)

        next_page = QPushButton("Next")
        next_page.clicked.connect(self.next_page)
        filter_layout.addWidget(next_page)

        layout.addLayout(filter_layout)

        # Tokenizer data table
        self.tokens_table = QTableWidget()
        self.tokens_table.setColumnCount(4)
        self.tokens_table.setHorizontalHeaderLabels(["Index", "Token", "Type", "Score"])
        self.tokens_table.horizontalHeader().setSectionResizeMode(0, QHeaderView.ResizeMode.ResizeToContents)
        self.tokens_table.horizontalHeader().setSectionResizeMode(1, QHeaderView.ResizeMode.Stretch)
        self.tokens_table.horizontalHeader().setSectionResizeMode(2, QHeaderView.ResizeMode.ResizeToContents)
        self.tokens_table.horizontalHeader().setSectionResizeMode(3, QHeaderView.ResizeMode.ResizeToContents)

        layout.addWidget(self.tokens_table)

        # Controls
        controls_layout = QHBoxLayout()
        add_button = QPushButton("Add Token")
        add_button.clicked.connect(self.add_token)
        controls_layout.addWidget(add_button)
        remove_button = QPushButton("Remove Selected")
        remove_button.clicked.connect(self.remove_selected)
        controls_layout.addWidget(remove_button)
        controls_layout.addStretch()
        layout.addLayout(controls_layout)

        # Buttons
        buttons = QDialogButtonBox(QDialogButtonBox.StandardButton.Ok | QDialogButtonBox.StandardButton.Cancel)
        buttons.accepted.connect(self.accept)
        buttons.rejected.connect(self.reject)
        layout.addWidget(buttons)

        # Initialize the filtered values
        self.filtered_indices = list(range(len(self.tokens)))

        # Load data for the first page
        self.load_page()

    def apply_filter(self):
        """Filter the tokens based on the search text."""
        filter_text = self.filter_edit.text().lower()

        if not filter_text:
            # No filter, show all values
            self.filtered_indices = list(range(len(self.tokens)))
        else:
            # Apply filter
            self.filtered_indices = []
            for i, token in enumerate(self.tokens):
                if filter_text in str(token).lower():
                    self.filtered_indices.append(i)

        # Reset to first page and reload
        self.total_pages = max(1, (len(self.filtered_indices) + self.page_size - 1) // self.page_size)
        self.current_page = 0
        self.page_label.setText(f"Page 1 of {self.total_pages}")
        self.load_page()

    def previous_page(self):
        """Go to the previous page of results."""
        if self.current_page > 0:
            self.current_page -= 1
            self.page_label.setText(f"Page {self.current_page + 1} of {self.total_pages}")
            self.load_page()

    def next_page(self):
        """Go to the next page of results."""
        if self.current_page < self.total_pages - 1:
            self.current_page += 1
            self.page_label.setText(f"Page {self.current_page + 1} of {self.total_pages}")
            self.load_page()

    def load_page(self):
        """Load the current page of tokenizer data."""
        self.tokens_table.setRowCount(0)  # Clear the table

        # Calculate start and end indices for the current page
        start_idx = self.current_page * self.page_size
        end_idx = min(start_idx + self.page_size, len(self.filtered_indices))

        # Pre-allocate rows for better performance
        self.tokens_table.setRowCount(end_idx - start_idx)

        for row, i in enumerate(range(start_idx, end_idx)):
            orig_idx = self.filtered_indices[i]

            # Index
            index_item = QTableWidgetItem(str(orig_idx))
            index_item.setData(Qt.ItemDataRole.UserRole, orig_idx)  # Store original index
            index_item.setFlags(index_item.flags() & ~Qt.ItemFlag.ItemIsEditable)
            self.tokens_table.setItem(row, 0, index_item)

            # Token
            token_item = QTableWidgetItem(str(self.tokens[orig_idx]))
            self.tokens_table.setItem(row, 1, token_item)

            # Token Type
            token_type = self.token_types[orig_idx] if orig_idx < len(self.token_types) else 0
            try:
                enum_val = TokenType(token_type)
                display_text = f"{enum_val.name} ({token_type})"
            except (ValueError, KeyError):
                display_text = f"Unknown ({token_type})"

            type_item = QTableWidgetItem(display_text)
            type_item.setData(Qt.ItemDataRole.UserRole, token_type)

            # Make type cell editable with a double-click handler
            type_item.setFlags(type_item.flags() & ~Qt.ItemFlag.ItemIsEditable)
            self.tokens_table.setItem(row, 2, type_item)

            # Score
            score = self.scores[orig_idx] if orig_idx < len(self.scores) else 0.0
            score_item = QTableWidgetItem(str(score))
            self.tokens_table.setItem(row, 3, score_item)

        # Connect double-click handler for token type cells
        self.tokens_table.cellDoubleClicked.connect(self.handle_cell_double_click)

    def handle_cell_double_click(self, row, column):
        """Handle double-click on a cell, specifically for token type editing."""
        if column == 2:  # Token Type column
            orig_item = self.tokens_table.item(row, 0)
            if orig_item:
                orig_idx = orig_item.data(Qt.ItemDataRole.UserRole)
                self.edit_token_type(row, orig_idx)

    def edit_token_type(self, row, orig_idx):
        """Edit a token type using a dialog with a dropdown of all enum options."""
        current_value = self.token_types[orig_idx] if orig_idx < len(self.token_types) else 0

        # Create a dialog with enum options
        dialog = QDialog(self)
        dialog.setWindowTitle("Select Token Type")
        layout = QVBoxLayout(dialog)

        combo = QComboBox()
        for enum_val in TokenType:
            combo.addItem(f"{enum_val.name} ({enum_val.value})", enum_val.value)

        # Set current value
        try:
            if isinstance(current_value, int):
                enum_val = TokenType(current_value)
                combo.setCurrentText(f"{enum_val.name} ({current_value})")
        except (ValueError, KeyError):
            pass

        layout.addWidget(combo)

        buttons = QDialogButtonBox(QDialogButtonBox.StandardButton.Ok | QDialogButtonBox.StandardButton.Cancel)
        buttons.accepted.connect(dialog.accept)
        buttons.rejected.connect(dialog.reject)
        layout.addWidget(buttons)

        if dialog.exec() == QDialog.DialogCode.Accepted:
            # Get the selected value
            new_value = combo.currentData()
            enum_val = TokenType(new_value)
            display_text = f"{enum_val.name} ({new_value})"

            # Update the display
            type_item = self.tokens_table.item(row, 2)
            if type_item:
                type_item.setText(display_text)
                type_item.setData(Qt.ItemDataRole.UserRole, new_value)

            # Update the actual value
            self.token_types[orig_idx] = new_value

    def add_token(self):
        """Add a new token to the end of the list."""
        # Add to the end of the arrays
        self.tokens.append("")
        self.token_types.append(0)  # Default to normal token
        self.scores.append(0.0)

        orig_idx = len(self.tokens) - 1

        # Add to filtered indices if it matches the current filter
        filter_text = self.filter_edit.text().lower()
        if not filter_text or filter_text in "":
            self.filtered_indices.append(orig_idx)

        # Update pagination
        self.total_pages = max(1, (len(self.filtered_indices) + self.page_size - 1) // self.page_size)

        # Go to the last page to show the new item
        self.current_page = self.total_pages - 1
        self.page_label.setText(f"Page {self.current_page + 1} of {self.total_pages}")

        # Reload the page
        self.load_page()

    def remove_selected(self):
        """Remove selected tokens from all arrays."""
        selected_rows = []
        for item in self.tokens_table.selectedItems():
            row = item.row()
            if row not in selected_rows:
                selected_rows.append(row)

        if not selected_rows:
            return

        # Get original indices in descending order to avoid index shifting
        orig_indices = []
        for row in selected_rows:
            orig_item = self.tokens_table.item(row, 0)
            if orig_item:
                orig_indices.append(orig_item.data(Qt.ItemDataRole.UserRole))
        orig_indices.sort(reverse=True)

        # Remove from all arrays
        for idx in orig_indices:
            if idx < len(self.tokens):
                del self.tokens[idx]
            if idx < len(self.token_types):
                del self.token_types[idx]
            if idx < len(self.scores):
                del self.scores[idx]

        # Rebuild filtered_indices
        self.filtered_indices = []
        filter_text = self.filter_edit.text().lower()

        for i, token in enumerate(self.tokens):
            if not filter_text or filter_text in str(token).lower():
                self.filtered_indices.append(i)

        # Update pagination
        self.total_pages = max(1, (len(self.filtered_indices) + self.page_size - 1) // self.page_size)
        self.current_page = min(self.current_page, self.total_pages - 1)
        self.page_label.setText(f"Page {self.current_page + 1} of {self.total_pages}")

        # Reload the page
        self.load_page()

    def get_data(self):
        """Return the edited tokenizer data."""
        return self.tokens, self.token_types, self.scores


class ArrayEditorDialog(QDialog):
    def __init__(self, array_values, element_type, key=None, parent=None):
        super().__init__(parent)
        self.setWindowTitle("Edit Array Values")
        self.resize(700, 500)

        self.array_values = array_values
        self.element_type = element_type
        self.key = key

        # Get enum type for this array if applicable
        self.enum_type = None
        if key in KEY_TO_ENUM_TYPE and element_type == GGUFValueType.INT32:
            self.enum_type = KEY_TO_ENUM_TYPE[key]

        layout = QVBoxLayout(self)

        # Add enum type information if applicable
        if self.enum_type is not None:
            enum_info_layout = QHBoxLayout()
            enum_label = QLabel(f"Editing {self.enum_type.__name__} values:")
            enum_info_layout.addWidget(enum_label)

            # Add a legend for the enum values
            enum_values = ", ".join([f"{e.name}={e.value}" for e in self.enum_type])
            enum_values_label = QLabel(f"Available values: {enum_values}")
            enum_values_label.setWordWrap(True)
            enum_info_layout.addWidget(enum_values_label, 1)

            layout.addLayout(enum_info_layout)

        # Add search/filter controls
        filter_layout = QHBoxLayout()
        filter_layout.addWidget(QLabel("Filter:"))
        self.filter_edit = QLineEdit()
        self.filter_edit.setPlaceholderText("Type to filter values...")
        self.filter_edit.textChanged.connect(self.apply_filter)
        filter_layout.addWidget(self.filter_edit)

        # Add page controls for large arrays
        self.page_size = 100  # Show 100 items per page
        self.current_page = 0
        self.total_pages = max(1, (len(array_values) + self.page_size - 1) // self.page_size)

        self.page_label = QLabel(f"Page 1 of {self.total_pages}")
        filter_layout.addWidget(self.page_label)

        prev_page = QPushButton("Previous")
        prev_page.clicked.connect(self.previous_page)
        filter_layout.addWidget(prev_page)

        next_page = QPushButton("Next")
        next_page.clicked.connect(self.next_page)
        filter_layout.addWidget(next_page)

        layout.addLayout(filter_layout)

        # Array items table
        self.items_table = QTableWidget()

        # Set up columns based on whether we have an enum type
        if self.enum_type is not None:
            self.items_table.setColumnCount(3)
            self.items_table.setHorizontalHeaderLabels(["Index", "Value", "Actions"])
            self.items_table.horizontalHeader().setSectionResizeMode(0, QHeaderView.ResizeMode.ResizeToContents)
            self.items_table.horizontalHeader().setSectionResizeMode(1, QHeaderView.ResizeMode.Stretch)
            self.items_table.horizontalHeader().setSectionResizeMode(2, QHeaderView.ResizeMode.ResizeToContents)
        else:
            self.items_table.setColumnCount(2)
            self.items_table.setHorizontalHeaderLabels(["Index", "Value"])
            self.items_table.horizontalHeader().setSectionResizeMode(0, QHeaderView.ResizeMode.ResizeToContents)
            self.items_table.horizontalHeader().setSectionResizeMode(1, QHeaderView.ResizeMode.Stretch)

        layout.addWidget(self.items_table)

        # Controls
        controls_layout = QHBoxLayout()
        add_button = QPushButton("Add Item")
        add_button.clicked.connect(self.add_item)
        controls_layout.addWidget(add_button)
        remove_button = QPushButton("Remove Selected")
        remove_button.clicked.connect(self.remove_selected)
        controls_layout.addWidget(remove_button)

        # Add bulk edit button for enum arrays
        if self.enum_type is not None:
            bulk_edit_button = QPushButton("Bulk Edit Selected")
            bulk_edit_button.clicked.connect(self.bulk_edit_selected)
            controls_layout.addWidget(bulk_edit_button)

        controls_layout.addStretch()
        layout.addLayout(controls_layout)

        # Buttons
        buttons = QDialogButtonBox(QDialogButtonBox.StandardButton.Ok | QDialogButtonBox.StandardButton.Cancel)
        buttons.accepted.connect(self.accept)
        buttons.rejected.connect(self.reject)
        layout.addWidget(buttons)

        # Initialize the filtered values
        self.filtered_indices = list(range(len(self.array_values)))

        # Load array values for the first page
        self.load_page()

    def apply_filter(self):
        """Filter the array values based on the search text."""
        filter_text = self.filter_edit.text().lower()

        if not filter_text:
            # No filter, show all values
            self.filtered_indices = list(range(len(self.array_values)))
        else:
            # Apply filter
            self.filtered_indices = []
            for i, value in enumerate(self.array_values):
                # For enum values, search in both name and value
                if self.enum_type is not None and isinstance(value, int):
                    try:
                        enum_val = self.enum_type(value)
                        display_text = f"{enum_val.name} ({value})".lower()
                        if filter_text in display_text:
                            self.filtered_indices.append(i)
                    except (ValueError, KeyError):
                        # If not a valid enum value, just check the raw value
                        if filter_text in str(value).lower():
                            self.filtered_indices.append(i)
                else:
                    # For non-enum values, just check the string representation
                    if filter_text in str(value).lower():
                        self.filtered_indices.append(i)

        # Reset to first page and reload
        self.total_pages = max(1, (len(self.filtered_indices) + self.page_size - 1) // self.page_size)
        self.current_page = 0
        self.page_label.setText(f"Page 1 of {self.total_pages}")
        self.load_page()

    def previous_page(self):
        """Go to the previous page of results."""
        if self.current_page > 0:
            self.current_page -= 1
            self.page_label.setText(f"Page {self.current_page + 1} of {self.total_pages}")
            self.load_page()

    def next_page(self):
        """Go to the next page of results."""
        if self.current_page < self.total_pages - 1:
            self.current_page += 1
            self.page_label.setText(f"Page {self.current_page + 1} of {self.total_pages}")
            self.load_page()

    def load_page(self):
        """Load the current page of array values."""
        self.items_table.setRowCount(0)  # Clear the table

        # Calculate start and end indices for the current page
        start_idx = self.current_page * self.page_size
        end_idx = min(start_idx + self.page_size, len(self.filtered_indices))

        # Pre-allocate rows for better performance
        self.items_table.setRowCount(end_idx - start_idx)

        for row, i in enumerate(range(start_idx, end_idx)):
            orig_idx = self.filtered_indices[i]
            value = self.array_values[orig_idx]

            # Index
            index_item = QTableWidgetItem(str(orig_idx))
            index_item.setData(Qt.ItemDataRole.UserRole, orig_idx)  # Store original index
            index_item.setFlags(index_item.flags() & ~Qt.ItemFlag.ItemIsEditable)
            self.items_table.setItem(row, 0, index_item)

            # Value
            if self.enum_type is not None:
                # Display enum value and name
                try:
                    if isinstance(value, (int, np.signedinteger)):
                        enum_val = self.enum_type(value)
                        display_text = f"{enum_val.name} ({value})"
                    else:
                        display_text = str(value)
                except (ValueError, KeyError):
                    display_text = f"Unknown ({value})"

                # Store the enum value in the item
                value_item = QTableWidgetItem(display_text)
                value_item.setData(Qt.ItemDataRole.UserRole, value)
                value_item.setFlags(value_item.flags() & ~Qt.ItemFlag.ItemIsEditable)
                self.items_table.setItem(row, 1, value_item)

                # Add an edit button in a separate column
                edit_button = QPushButton("Edit")
                edit_button.setProperty("row", row)
                edit_button.clicked.connect(self.edit_array_enum_value)

                # Create a widget to hold the button
                button_widget = QWidget()
                button_layout = QHBoxLayout(button_widget)
                button_layout.setContentsMargins(2, 2, 2, 2)
                button_layout.addWidget(edit_button)
                button_layout.addStretch()

                self.items_table.setCellWidget(row, 2, button_widget)
            else:
                value_item = QTableWidgetItem(str(value))
                self.items_table.setItem(row, 1, value_item)

    def edit_array_enum_value(self):
        """Handle editing an enum value in the array editor."""
        button = self.sender()
        row = button.property("row")

        # Get the original index from the table item
        orig_item = self.items_table.item(row, 0)
        new_item = self.items_table.item(row, 1)
        if orig_item and new_item and self.enum_type and self.edit_enum_value(row, self.enum_type):
            orig_idx = orig_item.data(Qt.ItemDataRole.UserRole)
            new_value = new_item.data(Qt.ItemDataRole.UserRole)
            # Update the stored value in the array
            if isinstance(new_value, (int, float, str, bool)):
                self.array_values[orig_idx] = new_value

    def bulk_edit_selected(self):
        """Edit multiple enum values at once."""
        if not self.enum_type:
            return

        selected_rows = set()
        for item in self.items_table.selectedItems():
            selected_rows.add(item.row())

        if not selected_rows:
            QMessageBox.information(self, "No Selection", "Please select at least one row to edit.")
            return

        # Create a dialog with enum options
        dialog = QDialog(self)
        dialog.setWindowTitle(f"Bulk Edit {self.enum_type.__name__} Values")
        layout = QVBoxLayout(dialog)

        layout.addWidget(QLabel(f"Set {len(selected_rows)} selected items to:"))

        combo = QComboBox()
        for enum_val in self.enum_type:
            combo.addItem(f"{enum_val.name} ({enum_val.value})", enum_val.value)

        layout.addWidget(combo)

        buttons = QDialogButtonBox(QDialogButtonBox.StandardButton.Ok | QDialogButtonBox.StandardButton.Cancel)
        buttons.accepted.connect(dialog.accept)
        buttons.rejected.connect(dialog.reject)
        layout.addWidget(buttons)

        if dialog.exec() == QDialog.DialogCode.Accepted:
            # Get the selected value
            new_value = combo.currentData()
            enum_val = self.enum_type(new_value)
            display_text = f"{enum_val.name} ({new_value})"

            # Update all selected rows
            for row in selected_rows:
                orig_item = self.items_table.item(row, 0)
                new_item = self.items_table.item(row, 1)
                if orig_item and new_item:
                    orig_idx = orig_item.data(Qt.ItemDataRole.UserRole)
                    self.array_values[orig_idx] = new_value

                    # Update the display
                    new_item.setText(display_text)
                    new_item.setData(Qt.ItemDataRole.UserRole, new_value)

    def add_item(self):
        # Add to the end of the array
        orig_idx = len(self.array_values)
        # Add default value based on type
        if self.enum_type is not None:
            # Default to first enum value
            default_value = list(self.enum_type)[0].value
            self.array_values.append(default_value)
        else:
            if self.element_type == GGUFValueType.STRING:
                self.array_values.append("")
            else:
                self.array_values.append(0)

        # Add to filtered indices if it matches the current filter
        self.filtered_indices.append(orig_idx)

        # Update pagination
        self.total_pages = max(1, (len(self.filtered_indices) + self.page_size - 1) // self.page_size)

        # Go to the last page to show the new item
        self.current_page = self.total_pages - 1
        self.page_label.setText(f"Page {self.current_page + 1} of {self.total_pages}")

        # Reload the page
        self.load_page()

    def remove_selected(self):
        selected_rows = []
        for item in self.items_table.selectedItems():
            row = item.row()
            if row not in selected_rows:
                selected_rows.append(row)

        if not selected_rows:
            return

        # Get original indices in descending order to avoid index shifting
        orig_indices = list()
        for row in selected_rows:
            orig_item = self.items_table.item(row, 0)
            if orig_item:
                orig_indices.append(orig_item.data(Qt.ItemDataRole.UserRole))
        orig_indices.sort(reverse=True)

        # Remove from array_values
        for idx in orig_indices:
            del self.array_values[idx]

        # Rebuild filtered_indices
        self.filtered_indices = []
        filter_text = self.filter_edit.text().lower()

        for i, value in enumerate(self.array_values):
            if not filter_text:
                self.filtered_indices.append(i)
            else:
                # Apply filter
                if self.enum_type is not None and isinstance(value, int):
                    try:
                        enum_val = self.enum_type(value)
                        display_text = f"{enum_val.name} ({value})".lower()
                        if filter_text in display_text:
                            self.filtered_indices.append(i)
                    except (ValueError, KeyError):
                        if filter_text in str(value).lower():
                            self.filtered_indices.append(i)
                else:
                    if filter_text in str(value).lower():
                        self.filtered_indices.append(i)

        # Update pagination
        self.total_pages = max(1, (len(self.filtered_indices) + self.page_size - 1) // self.page_size)
        self.current_page = min(self.current_page, self.total_pages - 1)
        self.page_label.setText(f"Page {self.current_page + 1} of {self.total_pages}")

        # Reload the page
        self.load_page()

    def edit_enum_value(self, row: int, enum_type: Type[enum.Enum]):
        """Edit an enum value using a dialog with a dropdown of all enum options."""
        # Get the original index from the table item
        orig_item = self.items_table.item(row, 0)
        if orig_item:
            orig_idx = orig_item.data(Qt.ItemDataRole.UserRole)
        else:
            return
        current_value = self.array_values[orig_idx]

        # Create a dialog with enum options
        dialog = QDialog(self)
        dialog.setWindowTitle(f"Select {enum_type.__name__} Value")
        layout = QVBoxLayout(dialog)

        # Add description
        description = QLabel(f"Select a {enum_type.__name__} value:")
        layout.addWidget(description)

        # Use a combo box for quick selection
        combo = QComboBox()
        for enum_val in enum_type:
            combo.addItem(f"{enum_val.name} ({enum_val.value})", enum_val.value)

        # Set current value
        try:
            if isinstance(current_value, int):
                enum_val = enum_type(current_value)
                combo.setCurrentText(f"{enum_val.name} ({current_value})")
        except (ValueError, KeyError):
            pass

        layout.addWidget(combo)

        buttons = QDialogButtonBox(QDialogButtonBox.StandardButton.Ok | QDialogButtonBox.StandardButton.Cancel)
        buttons.accepted.connect(dialog.accept)
        buttons.rejected.connect(dialog.reject)
        layout.addWidget(buttons)

        if dialog.exec() == QDialog.DialogCode.Accepted:
            # Update the value display and stored data
            new_value = combo.currentData()
            enum_val = enum_type(new_value)
            display_text = f"{enum_val.name} ({new_value})"

            new_item = self.items_table.item(row, 1)
            if new_item:
                new_item.setText(display_text)
                new_item.setData(Qt.ItemDataRole.UserRole, new_value)

            # Update the actual array value
            self.array_values[orig_idx] = new_value
            return True
        return False

    def get_array_values(self):
        # The array_values list is kept up-to-date as edits are made
        return self.array_values


class AddMetadataDialog(QDialog):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.setWindowTitle("Add Metadata")
        self.resize(400, 200)

        layout = QVBoxLayout(self)

        form_layout = QFormLayout()

        self.key_edit = QLineEdit()
        form_layout.addRow("Key:", self.key_edit)

        self.type_combo = QComboBox()
        for value_type in GGUFValueType:
            if value_type != GGUFValueType.ARRAY:  # Skip array type for simplicity
                self.type_combo.addItem(value_type.name, value_type)
        form_layout.addRow("Type:", self.type_combo)

        self.value_edit = QTextEdit()
        form_layout.addRow("Value:", self.value_edit)

        layout.addLayout(form_layout)

        buttons = QDialogButtonBox(QDialogButtonBox.StandardButton.Ok | QDialogButtonBox.StandardButton.Cancel)
        buttons.accepted.connect(self.accept)
        buttons.rejected.connect(self.reject)
        layout.addWidget(buttons)

    def get_data(self) -> Tuple[str, GGUFValueType, Any]:
        key = self.key_edit.text()
        value_type = self.type_combo.currentData()
        value_text = self.value_edit.toPlainText()

        # Convert value based on type
        if value_type == GGUFValueType.UINT8:
            value = np.uint8(int(value_text))
        elif value_type == GGUFValueType.INT8:
            value = np.int8(int(value_text))
        elif value_type == GGUFValueType.UINT16:
            value = np.uint16(int(value_text))
        elif value_type == GGUFValueType.INT16:
            value = np.int16(int(value_text))
        elif value_type == GGUFValueType.UINT32:
            value = np.uint32(int(value_text))
        elif value_type == GGUFValueType.INT32:
            value = np.int32(int(value_text))
        elif value_type == GGUFValueType.FLOAT32:
            value = np.float32(float(value_text))
        elif value_type == GGUFValueType.BOOL:
            value = value_text.lower() in ('true', 'yes', '1')
        elif value_type == GGUFValueType.STRING:
            value = value_text
        else:
            value = value_text

        return key, value_type, value


class GGUFEditorWindow(QMainWindow):
    def __init__(self):
        super().__init__()

        app_icon = QIcon(icon_path)
        self.setWindowIcon(app_icon)
        self.setWindowTitle("pig-reader")
        self.resize(1000, 800)

        self.current_file = None
        self.reader = None
        self.modified = False
        self.metadata_changes = {}  # Store changes to apply when saving
        self.metadata_to_remove = set()  # Store keys to remove when saving
        self.on_metadata_changed_is_connected = False

        self.setup_ui()

    def setup_ui(self):
        central_widget = QWidget()
        self.setCentralWidget(central_widget)

        main_layout = QVBoxLayout(central_widget)

        # File controls
        file_layout = QHBoxLayout()

        self.file_path_edit = QLineEdit()
        self.file_path_edit.setReadOnly(True)
        file_layout.addWidget(self.file_path_edit)

        open_button = QPushButton("Open GGUF")
        open_button.clicked.connect(self.open_file)
        file_layout.addWidget(open_button)

        save_button = QPushButton("Save As...")
        save_button.clicked.connect(self.save_file)
        file_layout.addWidget(save_button)

        main_layout.addLayout(file_layout)

        # Tabs for different views
        self.tabs = QTabWidget()

        # Metadata tab
        self.metadata_tab = QWidget()
        metadata_layout = QVBoxLayout(self.metadata_tab)

        # Metadata table
        self.metadata_table = QTableWidget()
        self.metadata_table.setColumnCount(4)
        self.metadata_table.setHorizontalHeaderLabels(["Key", "Type", "Value", "Actions"])
        self.metadata_table.horizontalHeader().setSectionResizeMode(0, QHeaderView.ResizeMode.Stretch)
        self.metadata_table.horizontalHeader().setSectionResizeMode(1, QHeaderView.ResizeMode.ResizeToContents)
        self.metadata_table.horizontalHeader().setSectionResizeMode(2, QHeaderView.ResizeMode.Stretch)
        self.metadata_table.horizontalHeader().setSectionResizeMode(3, QHeaderView.ResizeMode.ResizeToContents)
        metadata_layout.addWidget(self.metadata_table)

        # Metadata controls
        metadata_controls = QHBoxLayout()
        add_metadata_button = QPushButton("Add Metadata")
        add_metadata_button.clicked.connect(self.add_metadata)
        metadata_controls.addWidget(add_metadata_button)
        metadata_controls.addStretch()
        metadata_layout.addLayout(metadata_controls)

        # Tensors tab
        self.tensors_tab = QWidget()
        tensors_layout = QVBoxLayout(self.tensors_tab)
        self.tensors_table = QTableWidget()
        self.tensors_table.setColumnCount(5)
        self.tensors_table.setHorizontalHeaderLabels(["Name", "Type", "Shape", "Elements", "Size (bytes)"])
        self.tensors_table.horizontalHeader().setSectionResizeMode(0, QHeaderView.ResizeMode.Stretch)
        self.tensors_table.horizontalHeader().setSectionResizeMode(1, QHeaderView.ResizeMode.ResizeToContents)
        self.tensors_table.horizontalHeader().setSectionResizeMode(2, QHeaderView.ResizeMode.ResizeToContents)
        self.tensors_table.horizontalHeader().setSectionResizeMode(3, QHeaderView.ResizeMode.ResizeToContents)
        self.tensors_table.horizontalHeader().setSectionResizeMode(4, QHeaderView.ResizeMode.ResizeToContents)
        tensors_layout.addWidget(self.tensors_table)

        # Add tabs to tab widget
        self.tabs.addTab(self.metadata_tab, "Metadata")
        self.tabs.addTab(self.tensors_tab, "Tensors")
        main_layout.addWidget(self.tabs)

        # Status bar
        self.statusBar().showMessage("Ready")

    def load_file(self, file_path):
        """Load a GGUF file by path"""
        try:
            self.statusBar().showMessage(f"Loading {file_path}...")
            QApplication.processEvents()

            self.reader = GGUFReader(file_path, 'r')
            self.current_file = file_path
            self.file_path_edit.setText(file_path)

            self.load_metadata()
            self.load_tensors()

            self.metadata_changes = {}
            self.metadata_to_remove = set()
            self.modified = False

            self.statusBar().showMessage(f"Loaded {file_path}")
            return True
        except Exception as e:
            QMessageBox.critical(self, "Error", f"Failed to open file: {str(e)}")
            self.statusBar().showMessage("Error loading file")
            return False

    def open_file(self):
        file_path, _ = QFileDialog.getOpenFileName(
            self, "Open GGUF File", "", "GGUF Files (*.gguf);;All Files (*)"
        )

        if not file_path:
            return

        self.load_file(file_path)

    def load_metadata(self):
        self.metadata_table.setRowCount(0)

        if not self.reader:
            return

        # Disconnect to prevent triggering during loading
        if self.on_metadata_changed_is_connected:
            with warnings.catch_warnings():
                warnings.filterwarnings('ignore')
                self.metadata_table.itemChanged.disconnect(self.on_metadata_changed)
            self.on_metadata_changed_is_connected = False

        for i, (key, field) in enumerate(self.reader.fields.items()):
            self.metadata_table.insertRow(i)

            # Key
            key_item = QTableWidgetItem(key)
            key_item.setFlags(key_item.flags() & ~Qt.ItemFlag.ItemIsEditable)
            self.metadata_table.setItem(i, 0, key_item)

            # Type
            if not field.types:
                type_str = "N/A"
            elif field.types[0] == GGUFValueType.ARRAY:
                nest_count = len(field.types) - 1
                element_type = field.types[-1].name
                # Check if this is an enum array
                enum_type = self.get_enum_for_key(key)
                if enum_type is not None and field.types[-1] == GGUFValueType.INT32:
                    element_type = enum_type.__name__
                type_str = '[' * nest_count + element_type + ']' * nest_count
            else:
                type_str = str(field.types[0].name)
                # Check if this is an enum field
                enum_type = self.get_enum_for_key(key)
                if enum_type is not None and field.types[0] == GGUFValueType.INT32:
                    type_str = enum_type.__name__

            type_item = QTableWidgetItem(type_str)
            type_item.setFlags(type_item.flags() & ~Qt.ItemFlag.ItemIsEditable)
            self.metadata_table.setItem(i, 1, type_item)

            # Value
            value_str = self.format_field_value(field)
            value_item = QTableWidgetItem(value_str)

            # Make only simple values editable
            if len(field.types) == 1 and field.types[0] != GGUFValueType.ARRAY:
                value_item.setFlags(value_item.flags() | Qt.ItemFlag.ItemIsEditable)
            else:
                value_item.setFlags(value_item.flags() & ~Qt.ItemFlag.ItemIsEditable)

            self.metadata_table.setItem(i, 2, value_item)

            # Actions
            actions_widget = QWidget()
            actions_layout = QHBoxLayout(actions_widget)
            actions_layout.setContentsMargins(2, 2, 2, 2)

            # Add Edit button for arrays and enum fields
            if field.types and field.types[0] == GGUFValueType.ARRAY:
                edit_button = QPushButton("Edit")
                edit_button.setProperty("row", i)
                edit_button.setProperty("key", key)
                edit_button.clicked.connect(self.edit_array_metadata)
                actions_layout.addWidget(edit_button)

                # Add special label for tokenizer linked fields
                if key in TOKENIZER_LINKED_KEYS:
                    edit_button.setText("Edit Tokenizer")
                    edit_button.setToolTip("Edit all tokenizer data together")
            elif len(field.types) == 1 and self.get_enum_for_key(key) is not None:
                edit_button = QPushButton("Edit")
                edit_button.setProperty("row", i)
                edit_button.setProperty("key", key)
                edit_button.clicked.connect(self.edit_metadata_enum)
                actions_layout.addWidget(edit_button)

            remove_button = QPushButton("Remove")
            remove_button.setProperty("row", i)
            remove_button.setProperty("key", key)
            remove_button.clicked.connect(self.remove_metadata)
            actions_layout.addWidget(remove_button)

            self.metadata_table.setCellWidget(i, 3, actions_widget)

        # Reconnect after loading
        self.metadata_table.itemChanged.connect(self.on_metadata_changed)
        self.on_metadata_changed_is_connected = True

    def extract_array_values(self, field: ReaderField) -> list:
        """Extract all values from an array field."""
        if not field.types or field.types[0] != GGUFValueType.ARRAY:
            return []

        curr_type = field.types[1]
        array_values = []
        total_elements = len(field.data)

        if curr_type == GGUFValueType.STRING:
            for element_pos in range(total_elements):
                value_string = str(bytes(field.parts[-1 - (total_elements - element_pos - 1) * 2]), encoding='utf-8')
                array_values.append(value_string)
        elif self.reader and curr_type in self.reader.gguf_scalar_to_np:
            for element_pos in range(total_elements):
                array_values.append(field.parts[-1 - (total_elements - element_pos - 1)][0])

        return array_values

    def get_enum_for_key(self, key: str) -> Optional[Type[enum.Enum]]:
        """Get the enum type for a given key if it exists."""
        return KEY_TO_ENUM_TYPE.get(key)

    def format_enum_value(self, value: Any, enum_type: Type[enum.Enum]) -> str:
        """Format a value as an enum if possible."""
        try:
            if isinstance(value, (int, str)):
                enum_value = enum_type(value)
                return f"{enum_value.name} ({value})"
        except (ValueError, KeyError):
            pass
        return str(value)

    def format_field_value(self, field: ReaderField) -> str:
        if not field.types:
            return "N/A"

        if len(field.types) == 1:
            curr_type = field.types[0]
            if curr_type == GGUFValueType.STRING:
                return str(bytes(field.parts[-1]), encoding='utf-8')
            elif self.reader and curr_type in self.reader.gguf_scalar_to_np:
                value = field.parts[-1][0]
                # Check if this field has an enum type
                enum_type = self.get_enum_for_key(field.name)
                if enum_type is not None:
                    return self.format_enum_value(value, enum_type)
                return str(value)

        if field.types[0] == GGUFValueType.ARRAY:
            array_values = self.extract_array_values(field)
            render_element = min(5, len(array_values))

            # Get enum type for this array if applicable
            enum_type = self.get_enum_for_key(field.name)

            if enum_type is not None:
                array_elements = []
                for i in range(render_element):
                    array_elements.append(self.format_enum_value(array_values[i], enum_type))
            else:
                array_elements = [str(array_values[i]) for i in range(render_element)]

            return f"[ {', '.join(array_elements).strip()}{', ...' if len(array_values) > len(array_elements) else ''} ]"

        return "Complex value"

    def load_tensors(self):
        self.tensors_table.setRowCount(0)

        if not self.reader:
            return

        for i, tensor in enumerate(self.reader.tensors):
            self.tensors_table.insertRow(i)

            # Name
            name_item = QTableWidgetItem(tensor.name)
            name_item.setFlags(name_item.flags() & ~Qt.ItemFlag.ItemIsEditable)
            self.tensors_table.setItem(i, 0, name_item)

            # Type
            type_item = QTableWidgetItem(tensor.tensor_type.name)
            type_item.setFlags(type_item.flags() & ~Qt.ItemFlag.ItemIsEditable)
            self.tensors_table.setItem(i, 1, type_item)

            # Shape
            shape_str = " × ".join(str(d) for d in tensor.shape)
            shape_item = QTableWidgetItem(shape_str)
            shape_item.setFlags(shape_item.flags() & ~Qt.ItemFlag.ItemIsEditable)
            self.tensors_table.setItem(i, 2, shape_item)

            # Elements
            elements_item = QTableWidgetItem(str(tensor.n_elements))
            elements_item.setFlags(elements_item.flags() & ~Qt.ItemFlag.ItemIsEditable)
            self.tensors_table.setItem(i, 3, elements_item)

            # Size
            size_item = QTableWidgetItem(f"{tensor.n_bytes:,}")
            size_item.setFlags(size_item.flags() & ~Qt.ItemFlag.ItemIsEditable)
            self.tensors_table.setItem(i, 4, size_item)

    def on_metadata_changed(self, item):
        if item.column() != 2:  # Only handle value column changes
            return

        row = item.row()
        orig_item = self.metadata_table.item(row, 0)
        key = None
        if orig_item:
            key = orig_item.text()
        new_value = item.text()

        field = None
        if self.reader and key:
            field = self.reader.get_field(key)
        if not field or not field.types or not key:
            return

        value_type = field.types[0]

        # Check if this is an enum field
        enum_type = self.get_enum_for_key(key)
        if enum_type is not None and value_type == GGUFValueType.INT32:
            # Try to parse the enum value from the text
            try:
                # Check if it's a name
                try:
                    enum_val = enum_type[new_value]
                    converted_value = enum_val.value
                except (KeyError, AttributeError):
                    # Check if it's a number or "NAME (value)" format
                    if '(' in new_value and ')' in new_value:
                        # Extract the value from "NAME (value)" format
                        value_part = new_value.split('(')[1].split(')')[0].strip()
                        converted_value = int(value_part)
                    else:
                        # Try to convert directly to int
                        converted_value = int(new_value)

                # Validate that it's a valid enum value
                enum_type(converted_value)

                # Store the change
                self.metadata_changes[key] = (value_type, converted_value)
                self.modified = True

                # Update display with formatted enum value
                formatted_value = self.format_enum_value(converted_value, enum_type)
                item.setText(formatted_value)

                self.statusBar().showMessage(f"Changed {key} to {formatted_value}")
                return
            except (ValueError, KeyError) as e:
                QMessageBox.warning(
                    self,
                    f"Invalid Enum Value ({e})",
                    f"'{new_value}' is not a valid {enum_type.__name__} value.\n"
                    f"Valid values are: {', '.join(v.name for v in enum_type)}")

                # Revert to original value
                original_value = self.format_field_value(field)
                item.setText(original_value)
                return

        try:
            # Convert the string value to the appropriate type
            if value_type == GGUFValueType.UINT8:
                converted_value = np.uint8(int(new_value))
            elif value_type == GGUFValueType.INT8:
                converted_value = np.int8(int(new_value))
            elif value_type == GGUFValueType.UINT16:
                converted_value = np.uint16(int(new_value))
            elif value_type == GGUFValueType.INT16:
                converted_value = np.int16(int(new_value))
            elif value_type == GGUFValueType.UINT32:
                converted_value = np.uint32(int(new_value))
            elif value_type == GGUFValueType.INT32:
                converted_value = np.int32(int(new_value))
            elif value_type == GGUFValueType.FLOAT32:
                converted_value = np.float32(float(new_value))
            elif value_type == GGUFValueType.BOOL:
                converted_value = new_value.lower() in ('true', 'yes', '1')
            elif value_type == GGUFValueType.STRING:
                converted_value = new_value
            else:
                # Unsupported type for editing
                return

            # Store the change
            self.metadata_changes[key] = (value_type, converted_value)
            self.modified = True

            self.statusBar().showMessage(f"Changed {key} to {new_value}")
        except ValueError:
            QMessageBox.warning(self, "Invalid Value", f"The value '{new_value}' is not valid for type {value_type.name}")

            # Revert to original value
            original_value = self.format_field_value(field)
            item.setText(original_value)

    def remove_metadata(self):
        button = self.sender()
        key = button.property("key")
        row = button.property("row")

        reply = QMessageBox.question(
            self, "Confirm Removal",
            f"Are you sure you want to remove the metadata key '{key}'?",
            QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No, QMessageBox.StandardButton.No
        )

        if reply == QMessageBox.StandardButton.Yes:
            self.metadata_table.removeRow(row)
            self.metadata_to_remove.add(key)

            # If we previously had changes for this key, remove them
            if key in self.metadata_changes:
                del self.metadata_changes[key]

            self.modified = True
            self.statusBar().showMessage(f"Marked {key} for removal")

    def edit_metadata_enum(self):
        """Edit an enum metadata field."""
        button = self.sender()
        key = button.property("key")
        row = button.property("row")

        field = None
        if self.reader:
            field = self.reader.get_field(key)
        if not field or not field.types:
            return

        enum_type = self.get_enum_for_key(key)
        if enum_type is None:
            return

        # Get current value
        current_value = field.contents()

        # Create a dialog with enum options
        dialog = QDialog(self)
        dialog.setWindowTitle(f"Select {enum_type.__name__} Value")
        layout = QVBoxLayout(dialog)

        combo = QComboBox()
        for enum_val in enum_type:
            combo.addItem(f"{enum_val.name} ({enum_val.value})", enum_val.value)

        # Set current value
        try:
            if isinstance(current_value, (int, str)):
                enum_val = enum_type(current_value)
                combo.setCurrentText(f"{enum_val.name} ({current_value})")
        except (ValueError, KeyError):
            pass

        layout.addWidget(combo)

        buttons = QDialogButtonBox(QDialogButtonBox.StandardButton.Ok | QDialogButtonBox.StandardButton.Cancel)
        buttons.accepted.connect(dialog.accept)
        buttons.rejected.connect(dialog.reject)
        layout.addWidget(buttons)

        if dialog.exec() == QDialog.DialogCode.Accepted:
            # Get the selected value
            new_value = combo.currentData()
            enum_val = enum_type(new_value)

            # Store the change
            self.metadata_changes[key] = (field.types[0], new_value)
            self.modified = True

            # Update display
            display_text = f"{enum_val.name} ({new_value})"
            target_item = self.metadata_table.item(row, 2)
            if target_item:
                target_item.setText(display_text)

            self.statusBar().showMessage(f"Changed {key} to {display_text}")

    def edit_array_metadata(self):
        button = self.sender()
        key = button.property("key")
        row = button.property("row")

        # Check if this is one of the linked tokenizer keys
        if key in TOKENIZER_LINKED_KEYS:
            self.edit_tokenizer_metadata(key)
            return

        field = None
        if self.reader:
            field = self.reader.get_field(key)
        if not field or not field.types or field.types[0] != GGUFValueType.ARRAY:
            return

        # Get array element type
        element_type = field.types[1]

        # Extract array values
        array_values = self.extract_array_values(field)

        # Open array editor dialog
        dialog = ArrayEditorDialog(array_values, element_type, key, self)
        if dialog.exec() == QDialog.DialogCode.Accepted:
            new_values = dialog.get_array_values()

            # Store the change
            self.metadata_changes[key] = (GGUFValueType.ARRAY, (element_type, new_values))
            self.modified = True

            # Update display
            enum_type = self.get_enum_for_key(key)
            if enum_type is not None and element_type == GGUFValueType.INT32:
                value_str = f"[ {', '.join(self.format_enum_value(v, enum_type) for v in new_values[:5])}{', ...' if len(new_values) > 5 else ''} ]"
            else:
                value_str = f"[ {', '.join(str(v) for v in new_values[:5])}{', ...' if len(new_values) > 5 else ''} ]"
            target_item = self.metadata_table.item(row, 2)
            if target_item:
                target_item.setText(value_str)

            self.statusBar().showMessage(f"Updated array values for {key}")

    def edit_tokenizer_metadata(self, trigger_key):
        """Edit the linked tokenizer metadata arrays together."""
        if not self.reader:
            return

        # Get all three fields
        tokens_field = self.reader.get_field(Keys.Tokenizer.LIST)
        token_types_field = self.reader.get_field(Keys.Tokenizer.TOKEN_TYPE)
        scores_field = self.reader.get_field(Keys.Tokenizer.SCORES)

        # Extract values from each field
        tokens = self.extract_array_values(tokens_field) if tokens_field else []
        token_types = self.extract_array_values(token_types_field) if token_types_field else []
        scores = self.extract_array_values(scores_field) if scores_field else []

        # Apply any pending changes
        if Keys.Tokenizer.LIST in self.metadata_changes:
            _, (_, tokens) = self.metadata_changes[Keys.Tokenizer.LIST]
        if Keys.Tokenizer.TOKEN_TYPE in self.metadata_changes:
            _, (_, token_types) = self.metadata_changes[Keys.Tokenizer.TOKEN_TYPE]
        if Keys.Tokenizer.SCORES in self.metadata_changes:
            _, (_, scores) = self.metadata_changes[Keys.Tokenizer.SCORES]

        # Open the tokenizer editor dialog
        dialog = TokenizerEditorDialog(tokens, token_types, scores, self)
        if dialog.exec() == QDialog.DialogCode.Accepted:
            new_tokens, new_token_types, new_scores = dialog.get_data()

            # Store changes for all three arrays
            if tokens_field:
                self.metadata_changes[Keys.Tokenizer.LIST] = (
                    GGUFValueType.ARRAY,
                    (tokens_field.types[1], new_tokens)
                )

            if token_types_field:
                self.metadata_changes[Keys.Tokenizer.TOKEN_TYPE] = (
                    GGUFValueType.ARRAY,
                    (token_types_field.types[1], new_token_types)
                )

            if scores_field:
                self.metadata_changes[Keys.Tokenizer.SCORES] = (
                    GGUFValueType.ARRAY,
                    (scores_field.types[1], new_scores)
                )

            self.modified = True

            # Update display for all three fields
            self.update_tokenizer_display(Keys.Tokenizer.LIST, new_tokens)
            self.update_tokenizer_display(Keys.Tokenizer.TOKEN_TYPE, new_token_types)
            self.update_tokenizer_display(Keys.Tokenizer.SCORES, new_scores)

            self.statusBar().showMessage("Updated tokenizer data")

    def update_tokenizer_display(self, key, values):
        """Update the display of a tokenizer field in the metadata table."""
        for row in range(self.metadata_table.rowCount()):
            key_item = self.metadata_table.item(row, 0)
            if key_item and key_item.text() == key:
                value_str = f"[ {', '.join(str(v) for v in values[:5])}{', ...' if len(values) > 5 else ''} ]"
                value_item = self.metadata_table.item(row, 2)
                if value_item:
                    value_item.setText(value_str)
                break

    def add_metadata(self):
        dialog = AddMetadataDialog(self)
        if dialog.exec() == QDialog.DialogCode.Accepted:
            key, value_type, value = dialog.get_data()

            if not key:
                QMessageBox.warning(self, "Invalid Key", "Key cannot be empty")
                return

            # Check if key already exists
            for row in range(self.metadata_table.rowCount()):
                orig_item = self.metadata_table.item(row, 0)
                if orig_item and orig_item.text() == key:
                    QMessageBox.warning(self, "Duplicate Key", f"Key '{key}' already exists")
                    return

            # Add to table
            row = self.metadata_table.rowCount()
            self.metadata_table.insertRow(row)

            # Key
            key_item = QTableWidgetItem(key)
            key_item.setFlags(key_item.flags() & ~Qt.ItemFlag.ItemIsEditable)
            self.metadata_table.setItem(row, 0, key_item)

            # Type
            type_item = QTableWidgetItem(value_type.name)
            type_item.setFlags(type_item.flags() & ~Qt.ItemFlag.ItemIsEditable)
            self.metadata_table.setItem(row, 1, type_item)

            # Value
            value_item = QTableWidgetItem(str(value))
            value_item.setFlags(value_item.flags() | Qt.ItemFlag.ItemIsEditable)
            self.metadata_table.setItem(row, 2, value_item)

            # Actions
            actions_widget = QWidget()
            actions_layout = QHBoxLayout(actions_widget)
            actions_layout.setContentsMargins(2, 2, 2, 2)

            remove_button = QPushButton("Remove")
            remove_button.setProperty("row", row)
            remove_button.setProperty("key", key)
            remove_button.clicked.connect(self.remove_metadata)
            actions_layout.addWidget(remove_button)

            self.metadata_table.setCellWidget(row, 3, actions_widget)

            # Store the change
            self.metadata_changes[key] = (value_type, value)
            self.modified = True

            self.statusBar().showMessage(f"Added new metadata key {key}")

    def save_file(self):
        if not self.reader:
            QMessageBox.warning(self, "No File Open", "Please open a GGUF file first")
            return

        if not self.modified and not self.metadata_changes and not self.metadata_to_remove:
            QMessageBox.information(self, "No Changes", "No changes to save")
            return

        file_path, _ = QFileDialog.getSaveFileName(
            self, "Save GGUF File As", "", "GGUF Files (*.gguf);;All Files (*)"
        )

        if not file_path:
            return

        try:
            self.statusBar().showMessage(f"Saving to {file_path}...")
            QApplication.processEvents()

            # Get architecture and endianness from the original file
            arch = 'unknown'
            field = self.reader.get_field(Keys.General.ARCHITECTURE)
            if field:
                arch = field.contents()

            # Create writer
            writer = GGUFWriter(file_path, arch=arch, endianess=self.reader.endianess)

            # Get alignment if present
            alignment = None
            field = self.reader.get_field(Keys.General.ALIGNMENT)
            if field:
                alignment = field.contents()
                if alignment is not None:
                    writer.data_alignment = alignment

            # Copy metadata with changes
            for field in self.reader.fields.values():
                # Skip virtual fields and fields written by GGUFWriter
                if field.name == Keys.General.ARCHITECTURE or field.name.startswith('GGUF.'):
                    continue

                # Skip fields marked for removal
                if field.name in self.metadata_to_remove:
                    continue

                # Apply changes if any
                sub_type = None
                if field.name in self.metadata_changes:
                    value_type, value = self.metadata_changes[field.name]
                    if value_type == GGUFValueType.ARRAY:
                        # Handle array values
                        sub_type, value = value
                else:
                    # Copy original value
                    value = field.contents()
                    value_type = field.types[0]
                    if value_type == GGUFValueType.ARRAY:
                        sub_type = field.types[-1]

                if value is not None:
                    writer.add_key_value(field.name, value, value_type, sub_type=sub_type)

            # Add new metadata
            for key, (value_type, value) in self.metadata_changes.items():
                # Skip if the key already existed (we handled it above)
                if self.reader.get_field(key) is not None:
                    continue

                sub_type = None
                if value_type == GGUFValueType.ARRAY:
                    # Handle array values
                    sub_type, value = value

                writer.add_key_value(key, value, value_type, sub_type=sub_type)

            # Add tensors (including data)
            for tensor in self.reader.tensors:
                writer.add_tensor(tensor.name, tensor.data, raw_shape=tensor.data.shape, raw_dtype=tensor.tensor_type)

            # Write header and metadata
            writer.open_output_file(Path(file_path))
            writer.write_header_to_file()
            writer.write_kv_data_to_file()

            # Write tensor data using the optimized method
            writer.write_tensors_to_file(progress=False)

            writer.close()

            self.statusBar().showMessage(f"Saved to {file_path}")

            # Ask if user wants to open the new file
            reply = QMessageBox.question(
                self, "Open Saved File",
                "Would you like to open the newly saved file?",
                QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No, QMessageBox.StandardButton.Yes
            )

            if reply == QMessageBox.StandardButton.Yes:
                self.reader = GGUFReader(file_path, 'r')
                self.current_file = file_path
                self.file_path_edit.setText(file_path)

                self.load_metadata()
                self.load_tensors()

                self.metadata_changes = {}
                self.metadata_to_remove = set()
                self.modified = False

        except Exception as e:
            QMessageBox.critical(self, "Error", f"Failed to save file: {str(e)}")
            self.statusBar().showMessage("Error saving file")

def main() -> None:
    parser = argparse.ArgumentParser(description="GUI GGUF Editor")
    parser.add_argument("model_path", nargs="?", help="path to GGUF model file to load at startup")
    parser.add_argument("--verbose", action="store_true", help="increase output verbosity")
    args = parser.parse_args()
    logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO)
    app = QApplication(sys.argv)
    window = GGUFEditorWindow()
    window.show()
    # Load model if specified
    if args.model_path:
        if os.path.isfile(args.model_path) and args.model_path.endswith('.gguf'):
            window.load_file(args.model_path)
        else:
            logger.error(f"Invalid model path: {args.model_path}")
            QMessageBox.warning(
                window,
                "Invalid Model Path",
                f"The specified file does not exist or is not a GGUF file: {args.model_path}")
    sys.exit(app.exec())

if __name__ == '__main__':
    main()