"""Pixel metadata editor view for PyQt6 GUI."""

import logging
from typing import Optional, Dict
from pathlib import Path

from PyQt6.QtWidgets import (
    QWidget, QVBoxLayout, QHBoxLayout, QGroupBox, QLabel,
    QLineEdit, QComboBox, QPushButton, QTableWidget, QTableWidgetItem,
    QMessageBox, QScrollArea
)
from PyQt6.QtCore import Qt

from eubi_bridge.ngff.multiscales import Pyramid, NGFFMetadataHandler
from eubi_bridge.utils.path_utils import is_ome_zarr
from eubi_bridge.gui.workers import run_async
from eubi_bridge.gui.state import MetadataEditorState
from eubi_bridge.gui.widgets import LoadingIndicator

logger = logging.getLogger(__name__)


class PixelMetadataEditor(QWidget):
    """Editor for OME-Zarr pixel metadata (scales, units)."""
    
    # Standard unit options
    SPACE_UNITS = [
        "angstrom", "attometer", "centimeter", "decameter", "decimeter",
        "exameter", "femtometer", "foot", "gigameter", "hectometer",
        "inch", "kilometer", "megameter", "meter", "micrometer",
        "mile", "millimeter", "nanometer", "parsec", "petameter",
        "picometer", "terameter", "thou", "yard", "yoctometer",
        "yottameter", "zeptometer", "zettameter"
    ]
    
    TIME_UNITS = [
        "attosecond", "centisecond", "day", "decisecond", "exasecond",
        "femtosecond", "gigasecond", "hectosecond", "hour", "kilosecond",
        "megasecond", "microsecond", "millisecond", "minute", "nanosecond",
        "petasecond", "picosecond", "second", "terasecond", "yoctosecond",
        "yottasecond", "zeptosecond", "zettasecond"
    ]
    
    def __init__(self, bridge, parent=None):
        super().__init__(parent)
        
        self.bridge = bridge
        self.state = MetadataEditorState()
        
        # Metadata objects
        self.pyramid: Optional[Pyramid] = None
        self.handler: Optional[NGFFMetadataHandler] = None
        
        self._setup_ui()
    
    def _setup_ui(self):
        """Build UI layout."""
        main_layout = QVBoxLayout(self)
        
        # File selection
        main_layout.addWidget(self._create_file_selection_group())
        
        # Loading indicator
        self.loading_indicator = LoadingIndicator()
        self.loading_indicator.hide()
        main_layout.addWidget(self.loading_indicator)
        
        # Metadata display
        scroll = QScrollArea()
        scroll.setWidgetResizable(True)
        scroll_widget = QWidget()
        scroll_layout = QVBoxLayout(scroll_widget)
        
        scroll_layout.addWidget(self._create_omezarr_info_group())
        scroll_layout.addWidget(self._create_axis_metadata_group())
        scroll_layout.addWidget(self._create_storage_info_group())
        scroll_layout.addStretch()
        
        scroll.setWidget(scroll_widget)
        main_layout.addWidget(scroll, 1)
        
        # Action buttons
        main_layout.addWidget(self._create_action_buttons_group())
        
        self.setLayout(main_layout)
    
    def _create_file_selection_group(self) -> QGroupBox:
        """Create file selection section."""
        group = QGroupBox("Load OME-Zarr")
        layout = QHBoxLayout()
        
        self.path_input = QLineEdit()
        self.path_input.setPlaceholderText("Enter path or click Browse...")
        layout.addWidget(self.path_input)
        
        browse_btn = QPushButton("Browse...")
        browse_btn.clicked.connect(self._on_browse_clicked)
        layout.addWidget(browse_btn)
        
        load_btn = QPushButton("Load")
        load_btn.clicked.connect(self._on_load_clicked)
        layout.addWidget(load_btn)
        
        group.setLayout(layout)
        return group
    
    def _create_omezarr_info_group(self) -> QGroupBox:
        """Create OME-Zarr info display."""
        group = QGroupBox("OME-Zarr Information")
        layout = QVBoxLayout()
        
        # Version
        version_layout = QHBoxLayout()
        version_layout.addWidget(QLabel("NGFF Version:"))
        self.version_label = QLabel("Not loaded")
        version_layout.addWidget(self.version_label)
        version_layout.addStretch()
        layout.addLayout(version_layout)
        
        # Layers
        layers_layout = QHBoxLayout()
        layers_layout.addWidget(QLabel("Layers:"))
        self.layers_label = QLabel("0")
        layers_layout.addWidget(self.layers_label)
        layers_layout.addStretch()
        layout.addLayout(layers_layout)
        
        # Shape
        shape_layout = QHBoxLayout()
        shape_layout.addWidget(QLabel("Shape:"))
        self.shape_label = QLabel("Not loaded")
        self.shape_label.setWordWrap(True)
        shape_layout.addWidget(self.shape_label)
        layout.addLayout(shape_layout)
        
        # Data type
        dtype_layout = QHBoxLayout()
        dtype_layout.addWidget(QLabel("Data Type:"))
        self.dtype_label = QLabel("Not loaded")
        dtype_layout.addWidget(self.dtype_label)
        dtype_layout.addStretch()
        layout.addLayout(dtype_layout)
        
        group.setLayout(layout)
        return group
    
    def _create_axis_metadata_group(self) -> QGroupBox:
        """Create axis metadata editing section."""
        group = QGroupBox("Axis Metadata")
        layout = QVBoxLayout()
        
        # Axis table
        self.axis_table = QTableWidget()
        self.axis_table.setColumnCount(5)
        self.axis_table.setHorizontalHeaderLabels([
            "Axis", "Scale", "Unit", "Size", "Type"
        ])
        self.axis_table.horizontalHeader().setStretchLastSection(True)
        layout.addWidget(self.axis_table)
        
        group.setLayout(layout)
        return group
    
    def _create_storage_info_group(self) -> QGroupBox:
        """Create storage information display."""
        group = QGroupBox("Storage Information")
        layout = QVBoxLayout()
        
        # Chunks
        chunks_layout = QHBoxLayout()
        chunks_layout.addWidget(QLabel("Chunk Size:"))
        self.chunks_label = QLabel("Not loaded")
        chunks_layout.addWidget(self.chunks_label)
        chunks_layout.addStretch()
        layout.addLayout(chunks_layout)
        
        # Compression
        compression_layout = QHBoxLayout()
        compression_layout.addWidget(QLabel("Compression:"))
        self.compression_label = QLabel("Not loaded")
        compression_layout.addWidget(self.compression_label)
        compression_layout.addStretch()
        layout.addLayout(compression_layout)
        
        # Disk size
        size_layout = QHBoxLayout()
        size_layout.addWidget(QLabel("Disk Size:"))
        self.size_label = QLabel("Not loaded")
        size_layout.addWidget(self.size_label)
        size_layout.addStretch()
        layout.addLayout(size_layout)
        
        group.setLayout(layout)
        return group
    
    def _create_action_buttons_group(self) -> QGroupBox:
        """Create action buttons."""
        group = QGroupBox("Actions")
        layout = QHBoxLayout()
        
        update_btn = QPushButton("Update Metadata")
        update_btn.clicked.connect(self._on_update_metadata)
        layout.addWidget(update_btn)
        
        reset_btn = QPushButton("Reset Changes")
        reset_btn.clicked.connect(self._on_reset_changes)
        layout.addWidget(reset_btn)
        
        layout.addStretch()
        group.setLayout(layout)
        return group
    
    def _on_browse_clicked(self):
        """Handle browse button."""
        from PyQt6.QtWidgets import QFileDialog
        
        path = QFileDialog.getExistingDirectory(self, "Select OME-Zarr Folder")
        if path:
            self.path_input.setText(path)
    
    def _on_load_clicked(self):
        """Handle load button."""
        path = self.path_input.text().strip()
        if not path:
            QMessageBox.warning(self, "Warning", "Please enter a path")
            return
        
        if not Path(path).exists():
            QMessageBox.warning(self, "Warning", "Path does not exist")
            return
        
        if not is_ome_zarr(path):
            QMessageBox.warning(self, "Warning", "Not a valid OME-Zarr file")
            return
        
        self._load_metadata(path)
    
    def _load_metadata(self, path: str):
        """Load metadata from OME-Zarr file."""
        self.loading_indicator.set_message("Loading metadata...")
        self.loading_indicator.start()
        
        def load():
            try:
                pyr = Pyramid(path)
                handler = NGFFMetadataHandler(path)
                return pyr, handler, None
            except Exception as e:
                logger.error(f"Error loading metadata: {e}")
                return None, None, str(e)
        
        def on_result(result):
            pyr, handler, error = result
            
            if error:
                QMessageBox.critical(self, "Error", f"Failed to load: {error}")
                self.loading_indicator.stop()
                return
            
            self.pyramid = pyr
            self.handler = handler
            self.state.current_zarr_path = path
            self._update_ui_from_metadata()
            self.loading_indicator.stop()
        
        run_async(load, on_result=on_result)
    
    def _update_ui_from_metadata(self):
        """Update UI from loaded metadata."""
        if not self.pyramid:
            return
        
        # Update info display
        self.version_label.setText(self.pyramid.version or "Unknown")
        self.layers_label.setText(str(len(self.pyramid.datasets)))
        self.shape_label.setText(str(self.pyramid.shape))
        self.dtype_label.setText(str(self.pyramid.dtype))
        
        # Update axis table
        self._populate_axis_table()
        
        # Update storage info
        self._populate_storage_info()
    
    def _populate_axis_table(self):
        """Populate axis metadata table."""
        self.axis_table.setRowCount(0)
        
        if not self.pyramid or not self.pyramid.axes:
            return
        
        axes = self.pyramid.axes
        shape = self.pyramid.shape
        
        for i, axis_name in enumerate(axes):
            row = self.axis_table.rowCount()
            self.axis_table.insertRow(row)
            
            # Axis name - axes is a string, so axis_name is a character
            name_item = QTableWidgetItem(axis_name or f"Dim {i}")
            name_item.setFlags(name_item.flags() & ~Qt.ItemFlag.ItemIsEditable)
            self.axis_table.setItem(row, 0, name_item)
            
            # Scale - get from metadata
            scale = pyramid.meta.scaledict.get('0', {}).get(axis_name, 1.0)
            scale_item = QLineEdit(str(scale))
            self.axis_table.setCellWidget(row, 1, scale_item)
            
            # Unit - get from metadata
            unit = pyramid.meta.unit_dict.get(axis_name, "")
            unit_combo = QComboBox()
            
            if axis_name.upper() in ['X', 'Y', 'Z']:
                unit_combo.addItems(self.SPACE_UNITS)
            elif axis_name.upper() == 'T':
                unit_combo.addItems(self.TIME_UNITS)
            else:
                unit_combo.addItems(["", "unknown"])
            
            if unit and unit in unit_combo.model():
                unit_combo.setCurrentText(unit)
            
            self.axis_table.setCellWidget(row, 2, unit_combo)
            
            # Size
            size_item = QTableWidgetItem(str(shape[i]))
            size_item.setFlags(size_item.flags() & ~Qt.ItemFlag.ItemIsEditable)
            self.axis_table.setItem(row, 3, size_item)
            
            # Type
            type_item = QTableWidgetItem(axis.type or "")
            type_item.setFlags(type_item.flags() & ~Qt.ItemFlag.ItemIsEditable)
            self.axis_table.setItem(row, 4, type_item)
    
    def _populate_storage_info(self):
        """Populate storage information."""
        if not self.pyramid:
            return
        
        try:
            # Get first dataset for info
            dataset = self.pyramid.datasets[0]
            
            # Chunks
            chunks = dataset.chunks if hasattr(dataset, 'chunks') else "Unknown"
            self.chunks_label.setText(str(chunks))
            
            # Compression
            compressor = dataset.compressor if hasattr(dataset, 'compressor') else "Unknown"
            self.compression_label.setText(str(compressor))
            
            # Size (estimate)
            import zarr
            root = zarr.open(self.state.current_zarr_path)
            size_bytes = 0
            
            def count_size(group):
                nonlocal size_bytes
                for key in group:
                    item = group[key]
                    if hasattr(item, 'nbytes'):
                        size_bytes += item.nbytes
                    elif hasattr(item, '__iter__'):
                        count_size(item)
            
            count_size(root)
            size_mb = size_bytes / (1024 * 1024)
            self.size_label.setText(f"{size_mb:.2f} MB")
        
        except Exception as e:
            logger.error(f"Error reading storage info: {e}")
            self.chunks_label.setText("Error reading")
    
    def _on_update_metadata(self):
        """Update metadata in OME-Zarr file."""
        if not self.pyramid:
            QMessageBox.warning(self, "Warning", "No metadata loaded")
            return
        
        # Collect changes from table
        updates = {}
        
        for row in range(self.axis_table.rowCount()):
            # Get scale
            scale_widget = self.axis_table.cellWidget(row, 1)
            if isinstance(scale_widget, QLineEdit):
                try:
                    scale = float(scale_widget.text())
                    updates[f"axis_{row}_scale"] = scale
                except ValueError:
                    pass
            
            # Get unit
            unit_widget = self.axis_table.cellWidget(row, 2)
            if isinstance(unit_widget, QComboBox):
                unit = unit_widget.currentText()
                if unit:
                    updates[f"axis_{row}_unit"] = unit
        
        if not updates:
            QMessageBox.information(self, "Info", "No changes to save")
            return
        
        self.state.pending_changes = updates
        
        self.loading_indicator.set_message("Updating metadata...")
        self.loading_indicator.start()
        
        def update():
            try:
                # Apply updates to OME-Zarr
                # This would call the actual update function
                return True, None
            except Exception as e:
                logger.error(f"Error updating metadata: {e}")
                return False, str(e)
        
        def on_result(result):
            success, error = result
            self.loading_indicator.stop()
            
            if success:
                QMessageBox.information(self, "Success", "Metadata updated successfully")
                self._load_metadata(self.state.current_zarr_path)
            else:
                QMessageBox.critical(self, "Error", f"Failed to update: {error}")
        
        run_async(update, on_result=on_result)
    
    def _on_reset_changes(self):
        """Reset pending changes."""
        self.state.pending_changes.clear()
        if self.state.current_zarr_path:
            self._load_metadata(self.state.current_zarr_path)
