"""
Generic ZMQ Server Manager Widget for PyQt6.

Provides a reusable UI component for managing any ZMQ server (execution servers,
Napari viewers, future servers) using the ZMQServer/ZMQClient ABC interface.

Features:
- Auto-discovery of running servers via port scanning
- Display server info (port, type, status, log file)
- Graceful shutdown and force kill
- Double-click to open log files
- Works with ANY ZMQServer subclass
- Tracks launching viewers with queued image counts
"""

import logging
from pathlib import Path
from typing import List, Dict, Any, Optional
from PyQt6.QtWidgets import (
    QWidget,
    QVBoxLayout,
    QHBoxLayout,
    QTreeWidget,
    QTreeWidgetItem,
    QPushButton,
    QGroupBox,
    QMessageBox,
    QAbstractItemView,
    QLabel,
    QSplitter,
    QFrame,
)
from PyQt6.QtCore import Qt, pyqtSignal, pyqtSlot, QTimer
from PyQt6.QtGui import QFont
from pyqt_reactive.theming import StyleSheetGenerator
import threading
from objectstate import spawn_thread_with_context
from zmqruntime.viewer_state import ViewerStateManager, ViewerState

logger = logging.getLogger(__name__)


# Global reference to active ZMQ server manager widgets (for triggering refreshes)
_active_managers_lock = threading.Lock()
_active_managers: List["ZMQServerManagerWidget"] = []


