StreamLens/analyzer/gui/main_window.py

"""
Main GUI window for StreamLens
"""

import sys
import os
from typing import Optional, List, TYPE_CHECKING

try:
    from PySide6.QtWidgets import (
        QMainWindow, QWidget, QVBoxLayout, QHBoxLayout, QSplitter,
        QTableWidget, QTableWidgetItem, QTextEdit, QMenuBar, QMenu,
        QFileDialog, QMessageBox, QProgressBar, QStatusBar, QLabel,
        QHeaderView, QPushButton, QGroupBox
    )
    from PySide6.QtCore import Qt, QThread, Signal, QTimer
    from PySide6.QtGui import QAction, QIcon, QFont
    
    # Matplotlib integration - lazy loaded
    matplotlib = None
    FigureCanvas = None
    NavigationToolbar = None
    Figure = None
    plt = None
    
    def _ensure_matplotlib_gui_loaded():
        """Lazy load matplotlib for GUI mode"""
        global matplotlib, FigureCanvas, NavigationToolbar, Figure, plt
        
        if matplotlib is not None:
            return True
            
        try:
            import matplotlib as mpl
            matplotlib = mpl
            matplotlib.use('Qt5Agg')  # Use Qt backend for matplotlib
            
            from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FC
            from matplotlib.backends.backend_qt5agg import NavigationToolbar2QT as NT
            from matplotlib.figure import Figure as Fig
            import matplotlib.pyplot as pyplot
            
            FigureCanvas = FC
            NavigationToolbar = NT
            Figure = Fig
            plt = pyplot
            
            return True
        except ImportError as e:
            print(f"Matplotlib GUI integration not available: {e}")
            return False
    
except ImportError as e:
    print(f"GUI dependencies not available: {e}")
    sys.exit(1)

if TYPE_CHECKING:
    from ..analysis.core import EthernetAnalyzer
    from ..models.flow_stats import FlowStats


class PlotWidget(QWidget):
    """Widget containing matplotlib plot with toolbar"""
    
    def __init__(self, parent=None):
        super().__init__(parent)
        self.figure = None
        self.canvas = None
        self.toolbar = None
        self._initialized = False
        
        # Just create basic layout initially
        self.layout = QVBoxLayout()
        self.setLayout(self.layout)
    
    def _ensure_initialized(self):
        """Initialize matplotlib components when first needed"""
        if self._initialized:
            return True
            
        if not _ensure_matplotlib_gui_loaded():
            return False
            
        self.figure = Figure(figsize=(12, 8))
        self.canvas = FigureCanvas(self.figure)
        self.toolbar = NavigationToolbar(self.canvas, self)
        
        self.layout.addWidget(self.toolbar)
        self.layout.addWidget(self.canvas)
        
        self._initialized = True
        self.clear_plot()
        return True
    
    def clear_plot(self):
        """Clear the plot and show initial message"""
        if not self._ensure_initialized():
            return
        self.figure.clear()
        
        # Add a centered message when no plot is shown
        ax = self.figure.add_subplot(111)
        ax.text(0.5, 0.5, 'Select a flow with Chapter 10 data to view signal visualization', 
               ha='center', va='center', transform=ax.transAxes, fontsize=12, 
               color='gray', style='italic')
        ax.set_xlim(0, 1)
        ax.set_ylim(0, 1)
        ax.axis('off')
        
        self.canvas.draw()
    
    def plot_flow_signals(self, flow: 'FlowStats', signal_data_list: List, flow_key: str):
        """Plot Chapter 10 signals for a flow"""
        if not self._ensure_initialized():
            return
        self.figure.clear()
        
        if not signal_data_list:
            # Show message when no data
            ax = self.figure.add_subplot(111)
            ax.text(0.5, 0.5, 'No Chapter 10 signal data found in selected flow', 
                   ha='center', va='center', transform=ax.transAxes, fontsize=14)
            ax.set_xlim(0, 1)
            ax.set_ylim(0, 1)
            ax.axis('off')
            self.canvas.draw()
            return
        
        # Create subplots for each signal
        n_signals = len(signal_data_list)
        axes = []
        
        for idx, signal_data in enumerate(signal_data_list):
            ax = self.figure.add_subplot(n_signals, 1, idx + 1)
            axes.append(ax)
            
            # Plot each channel in the signal data
            for channel_name, data in signal_data.channels.items():
                ax.plot(signal_data.timestamps, data, label=channel_name, linewidth=0.8)
            
            ax.set_xlabel('Time (s)')
            ax.set_ylabel('Amplitude')
            ax.grid(True, alpha=0.3)
            ax.legend()
            
            # Add metadata info
            if signal_data.metadata and signal_data.metadata.channel_configs:
                config_info = []
                for ch_id, config in signal_data.metadata.channel_configs.items():
                    if 'units' in config:
                        config_info.append(f"CH{ch_id}: {config.get('units', 'Unknown')}")
                if config_info:
                    ax.set_title(f"Channels: {', '.join(config_info)}")
        
        self.figure.suptitle(f'Chapter 10 Signals - Flow: {flow_key}', fontsize=14)
        self.figure.tight_layout()
        self.canvas.draw()


class PCAPLoadThread(QThread):
    """Thread for loading PCAP files without blocking UI"""
    
    progress_updated = Signal(int)
    loading_finished = Signal(object)  # Emits the analyzer
    error_occurred = Signal(str)
    
    def __init__(self, file_path: str, parent=None):
        super().__init__(parent)
        self.file_path = file_path
    
    def run(self):
        try:
            from ..analysis.core import EthernetAnalyzer
            from ..utils.pcap_loader import PCAPLoader
            
            # Create analyzer
            analyzer = EthernetAnalyzer()
            
            # Load PCAP
            loader = PCAPLoader(self.file_path)
            if not loader.validate_file():
                self.error_occurred.emit(f"Invalid PCAP file: {self.file_path}")
                return
            
            packets = loader.load_all()
            analyzer.all_packets = packets
            
            # Process packets with progress updates
            total_packets = len(packets)
            for i, packet in enumerate(packets, 1):
                analyzer._process_single_packet(packet, i)
                
                # Update progress every 1000 packets
                if i % 1000 == 0 or i == total_packets:
                    progress = int((i / total_packets) * 100)
                    self.progress_updated.emit(progress)
            
            # Calculate statistics
            analyzer.calculate_statistics()
            
            self.loading_finished.emit(analyzer)
            
        except Exception as e:
            self.error_occurred.emit(str(e))