class ZMQServerManagerWidget(QWidget):
    """
    Generic ZMQ server manager widget.

    Works with any ZMQServer subclass via the ABC interface.
    Displays running servers and provides management controls.
    """

    # Signals
    server_killed = pyqtSignal(int)  # Emitted when server is killed (port)
    log_file_opened = pyqtSignal(str)  # Emitted when log file is opened (path)
    _scan_complete = pyqtSignal(list)  # Internal signal for async scan completion
    _kill_complete = pyqtSignal(
        bool, str
    )  # Internal signal for async kill completion (success, message)

    def __init__(
        self,
        ports_to_scan: List[int],
        title: str = "ZMQ Servers",
        style_generator: Optional[StyleSheetGenerator] = None,
        parent: Optional[QWidget] = None,
    ):
        """
        Initialize ZMQ server manager widget.

        Args:
            ports_to_scan: List of ports to scan for servers
            title: Title for the group box
            style_generator: Style generator for consistent styling
            parent: Parent widget
        """
        super().__init__(parent)

        self.ports_to_scan = ports_to_scan
        self.title = title
        self.style_generator = style_generator

        # Server tracking
        self.servers: List[Dict[str, Any]] = []

        # Register this manager for launching viewer updates
        with _active_managers_lock:
            _active_managers.append(self)

        # Register a ViewerStateManager callback so the UI updates immediately
        # when viewer state changes. This is required for the hard migration.
        def _manager_callback(instance):
            # Forward to the UI thread using invokeMethod
            try:
                from PyQt6.QtCore import QMetaObject, Qt

                QMetaObject.invokeMethod(
                    self,
                    "_refresh_launching_viewers_only",
                    Qt.ConnectionType.QueuedConnection,
                )
            except Exception:
                pass

        mgr = ViewerStateManager.get_instance()
        mgr.register_state_callback(_manager_callback)
        self._viewer_state_callback = _manager_callback

        # Connect internal signal for async scanning
        self._scan_complete.connect(self._update_server_list)

        # Auto-refresh timer (async scanning won't block UI)
        self.refresh_timer = QTimer()
        self.refresh_timer.timeout.connect(self.refresh_servers)

        # Cleanup flag to prevent operations after cleanup
        self._is_cleaning_up = False

        self.setup_ui()

    def cleanup(self):
        """Cleanup resources before widget destruction."""
        if self._is_cleaning_up:
            return

        self._is_cleaning_up = True

        # Stop refresh timer first to prevent new refresh calls
        if hasattr(self, "refresh_timer") and self.refresh_timer:
            self.refresh_timer.stop()
            self.refresh_timer.deleteLater()
            self.refresh_timer = None

        # Unregister this manager from global list
        with _active_managers_lock:
            if self in _active_managers:
                _active_managers.remove(self)

        # Unregister viewer state callback
        try:
            mgr = ViewerStateManager.get_instance()
            if getattr(self, "_viewer_state_callback", None):
                mgr.unregister_state_callback(self._viewer_state_callback)
        except Exception:
            # Hard migration: ViewerStateManager should be present; ignore failures
            pass

        logger.debug("ZMQServerManagerWidget cleanup completed")

    def __del__(self):
        """Cleanup when widget is destroyed."""
        self.cleanup()

    def showEvent(self, a0):
        """Auto-scan for servers when widget is shown."""
        super().showEvent(a0)
        if not self._is_cleaning_up:
            # Scan for servers on first show
            self.refresh_servers()
            # Start auto-refresh (1 second interval - async scanning won't block UI)
            if self.refresh_timer:
                self.refresh_timer.start(1000)

    def hideEvent(self, a0):
        """Stop auto-refresh when widget is hidden."""
        super().hideEvent(a0)
        # Stop timer to prevent unnecessary background work
        if hasattr(self, "refresh_timer") and self.refresh_timer:
            self.refresh_timer.stop()

    def setup_ui(self):
        """Setup the user interface with AbstractManagerWidget-like styling."""
        main_layout = QVBoxLayout(self)
        main_layout.setContentsMargins(2, 2, 2, 2)
        main_layout.setSpacing(2)

        # Header (title + status) - similar to AbstractManagerWidget
        header = self._create_header()
        main_layout.addWidget(header)

        # QSplitter: tree widget ABOVE buttons (VERTICAL orientation)
        splitter = QSplitter(Qt.Orientation.Vertical)

        # Top: server tree
        self.server_tree = QTreeWidget()
        self.server_tree.setHeaderLabels(["Server / Worker", "Status", "Info"])
        self.server_tree.setSelectionMode(
            QAbstractItemView.SelectionMode.ExtendedSelection
        )
        self.server_tree.itemDoubleClicked.connect(self._on_item_double_clicked)
        self.server_tree.setColumnWidth(0, 250)
        self.server_tree.setColumnWidth(1, 100)
        splitter.addWidget(self.server_tree)

        # Bottom: button panel
        button_panel = self._create_button_panel()
        splitter.addWidget(button_panel)

        # Set initial sizes: tree takes all space, buttons collapse to minimum height
        splitter.setSizes([1000, 1])
        splitter.setStretchFactor(0, 1)  # Tree widget expands
        splitter.setStretchFactor(1, 0)  # Button panel stays at minimum height

        main_layout.addWidget(splitter)

        # Apply styling
        if self.style_generator:
            cs = self.style_generator.color_scheme

            # Apply tree widget style
            self.server_tree.setStyleSheet(
                self.style_generator.generate_tree_widget_style()
            )

        # Connect internal signals
        self._scan_complete.connect(self._update_server_list)
        self._kill_complete.connect(self._on_kill_complete)

    def _create_header(self) -> QWidget:
        """Create header with title and status label (similar to AbstractManagerWidget)."""
        header = QWidget()
        header_layout = QHBoxLayout(header)
        header_layout.setContentsMargins(5, 5, 5, 5)

        # Title label
        title_label = QLabel(self.title)
        title_label.setFont(QFont("Arial", 12, QFont.Weight.Bold))
        if self.style_generator:
            title_label.setStyleSheet(
                f"color: {self.style_generator.color_scheme.to_hex(self.style_generator.color_scheme.text_accent)};"
            )
        header_layout.addWidget(title_label)

        header_layout.addStretch()

        # Status label
        self.status_label = QLabel("Ready")
        if self.style_generator:
            self.status_label.setStyleSheet(
                f"color: {self.style_generator.color_scheme.to_hex(self.style_generator.color_scheme.status_success)}; "
                f"font-weight: bold;"
            )
        header_layout.addWidget(self.status_label)

        return header

    def _create_button_panel(self) -> QWidget:
        """Create button panel with consistent styling."""
        panel = QWidget()
        layout = QHBoxLayout(panel)
        layout.setContentsMargins(5, 5, 5, 5)
        layout.setSpacing(5)

        self.refresh_btn = QPushButton("Refresh")
        self.refresh_btn.setToolTip("Refresh server list")
        self.refresh_btn.clicked.connect(self.refresh_servers)
        layout.addWidget(self.refresh_btn)

        self.quit_btn = QPushButton("Quit")
        self.quit_btn.setToolTip("Gracefully quit selected servers")
        self.quit_btn.clicked.connect(self.quit_selected_servers)
        layout.addWidget(self.quit_btn)

        self.force_kill_btn = QPushButton("Force Kill")
        self.force_kill_btn.setToolTip("Force kill selected servers")
        self.force_kill_btn.clicked.connect(self.force_kill_selected_servers)
        layout.addWidget(self.force_kill_btn)

        # Apply button styles
        if self.style_generator:
            button_style = self.style_generator.generate_button_style()
            self.refresh_btn.setStyleSheet(button_style)
            self.quit_btn.setStyleSheet(button_style)
            self.force_kill_btn.setStyleSheet(button_style)

        return panel

    def refresh_servers(self):
        """Scan ports and refresh server list (async in background)."""
        # Guard against calls after cleanup
        if self._is_cleaning_up:
            return

        def scan_and_update():
            """Background thread to scan ports without blocking UI."""
            import concurrent.futures

            # Scan ports in parallel using thread pool (like Napari implementation)
            servers = []

            with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
                # Submit all ping tasks
                future_to_port = {
                    executor.submit(self._ping_server, port): port
                    for port in self.ports_to_scan
                }

                # Collect results as they complete
                for future in concurrent.futures.as_completed(future_to_port):
                    port = future_to_port[future]
                    try:
                        server_info = future.result()
                        if server_info:
                            servers.append(server_info)
                    except Exception as e:
                        logger.debug(f"Error scanning port {port}: {e}")

            # Update UI on main thread via signal
            self._scan_complete.emit(servers)

        # Start scan in background thread
        spawn_thread_with_context(scan_and_update, name="scan_servers")

    def _ping_server(self, port: int) -> Optional[Dict[str, Any]]:
        """
        Ping a server on the given port and return its info.

        Returns server info dict if responsive, None otherwise.
        """
        import zmq
        import pickle
        from openhcs.constants.constants import CONTROL_PORT_OFFSET
        from openhcs.runtime.zmq_config import OPENHCS_ZMQ_CONFIG
        from zmqruntime.transport import (
            get_zmq_transport_url,
            get_default_transport_mode,
        )

        control_port = port + CONTROL_PORT_OFFSET
        control_context = None
        control_socket = None

        try:
            control_context = zmq.Context()
            control_socket = control_context.socket(zmq.REQ)
            control_socket.setsockopt(zmq.LINGER, 0)
            control_socket.setsockopt(
                zmq.RCVTIMEO, 300
            )  # 300ms timeout for fast scanning

            # Use transport mode-aware URL (IPC or TCP)
            transport_mode = get_default_transport_mode()
            control_url = get_zmq_transport_url(
                control_port,
                host="localhost",
                mode=transport_mode,
                config=OPENHCS_ZMQ_CONFIG,
            )
            control_socket.connect(control_url)

            # Send ping
            ping_message = {"type": "ping"}
            control_socket.send(pickle.dumps(ping_message))

            # Wait for pong
            response = control_socket.recv()
            response_data = pickle.loads(response)

            # Return server info if valid pong
            if response_data.get("type") == "pong":
                return response_data

            return None

        except Exception:
            return None
        finally:
            if control_socket:
                try:
                    control_socket.close()
                except:
                    pass
            if control_context:
                try:
                    control_context.term()
                except:
                    pass

    @pyqtSlot()
    def _refresh_launching_viewers_only(self):
        """Fast refresh: Update UI with launching viewers only (no port scan).

        This is called when launching viewer state changes and provides instant feedback.
        """
        # Guard against calls after cleanup
        if self._is_cleaning_up:
            return

        # Keep existing scanned servers, just update the tree display
        self._update_server_list(self.servers)

    @pyqtSlot(list)
    def _update_server_list(self, servers: List[Dict[str, Any]]):
        """Update server tree on UI thread (called via signal)."""
        from zmqruntime.queue_tracker import GlobalQueueTrackerRegistry

        self.servers = servers

        # Save current selection (by port) before clearing
        selected_ports = set()
        for item in self.server_tree.selectedItems():
            server_data = item.data(0, Qt.ItemDataRole.UserRole)
            if server_data and "port" in server_data:
                selected_ports.add(server_data["port"])

        self.server_tree.clear()

        # Get queue tracker registry for progress info
        registry = GlobalQueueTrackerRegistry()

        # First, add launching viewers
        mgr = ViewerStateManager.get_instance()
        launching_viewers = {
            v.port: {"type": v.viewer_type, "queued_images": v.queued_images}
            for v in mgr.list_viewers()
            if v.state == ViewerState.LAUNCHING
        }

        for port, info in launching_viewers.items():
            viewer_type = info["type"].capitalize()
            queued_images = info["queued_images"]

            display_text = f"Port {port} - {viewer_type} Viewer"
            status_text = "🚀 Launching"
            info_text = (
                f"{queued_images} images queued" if queued_images > 0 else "Starting..."
            )

            item = QTreeWidgetItem([display_text, status_text, info_text])
            item.setData(0, Qt.ItemDataRole.UserRole, {"port": port, "launching": True})
            self.server_tree.addTopLevelItem(item)

        # Add servers that are processing images (even if they didn't respond to ping)
        # This prevents busy servers from disappearing during image processing
        scanned_ports = {server.get("port") for server in servers}
        for tracker_port, tracker in registry.get_all_trackers().items():
            if tracker_port in scanned_ports or tracker_port in launching_viewers:
                continue  # Already in the list

            # Check if this tracker has pending images (server is busy processing)
            pending = tracker.get_pending_count()
            if pending > 0:
                # Server is busy processing - add it even though it didn't respond to ping
                processed, total = tracker.get_progress()
                viewer_type = tracker.viewer_type.capitalize()

                display_text = f"Port {tracker_port} - {viewer_type}ViewerServer"
                status_text = "⚙️"  # Busy icon
                info_text = f"Processing: {processed}/{total} images"

                # Check for stuck images
                if tracker.has_stuck_images():
                    status_text = "⚠️"
                    stuck_images = tracker.get_stuck_images()
                    info_text += f" (⚠️ {len(stuck_images)} stuck)"

                # Create pseudo-server dict for consistency
                pseudo_server = {
                    "port": tracker_port,
                    "server": f"{viewer_type}ViewerServer",
                    "ready": True,
                    "busy": True,  # Mark as busy
                }

                item = QTreeWidgetItem([display_text, status_text, info_text])
                item.setData(0, Qt.ItemDataRole.UserRole, pseudo_server)
                self.server_tree.addTopLevelItem(item)

        # Then add running servers
        for server in servers:
            port = server.get("port", "unknown")

            # Skip if this port is in launching registry (shouldn't happen, but be safe)
            if port in launching_viewers:
                continue

            server_type = server.get("server", "Unknown")
            ready = server.get("ready", False)

            # Determine status icon
            if ready:
                status_icon = "✅"
            else:
                status_icon = "🚀"

            # Handle execution servers specially - show workers as children
            if server_type.endswith("ExecutionServer"):
                running_executions = server.get("running_executions", [])
                workers = server.get("workers", [])

                # Create server item
                if running_executions:
                    server_text = f"Port {port} - Execution Server"
                    status_text = f"{status_icon} {len(running_executions)} exec"
                    info_text = f"{len(workers)} workers"
                else:
                    server_text = f"Port {port} - Execution Server"
                    status_text = f"{status_icon} Idle"
                    info_text = f"{len(workers)} workers" if workers else ""

                server_item = QTreeWidgetItem([server_text, status_text, info_text])
                server_item.setData(0, Qt.ItemDataRole.UserRole, server)
                self.server_tree.addTopLevelItem(server_item)

                # Add worker processes as children
                for worker in workers:
                    pid = worker.get("pid", "unknown")
                    status = worker.get("status", "unknown")
                    cpu = worker.get("cpu_percent", 0)
                    mem_mb = worker.get("memory_mb", 0)

                    worker_text = f"  Worker PID {pid}"
                    worker_status = f"⚙️ {status}"
                    worker_info = f"CPU: {cpu:.1f}% | Mem: {mem_mb:.0f}MB"

                    worker_item = QTreeWidgetItem(
                        [worker_text, worker_status, worker_info]
                    )
                    worker_item.setData(
                        0,
                        Qt.ItemDataRole.UserRole,
                        {"type": "worker", "pid": pid, "server": server},
                    )
                    server_item.addChild(worker_item)

                # Expand server item if it has workers
                if workers:
                    server_item.setExpanded(True)

            else:
                # Other server types (Napari, Fiji viewers) - show with progress if available
                display_text = f"Port {port} - {server_type}"
                status_text = status_icon
                info_text = ""

                # Check if this is a viewer with pending images
                tracker = registry.get_tracker(port)
                if tracker:
                    processed, total = tracker.get_progress()
                    pending = tracker.get_pending_count()

                    if pending > 0:
                        # Still processing images
                        info_text = f"Processing: {processed}/{total} images"

                        # Check for stuck images
                        if tracker.has_stuck_images():
                            status_text = "⚠️"  # Warning icon for stuck
                            stuck_images = tracker.get_stuck_images()
                            info_text += f" (⚠️ {len(stuck_images)} stuck)"
                    elif total > 0:
                        # All images processed
                        info_text = f"✅ Processed {total} images"

                # If no processing info, show memory usage
                if not info_text:
                    mem_mb = server.get("memory_mb")
                    cpu_percent = server.get("cpu_percent")
                    if mem_mb is not None:
                        info_text = f"Mem: {mem_mb:.0f}MB"
                        if cpu_percent is not None:
                            info_text += f" | CPU: {cpu_percent:.1f}%"

                item = QTreeWidgetItem([display_text, status_text, info_text])
                item.setData(0, Qt.ItemDataRole.UserRole, server)
                self.server_tree.addTopLevelItem(item)

        # Restore selection after refresh
        if selected_ports:
            for i in range(self.server_tree.topLevelItemCount()):
                item = self.server_tree.topLevelItem(i)
                server_data = item.data(0, Qt.ItemDataRole.UserRole)
                if server_data and server_data.get("port") in selected_ports:
                    item.setSelected(True)

        logger.debug(f"Found {len(servers)} ZMQ servers")

    @pyqtSlot(bool, str)
    def _on_kill_complete(self, success: bool, message: str):
        """Handle kill operation completion on UI thread."""
        if not success:
            QMessageBox.warning(self, "Kill Failed", message)
        # Refresh list after kill (quick refresh for better UX)
        QTimer.singleShot(200, self.refresh_servers)

    def quit_selected_servers(self):
        """Gracefully quit selected servers (async to avoid blocking UI)."""
        selected_items = self.server_tree.selectedItems()
        if not selected_items:
            QMessageBox.warning(self, "No Selection", "Please select servers to quit.")
            return

        # Collect ports to kill BEFORE showing dialog (items may be deleted by auto-refresh)
        ports_to_kill = []
        for item in selected_items:
            data = item.data(0, Qt.ItemDataRole.UserRole)
            # Skip worker items (they don't have ports)
            if data and data.get("type") == "worker":
                continue
            port = data.get("port") if data else None
            if port:
                ports_to_kill.append(port)

        if not ports_to_kill:
            QMessageBox.warning(
                self, "No Servers", "No servers selected (only workers selected)."
            )
            return

        # Confirm with user
        reply = QMessageBox.question(
            self,
            "Quit Confirmation",
            f"Gracefully quit {len(ports_to_kill)} server(s)?\n\n"
            "For execution servers: kills workers only, server stays alive.",
            QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No,
            QMessageBox.StandardButton.Yes,
        )

        if reply != QMessageBox.StandardButton.Yes:
            return

        # Kill in background thread to avoid blocking UI
        def kill_servers():
            from openhcs.runtime.zmq_config import OPENHCS_ZMQ_CONFIG
            from zmqruntime.client import ZMQClient
            from zmqruntime.queue_tracker import GlobalQueueTrackerRegistry

            failed_ports = []
            registry = GlobalQueueTrackerRegistry()

            for port in ports_to_kill:
                try:
                    logger.info(f"Attempting to quit server on port {port}...")
                    success = ZMQClient.kill_server_on_port(
                        port, graceful=True, config=OPENHCS_ZMQ_CONFIG
                    )
                    if success:
                        logger.info(f"✅ Successfully quit server on port {port}")
                        # Clear queue tracker for this viewer
                        registry.remove_tracker(port)
                        self.server_killed.emit(port)
                    else:
                        failed_ports.append(port)
                        logger.warning(
                            f"❌ Failed to quit server on port {port} (kill_server_on_port returned False)"
                        )
                except Exception as e:
                    failed_ports.append(port)
                    logger.error(f"❌ Error quitting server on port {port}: {e}")

            # Emit completion signal
            if failed_ports:
                self._kill_complete.emit(
                    False, f"Failed to quit servers on ports: {failed_ports}"
                )
            else:
                self._kill_complete.emit(True, "All servers quit successfully")

        spawn_thread_with_context(kill_servers, name="kill_servers")

    def force_kill_selected_servers(self):
        """Force kill selected servers (async to avoid blocking UI)."""
        selected_items = self.server_tree.selectedItems()
        if not selected_items:
            QMessageBox.warning(
                self, "No Selection", "Please select servers to force kill."
            )
            return

        # Collect ports to kill BEFORE showing dialog (items may be deleted by auto-refresh)
        ports_to_kill = []
        for item in selected_items:
            data = item.data(0, Qt.ItemDataRole.UserRole)
            # Skip worker items (they don't have ports)
            if data and data.get("type") == "worker":
                continue
            port = data.get("port") if data else None
            if port:
                ports_to_kill.append(port)

        if not ports_to_kill:
            QMessageBox.warning(
                self, "No Servers", "No servers selected (only workers selected)."
            )
            return

        # Confirm with user
        reply = QMessageBox.question(
            self,
            "Force Kill Confirmation",
            f"Force kill {len(ports_to_kill)} server(s)?\n\n"
            "For execution servers: kills workers AND server.\n"
            "For Napari viewers: kills the viewer process.",
            QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No,
            QMessageBox.StandardButton.No,
        )

        if reply != QMessageBox.StandardButton.Yes:
            return

        def kill_servers():
            """Kill selected servers gracefully or forcefully."""
            from openhcs.runtime.zmq_config import OPENHCS_ZMQ_CONFIG
            from zmqruntime.client import ZMQClient
            from zmqruntime.queue_tracker import GlobalQueueTrackerRegistry

            failed_ports = []
            registry = GlobalQueueTrackerRegistry()

            for port in ports_to_kill:
                try:
                    logger.info(
                        f"🔥 FORCE KILL: Force killing server on port {port} (kills workers AND server)"
                    )
                    # Use kill_server_on_port with graceful=False
                    # This handles both IPC and TCP modes correctly
                    success = ZMQClient.kill_server_on_port(
                        port, graceful=False, config=OPENHCS_ZMQ_CONFIG
                    )

                    if success:
                        logger.info(
                            f"✅ Successfully force killed server on port {port}"
                        )
                    else:
                        logger.warning(
                            f"⚠️ Force kill returned False for port {port}, but continuing cleanup"
                        )

                    # Clear queue tracker for this viewer (always, regardless of kill result)
                    registry.remove_tracker(port)
                    self.server_killed.emit(port)

                except Exception as e:
                    logger.error(f"❌ Error force killing server on port {port}: {e}")
                    # Still emit signal to trigger refresh and cleanup, even on error
                    registry.remove_tracker(port)
                    self.server_killed.emit(port)

            # Always emit success - we've done our best to kill the processes
            # The refresh will remove any dead entries from the list
            self._kill_complete.emit(
                True, "Force kill operation completed (list will refresh)"
            )

        spawn_thread_with_context(kill_servers, name="force_kill_servers")

    def _on_item_double_clicked(self, item: QTreeWidgetItem):
        """Handle double-click on tree item - open log file."""
        data = item.data(0, Qt.ItemDataRole.UserRole)

        # For worker items, get the server from parent
        if data and data.get("type") == "worker":
            data = data.get("server", {})

        log_file = data.get("log_file_path") if data else None

        if log_file and Path(log_file).exists():
            # Emit signal for parent to handle (e.g., open in log viewer)
            self.log_file_opened.emit(log_file)
            logger.info(f"Opened log file: {log_file}")
        else:
            QMessageBox.information(
                self,
                "No Log File",
                f"No log file available for this item.\n\nPort: {data.get('port', 'unknown') if data else 'unknown'}",
            )