class StreamLensMainWindow(QMainWindow):
    """Main application window"""
    
    def __init__(self):
        super().__init__()
        self.analyzer: Optional['EthernetAnalyzer'] = None
        self.current_file = None
        self.loading_thread = None
        
        self.setWindowTitle("StreamLens - Ethernet Traffic Analyzer")
        self.setGeometry(100, 100, 1400, 900)
        
        self.setup_ui()
        self.setup_menus()
        self.setup_status_bar()
        
    def setup_ui(self):
        """Set up the main UI layout"""
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        
        # Create main horizontal splitter
        main_splitter = QSplitter(Qt.Horizontal)
        
        # Left panel - Flow list and controls
        left_panel = self.create_left_panel()
        left_panel.setMaximumWidth(600)
        left_panel.setMinimumWidth(400)
        
        # Right panel - Plot area
        right_panel = self.create_right_panel()
        
        main_splitter.addWidget(left_panel)
        main_splitter.addWidget(right_panel)
        main_splitter.setStretchFactor(0, 0)  # Don't stretch left panel
        main_splitter.setStretchFactor(1, 1)  # Stretch right panel
        
        layout = QHBoxLayout()
        layout.addWidget(main_splitter)
        central_widget.setLayout(layout)
    
    def create_left_panel(self):
        """Create the left panel with flow list and info"""
        widget = QWidget()
        layout = QVBoxLayout()
        
        # File info group
        file_group = QGroupBox("File Information")
        file_layout = QVBoxLayout()
        self.file_info_label = QLabel("No file loaded")
        self.file_info_label.setWordWrap(True)
        file_layout.addWidget(self.file_info_label)
        file_group.setLayout(file_layout)
        layout.addWidget(file_group)
        
        # Flow list group
        flows_group = QGroupBox("IP Flows (sorted by max sigma deviation)")
        flows_layout = QVBoxLayout()
        
        self.flows_table = QTableWidget()
        self.flows_table.setColumnCount(5)
        self.flows_table.setHorizontalHeaderLabels([
            "Source → Destination", "Packets", "Max σ", "Protocols", "Frame Types"
        ])
        
        # Configure table
        header = self.flows_table.horizontalHeader()
        header.setStretchLastSection(True)
        header.setSectionResizeMode(0, QHeaderView.Stretch)
        
        self.flows_table.setSelectionBehavior(QTableWidget.SelectRows)
        self.flows_table.setAlternatingRowColors(True)
        self.flows_table.itemSelectionChanged.connect(self.on_flow_selected)
        
        flows_layout.addWidget(self.flows_table)
        
        # Status label for visualization feedback
        self.viz_status_label = QLabel("Select a flow to view signal visualization")
        self.viz_status_label.setStyleSheet("color: gray; font-style: italic;")
        flows_layout.addWidget(self.viz_status_label)
        
        flows_group.setLayout(flows_layout)
        layout.addWidget(flows_group)
        
        widget.setLayout(layout)
        return widget
    
    def create_right_panel(self):
        """Create the right panel with plot area"""
        widget = QWidget()
        layout = QVBoxLayout()
        
        plot_group = QGroupBox("Signal Visualization")
        plot_layout = QVBoxLayout()
        
        self.plot_widget = PlotWidget()
        plot_layout.addWidget(self.plot_widget)
        
        plot_group.setLayout(plot_layout)
        layout.addWidget(plot_group)
        
        widget.setLayout(layout)
        return widget
    
    def setup_menus(self):
        """Set up application menus"""
        menubar = self.menuBar()
        
        # File menu
        file_menu = menubar.addMenu("File")
        
        open_action = QAction("Open PCAP...", self)
        open_action.setShortcut("Ctrl+O")
        open_action.triggered.connect(self.open_pcap_file)
        file_menu.addAction(open_action)
        
        file_menu.addSeparator()
        
        exit_action = QAction("Exit", self)
        exit_action.setShortcut("Ctrl+Q")
        exit_action.triggered.connect(self.close)
        file_menu.addAction(exit_action)
        
        # View menu
        view_menu = menubar.addMenu("View")
        
        refresh_action = QAction("Refresh", self)
        refresh_action.setShortcut("F5")
        refresh_action.triggered.connect(self.refresh_data)
        view_menu.addAction(refresh_action)
    
    def setup_status_bar(self):
        """Set up status bar"""
        self.status_bar = QStatusBar()
        self.setStatusBar(self.status_bar)
        
        self.progress_bar = QProgressBar()
        self.progress_bar.setVisible(False)
        self.status_bar.addPermanentWidget(self.progress_bar)
        
        self.status_bar.showMessage("Ready")
    
    def open_pcap_file(self):
        """Open PCAP file dialog"""
        file_path, _ = QFileDialog.getOpenFileName(
            self, 
            "Open PCAP File", 
            "", 
            "PCAP Files (*.pcap *.pcapng);;All Files (*)"
        )
        
        if file_path:
            self.load_pcap_file(file_path)
    
    def load_pcap_file(self, file_path: str):
        """Load PCAP file in background thread"""
        if self.loading_thread and self.loading_thread.isRunning():
            return
        
        self.current_file = file_path
        self.status_bar.showMessage(f"Loading {os.path.basename(file_path)}...")
        self.progress_bar.setVisible(True)
        self.progress_bar.setValue(0)
        
        # Disable UI during loading
        self.flows_table.setEnabled(False)
        
        # Start loading thread
        self.loading_thread = PCAPLoadThread(file_path)
        self.loading_thread.progress_updated.connect(self.progress_bar.setValue)
        self.loading_thread.loading_finished.connect(self.on_pcap_loaded)
        self.loading_thread.error_occurred.connect(self.on_loading_error)
        self.loading_thread.start()
    
    def on_pcap_loaded(self, analyzer: 'EthernetAnalyzer'):
        """Handle PCAP loading completion"""
        self.analyzer = analyzer
        self.progress_bar.setVisible(False)
        self.flows_table.setEnabled(True)
        
        # Update file info
        summary = analyzer.get_summary()
        file_info = f"File: {os.path.basename(self.current_file)}\n"
        file_info += f"Packets: {summary['total_packets']:,}\n"
        file_info += f"Flows: {summary['unique_flows']}\n"
        file_info += f"IPs: {summary['unique_ips']}"
        self.file_info_label.setText(file_info)
        
        # Populate flows table
        self.populate_flows_table()
        
        self.status_bar.showMessage(f"Loaded {summary['total_packets']:,} packets from {os.path.basename(self.current_file)}")
    
    def on_loading_error(self, error_message: str):
        """Handle loading error"""
        self.progress_bar.setVisible(False)
        self.flows_table.setEnabled(True)
        
        QMessageBox.critical(self, "Loading Error", f"Error loading PCAP file:\n{error_message}")
        self.status_bar.showMessage("Error loading file")
    
    def populate_flows_table(self):
        """Populate the flows table with data"""
        if not self.analyzer:
            return
        
        summary = self.analyzer.get_summary()
        flows_list = list(summary['flows'].values())
        
        # Sort by maximum sigma deviation
        flows_list.sort(key=lambda x: (
            self.analyzer.statistics_engine.get_max_sigma_deviation(x),
            x.frame_count
        ), reverse=True)
        
        self.flows_table.setRowCount(len(flows_list))
        
        for row, flow in enumerate(flows_list):
            # Source -> Destination
            flow_item = QTableWidgetItem(f"{flow.src_ip} → {flow.dst_ip}")
            self.flows_table.setItem(row, 0, flow_item)
            
            # Packets
            packets_item = QTableWidgetItem(str(flow.frame_count))
            packets_item.setData(Qt.UserRole, flow)  # Store flow object
            self.flows_table.setItem(row, 1, packets_item)
            
            # Max sigma deviation
            max_sigma = self.analyzer.statistics_engine.get_max_sigma_deviation(flow)
            sigma_item = QTableWidgetItem(f"{max_sigma:.2f}σ")
            self.flows_table.setItem(row, 2, sigma_item)
            
            # Protocols
            protocols = ", ".join(flow.protocols)
            if flow.detected_protocol_types:
                protocols += f" ({', '.join(flow.detected_protocol_types)})"
            protocols_item = QTableWidgetItem(protocols)
            self.flows_table.setItem(row, 3, protocols_item)
            
            # Frame types
            frame_types = ", ".join(flow.frame_types.keys())
            frame_types_item = QTableWidgetItem(frame_types)
            self.flows_table.setItem(row, 4, frame_types_item)
        
        # Resize columns to content
        self.flows_table.resizeColumnsToContents()
    
    def on_flow_selected(self):
        """Handle flow selection and automatically render plots"""
        selected_rows = self.flows_table.selectionModel().selectedRows()
        
        if not selected_rows:
            # No selection - clear plot and show message
            self.plot_widget.clear_plot()
            self.viz_status_label.setText("Select a flow to view signal visualization")
            self.viz_status_label.setStyleSheet("color: gray; font-style: italic;")
            return
        
        # Get selected flow and automatically visualize
        self.visualize_selected_flow()
    
    def visualize_selected_flow(self):
        """Visualize the selected flow"""
        if not self.analyzer:
            return
        
        selected_rows = self.flows_table.selectionModel().selectedRows()
        if not selected_rows:
            return
        
        # Get selected flow
        row = selected_rows[0].row()
        flow_item = self.flows_table.item(row, 1)  # Packets column has flow data
        flow = flow_item.data(Qt.UserRole)
        
        if not flow:
            return
        
        # Check if flow has Chapter 10 data
        has_ch10 = any('CH10' in ft or 'TMATS' in ft for ft in flow.frame_types.keys())
        
        if not has_ch10:
            # Clear plot and show informative message
            self.plot_widget.clear_plot()
            self.viz_status_label.setText("Selected flow does not contain Chapter 10 telemetry data")
            self.viz_status_label.setStyleSheet("color: orange; font-style: italic;")
            return
        
        # Get flow packets
        flow_packets = self.get_flow_packets(flow)
        
        if not flow_packets:
            # Clear plot and show error message
            self.plot_widget.clear_plot()
            self.viz_status_label.setText("No packets found for selected flow")
            self.viz_status_label.setStyleSheet("color: red; font-style: italic;")
            return
        
        # Use signal visualizer to extract and process signals
        from ..utils.signal_visualizer import signal_visualizer
        
        flow_key = f"{flow.src_ip} → {flow.dst_ip}"
        self.status_bar.showMessage(f"Extracting signals from flow {flow_key}...")
        self.viz_status_label.setText("Processing Chapter 10 telemetry data...")
        self.viz_status_label.setStyleSheet("color: blue; font-style: italic;")
        
        try:
            # Extract signals
            tmats_metadata = signal_visualizer._extract_tmats_from_flow(flow_packets)
            signal_data_list = signal_visualizer._extract_signals_from_flow(flow_packets, tmats_metadata)
            
            # Plot in GUI
            self.plot_widget.plot_flow_signals(flow, signal_data_list, flow_key)
            
            if signal_data_list:
                self.status_bar.showMessage(f"Plotted {len(signal_data_list)} signal(s) from {flow_key}")
                self.viz_status_label.setText(f"Displaying {len(signal_data_list)} signal channel(s)")
                self.viz_status_label.setStyleSheet("color: green; font-weight: bold;")
            else:
                self.status_bar.showMessage(f"No decodable signal data found in {flow_key}")
                self.viz_status_label.setText("No decodable signal data found in selected flow")
                self.viz_status_label.setStyleSheet("color: orange; font-style: italic;")
                
        except Exception as e:
            # Show error in status label instead of popup
            self.plot_widget.clear_plot()
            self.viz_status_label.setText(f"Visualization error: {str(e)}")
            self.viz_status_label.setStyleSheet("color: red; font-style: italic;")
            self.status_bar.showMessage("Visualization error")
    
    def get_flow_packets(self, flow: 'FlowStats') -> List:
        """Get all packets for a specific flow"""
        if not self.analyzer or not self.analyzer.all_packets:
            return []
        
        flow_packets = []
        
        for packet in self.analyzer.all_packets:
            try:
                if hasattr(packet, 'haslayer'):
                    from scapy.all import IP
                    if packet.haslayer(IP):
                        ip_layer = packet[IP]
                        if ip_layer.src == flow.src_ip and ip_layer.dst == flow.dst_ip:
                            flow_packets.append(packet)
            except:
                continue
        
        return flow_packets
    
    def refresh_data(self):
        """Refresh the current data"""
        if self.current_file:
            self.load_pcap_file(self.current_file)
    
    def closeEvent(self, event):
        """Handle application close"""
        if self.loading_thread and self.loading_thread.isRunning():
            self.loading_thread.quit()
            self.loading_thread.wait()
        
        event.accept()