Update app.py

app.py

@@ -1,955 +1,459 @@
-import warnings
-warnings.filterwarnings("ignore", category=FutureWarning, message=".*torch.distributed.reduce_op.*")
 import gradio as gr
-import torch
-import time
+import os
 import sys
 import subprocess
-import os
-import pty
-import select
-import signal
-import fcntl
-import struct
-import termios
 import threading
-import re
+import time
 import ast
-import json
-import types
+import re
+import signal
+import shutil
 import tempfile
-from pathlib import Path
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, GenerationConfig
-from typing import List, Dict, Optional
-from functools import lru_cache
-
-os.environ["TOKENIZERS_PARALLELISM"] = "false"
-
-MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
-DEVICE = "cpu"
-DTYPE = torch.float32
-DISPLAY_NUM = ":99"
-SCREEN_W, SCREEN_H = 800, 600
-SNIPPETS_FILE = "/tmp/axon_snippets.json"
+import atexit
+from PIL import Image
 
 # ═══════════════════════════════════════
-# Virtual Display
+# 1. HELPER CLASSES
 # ═══════════════════════════════════════
-class VirtualDisplay:
+
+class VirtualDisplayManager:
+    """Manages a virtual display (Xvfb) for capturing GUI output."""
     def __init__(self):
-        self.xvfb_proc = None
-        self.display = DISPLAY_NUM
-        self._ensure_deps()
-        self._start_xvfb()
-
-    def _ensure_deps(self):
-        needed = []
-        for cmd in ["Xvfb", "xdpyinfo"]:
-            if subprocess.run(["which", cmd], capture_output=True).returncode != 0:
-                needed.extend(["xvfb", "x11-utils"]); break
-        for cmd in ["scrot", "import"]:
-            if subprocess.run(["which", cmd], capture_output=True).returncode == 0: break
-        else:
-            needed.extend(["scrot", "imagemagick"])
-        if needed:
-            print(f"[Xvfb] Installing: {set(needed)}")
-            subprocess.run(["apt-get", "update", "-qq"], capture_output=True, timeout=60)
-            subprocess.run(["apt-get", "install", "-y", "-qq"] + list(set(needed)),
-                           capture_output=True, timeout=120)
-
-    def _start_xvfb(self):
+        self.display = None
+        self.is_active = False
+        self.check_xvfb()
+
+    def check_xvfb(self):
+        # Check if Xvfb is installed (Linux only)
+        self.has_xvfb = shutil.which('Xvfb') is not None
+
+    def start(self):
+        if not self.has_xvfb:
+            print("Xvfb not found. Virtual display disabled.")
+            return
         try:
-            subprocess.run(["pkill", "-f", f"Xvfb {self.display}"], capture_output=True, timeout=5)
-            time.sleep(0.2)
-            self.xvfb_proc = subprocess.Popen(
-                ["Xvfb", self.display, "-screen", "0", f"{SCREEN_W}x{SCREEN_H}x24",
-                 "-ac", "-nolisten", "tcp"],
-                stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
-            time.sleep(0.5)
-            if self.xvfb_proc.poll() is None:
-                print(f"[Xvfb] Running on {self.display}")
-                os.environ["DISPLAY"] = self.display
-            else:
-                self.xvfb_proc = None
+            from pyvirtualdisplay import Display
+            self.display = Display(visible=0, size=(1280, 720))
+            self.display.start()
+            os.environ["DISPLAY"] = f":{self.display.display}"
+            self.is_active = True
+            print(f"Virtual Display started on {os.environ['DISPLAY']}")
         except Exception as e:
-            print(f"[Xvfb] {e}"); self.xvfb_proc = None
+            print(f"Failed to start virtual display: {e}")
 
     def capture(self):
-        if not self.is_running: return None
-        tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False); tmp.close()
-        for cmd in [
-            ["scrot", "--display", self.display, tmp.name],
-            f"xwd -root -display {self.display} | convert xwd:- png:{tmp.name}",
-            ["import", "-window", "root", "-display", self.display, tmp.name],
-        ]:
-            try:
-                kw = {"shell": True} if isinstance(cmd, str) else {}
-                r = subprocess.run(cmd, capture_output=True, timeout=5,
-                                   env={**os.environ, "DISPLAY": self.display}, **kw)
-                if r.returncode == 0 and os.path.exists(tmp.name) and os.path.getsize(tmp.name) > 200:
-                    return tmp.name
-            except: pass
-        try: os.unlink(tmp.name)
-        except: pass
-        return None
+        if not self.is_active:
+            return None
+        try:
+            # Capture screenshot of the virtual display
+            # Using xwd and converting to png is standard, but let's try PIL grab if possible
+            # or a simple screenshot tool if installed.
+            # Fallback: ImageGrab works on Linux if DISPLAY is set correctly and dependencies exist.
+            from PIL import ImageGrab
+            img = ImageGrab.grab()
+            
+            # Save to a temporary path for Gradio to display
+            fd, path = tempfile.mkstemp(suffix=".png")
+            os.close(fd)
+            img.save(path)
+            return path
+        except Exception as e:
+            print(f"Capture failed: {e}")
+            return None
 
-    @property
-    def is_running(self): return self.xvfb_proc is not None and self.xvfb_proc.poll() is None
     def cleanup(self):
-        if self.xvfb_proc:
-            self.xvfb_proc.terminate()
-            try: self.xvfb_proc.wait(timeout=3)
-            except: self.xvfb_proc.kill()
-
-vdisplay = VirtualDisplay()
+        if self.display:
+            self.display.stop()
 
-# ═══════════════════════════════════════
-# GUI Process Manager
-# ═══════════════════════════════════════
-class GUIProcessManager:
-    def __init__(self):
-        self.process = None
-    def launch(self, code):
-        self.stop()
-        tmp = "/tmp/_axon_gui_run.py"
-        with open(tmp, "w") as f: f.write(code)
-        env = os.environ.copy()
-        env.update({"DISPLAY": DISPLAY_NUM, "SDL_VIDEODRIVER": "x11"})
-        try:
-            self.process = subprocess.Popen([sys.executable, tmp], stdout=subprocess.PIPE,
-                stderr=subprocess.PIPE, env=env, preexec_fn=os.setsid)
-            time.sleep(1.0)
-            if self.process.poll() is not None:
-                out, err = self.process.communicate(timeout=2)
-                msg = (out.decode(errors="replace") + "\n" + err.decode(errors="replace")).strip()
-                return f"[Exited immediately]\n{msg}" if msg else "[Exited]"
-            return f"[GUI launched — PID {self.process.pid}]  CAPTURE to see output."
-        except Exception as e: return f"[Error] {e}"
-    def stop(self):
-        if self.process and self.process.poll() is None:
-            pid = self.process.pid
-            try: os.killpg(os.getpgid(pid), signal.SIGTERM); self.process.wait(timeout=3)
-            except:
-                try: os.killpg(os.getpgid(pid), signal.SIGKILL)
-                except: pass
-            self.process = None; return f"[Stopped PID {pid}]"
-        self.process = None; return "[No process]"
-    @property
-    def is_running(self): return self.process is not None and self.process.poll() is None
-    def get_status(self): return f"● PID {self.process.pid}" if self.is_running else "○ idle"
-
-gui_mgr = GUIProcessManager()
-
-# ═══════════════════════════════════════
-# File System (with cross-file imports)
-# ═══════════════════════════════════════
 class FileSystem:
-    def __init__(self):
-        self.files: Dict[str, str] = {
-            "main.py": '# Start coding here\nfrom utils import add\n\nresult = add(3, 7)\nprint(f"3 + 7 = {result}")',
-            "utils.py": "# Utility functions\ndef add(a, b):\n    return a + b\n\ndef multiply(a, b):\n    return a * b",
-            "notes.txt": "Project requirements:\n- Build a cool AI app\n- Use Gradio"
-        }
-        self.current_file = "main.py"
-        self._sync_dir = "/tmp/axon_workspace"
-        self._sync_to_disk()
-
-    def _sync_to_disk(self):
-        """Write all virtual files to disk so cross-file imports work."""
-        os.makedirs(self._sync_dir, exist_ok=True)
-        for name, content in self.files.items():
-            path = os.path.join(self._sync_dir, name)
-            with open(path, "w") as f:
-                f.write(content)
+    """Manages file operations."""
+    def __init__(self, root="."):
+        self.root = os.path.abspath(root)
+        self.current_file = None
+        self.files = []
+        self.refresh()
+
+    def refresh(self):
+        # List only .py files for simplicity, or all text files
+        self.files = [f for f in os.listdir(self.root) if f.endswith('.py') or f.endswith('.txt') or f.endswith('.md')]
+        self.files.sort()
+        if not self.current_file and self.files:
+            self.current_file = self.files[0]
+        return self.files
+
+    def get_all_files(self):
+        self.refresh()
+        return self.files
+
+    def set_current_file(self, filename):
+        if filename in self.files:
+            self.current_file = filename
+
+    def get_current_file_content(self):
+        if not self.current_file: return ""
+        try:
+            with open(os.path.join(self.root, self.current_file), 'r', encoding='utf-8') as f:
+                return f.read()
+        except:
+            return "Error reading file."
 
     def save_file(self, content):
-        if self.current_file:
-            self.files[self.current_file] = content
-            # Sync to disk for imports
-            path = os.path.join(self._sync_dir, self.current_file)
-            os.makedirs(os.path.dirname(path) if "/" in self.current_file else self._sync_dir, exist_ok=True)
-            with open(path, "w") as f:
-                f.write(content)
-
-    def get_current_file_content(self): return self.files.get(self.current_file, "")
-    def set_current_file(self, f):
-        if f in self.files: self.current_file = f
-    def create_file(self, f, c=""):
-        if f not in self.files:
-            self.files[f] = c; self.current_file = f
-            self._sync_to_disk()
-    def delete_file(self, f):
-        if f in self.files and len(self.files) > 1:
-            del self.files[f]
-            path = os.path.join(self._sync_dir, f)
-            if os.path.exists(path): os.unlink(path)
-            self.current_file = list(self.files.keys())[0]
-            return True
-        return False
-    def rename_file(self, old, new):
-        if old in self.files and new not in self.files:
-            self.files[new] = self.files.pop(old)
-            old_path = os.path.join(self._sync_dir, old)
-            if os.path.exists(old_path): os.unlink(old_path)
-            self._sync_to_disk()
-            if self.current_file == old: self.current_file = new
+        if not self.current_file: return
+        with open(os.path.join(self.root, self.current_file), 'w', encoding='utf-8') as f:
+            f.write(content)
+
+    def create_file(self, name):
+        if not name: return
+        path = os.path.join(self.root, name)
+        if not os.path.exists(path):
+            with open(path, 'w') as f:
+                f.write("# New file\n")
+            self.refresh()
+            self.current_file = name
+
+    def delete_file(self, name):
+        path = os.path.join(self.root, name)
+        if os.path.exists(path):
+            os.remove(path)
+            self.current_file = None
+            self.refresh()
             return True
         return False
-    def get_all_files(self): return list(self.files.keys())
 
-fs = FileSystem()
-
-# ═══════════════════════════════════════
-# Snippet Library
-# ═══════════════════════════════════════
-class SnippetLibrary:
+class TerminalManager:
+    """Manages running shell commands."""
     def __init__(self):
-        self.snippets: Dict[str, dict] = {}
-        self._load()
-
-    def _load(self):
+        self.history = []
+        
+    def run_command(self, command):
+        if not command.strip(): return "No command"
+        
         try:
-            if os.path.exists(SNIPPETS_FILE):
-                with open(SNIPPETS_FILE) as f:
-                    self.snippets = json.load(f)
-        except: self.snippets = {}
+            # Capture output
+            result = subprocess.run(
+                command, 
+                shell=True, 
+                cwd=os.getcwd(),
+                stdout=subprocess.PIPE, 
+                stderr=subprocess.PIPE,
+                text=True,
+                timeout=10
+            )
+            output = f"$ {command}\n{result.stdout}\n{result.stderr}"
+        except Exception as e:
+            output = f"$ {command}\nError: {str(e)}"
+            
+        self.history.append(command)
+        return output
+
+    def get_history_choices(self):
+        return list(set(self.history)) # Unique history
+    
+    def cleanup(self):
+        pass
+
+class GUIManager:
+    """Manages running the user's Python script as a subprocess."""
+    def __init__(self):
+        self.process = None
 
-    def _save(self):
+    def launch(self, code):
+        self.stop() # Stop any existing process
+        
+        # Write code to a temp file to run it
+        fd, path = tempfile.mkstemp(suffix=".py")
+        with os.fdopen(fd, 'w') as f:
+            f.write(code)
+            
         try:
-            with open(SNIPPETS_FILE, "w") as f:
-                json.dump(self.snippets, f, indent=2)
-        except: pass
-
-    def add(self, name: str, code: str, tags: str = "") -> str:
-        if not name.strip(): return "Error: name required"
-        self.snippets[name.strip()] = {
-            "code": code,
-            "tags": [t.strip() for t in tags.split(",") if t.strip()],
-            "created": time.strftime("%Y-%m-%d %H:%M"),
-        }
-        self._save()
-        return f"Saved snippet: {name.strip()}"
-
-    def get(self, name: str) -> str:
-        s = self.snippets.get(name)
-        return s["code"] if s else ""
-
-    def delete(self, name: str) -> str:
-        if name in self.snippets:
-            del self.snippets[name]
-            self._save()
-            return f"Deleted: {name}"
-        return "Not found"
-
-    def search(self, query: str) -> List[str]:
-        q = query.lower()
-        results = []
-        for name, data in self.snippets.items():
-            if q in name.lower() or any(q in t.lower() for t in data.get("tags", [])):
-                results.append(name)
-        return results if query else list(self.snippets.keys())
-
-    def list_all(self) -> str:
-        if not self.snippets: return "No snippets saved yet."
-        lines = []
-        for name, data in self.snippets.items():
-            tags = ", ".join(data.get("tags", []))
-            lines.append(f"  {name}  [{tags}]  ({data.get('created', '?')})")
-        return "\n".join(lines)
-
-    def get_names(self) -> List[str]:
-        return list(self.snippets.keys())
+            # Run asynchronously
+            self.process = subprocess.Popen(
+                [sys.executable, path],
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                text=True,
+                cwd=os.getcwd()
+            )
+            return f"Started process PID: {self.process.pid}"
+        except Exception as e:
+            return f"Failed to start: {e}"
 
-snippets = SnippetLibrary()
+    def stop(self):
+        if self.process:
+            self.process.terminate()
+            self.process = None
+            return "Process stopped."
+        return "No process running."
 
-# ═══════════════════════════════════════
-# Code Structure Analyzer
-# ═══════════════════════════════════════
 class CodeAnalyzer:
-    """Parses Python code and extracts structure (functions, classes, imports)."""
-
-    @staticmethod
-    def analyze(code: str) -> str:
-        if not code.strip(): return "Empty file"
+    """Parses Python code to extract structure (Classes/Functions)."""
+    def analyze(self, code):
+        if not code: return "No code."
         try:
             tree = ast.parse(code)
-        except SyntaxError as e:
-            return f"⚠ Syntax error: line {e.lineno}: {e.msg}"
-
-        lines = []
-        imports = []
-        functions = []
-        classes = []
-
-        for node in ast.walk(tree):
-            if isinstance(node, ast.Import):
-                for alias in node.names:
-                    imports.append(alias.name)
-            elif isinstance(node, ast.ImportFrom):
-                imports.append(f"from {node.module} import {', '.join(a.name for a in node.names)}")
-            elif isinstance(node, ast.FunctionDef) or isinstance(node, ast.AsyncFunctionDef):
-                # Check if it's a method (inside a class)
-                args = ", ".join(a.arg for a in node.args.args)
-                prefix = "async " if isinstance(node, ast.AsyncFunctionDef) else ""
-                decorators = [f"@{ast.dump(d)}" for d in node.decorator_list] if hasattr(node, 'decorator_list') else []
-                ret = ""
-                if node.returns:
-                    try: ret = f" → {ast.unparse(node.returns)}"
-                    except: ret = " → ?"
-                functions.append((node.lineno, f"{prefix}def {node.name}({args}){ret}"))
-            elif isinstance(node, ast.ClassDef):
-                bases = ", ".join(ast.unparse(b) for b in node.bases) if node.bases else ""
-                bases_str = f"({bases})" if bases else ""
-                methods = []
-                for item in node.body:
-                    if isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef)):
-                        methods.append(item.name)
-                classes.append((node.lineno, node.name + bases_str, methods))
-
-        # Count stats
-        total_lines = len(code.splitlines())
-        non_empty = len([l for l in code.splitlines() if l.strip()])
-
-        # Build output
-        lines.append(f"📊 {total_lines} lines ({non_empty} non-empty)")
-        lines.append("")
-
-        if imports:
-            lines.append(f"📦 IMPORTS ({len(imports)})")
-            for i in imports:
-                lines.append(f"  {i}")
-            lines.append("")
-
-        if classes:
-            lines.append(f"🏗 CLASSES ({len(classes)})")
-            for lineno, name, methods in classes:
-                lines.append(f"  L{lineno}: class {name}")
-                for m in methods:
-                    lines.append(f"    .{m}()")
-            lines.append("")
-
-        if functions:
-            # Filter out class methods (they're shown under classes)
-            class_methods = set()
-            for _, _, methods in classes:
-                class_methods.update(methods)
-            standalone = [(ln, f) for ln, f in functions if not any(
-                f.split("(")[0].replace("def ", "").replace("async ", "").strip() in class_methods
-                for _ in [None]
-            )]
-            # Actually just show all functions for simplicity
-            lines.append(f"⚡ FUNCTIONS ({len(functions)})")
-            for lineno, sig in functions:
-                lines.append(f"  L{lineno}: {sig}")
-
-        if not imports and not functions and not classes:
-            lines.append("(No functions, classes, or imports found)")
-
-        return "\n".join(lines)
+            lines = []
+            for node in ast.walk(tree):
+                if isinstance(node, ast.ClassDef):
+                    lines.append(f"📦 Class: {node.name}")
+                    for item in node.body:
+                        if isinstance(item, ast.FunctionDef):
+                            lines.append(f"  🔹 def {item.name}(...)")
+                elif isinstance(node, ast.FunctionDef):
+                    # Only list top-level functions here (those not in a class)
+                    # This logic is simplified; ast.walk flattens everything.
+                    # A proper visitor is better, but for a simple view:
+                    pass 
+            
+            # Simple Regex fallback for cleaner hierarchy if AST is too messy
+            summary = []
+            for line in code.splitlines():
+                if line.strip().startswith("class "):
+                    summary.append(f"📦 {line.split(':')[0]}")
+                elif line.strip().startswith("def "):
+                    indent = "  " if line.startswith("    ") else ""
+                    summary.append(f"{indent}🔹 {line.split(':')[0]}")
+            
+            return "\n".join(summary) if summary else "No classes/functions found."
+        except SyntaxError:
+            return "❌ Syntax Error"
+
+class SnippetManager:
+    """Manages code snippets."""
+    def __init__(self):
+        self.snippets = {
+            "Hello World": "print('Hello World')",
+            "Tkinter Basic": "import tkinter as tk\nroot = tk.Tk()\nroot.mainloop()",
+            "Gradio Demo": "import gradio as gr\ndef greet(n):\n return 'Hi '+n\ngr.Interface(fn=greet, inputs='text', outputs='text').launch()"
+        }
 
-analyzer = CodeAnalyzer()
+    def get_names(self):
+        return list(self.snippets.keys())
+
+    def get(self, name):
+        return self.snippets.get(name, "")
+
+    def add(self, name, code, tags):
+        self.snippets[name] = code
+        return "Snippet saved."
 
-# ═══════════════════════════════════════
-# Find & Replace
-# ═══════════════════════════════════════
 class FindReplace:
-    @staticmethod
-    def find_in_file(code: str, query: str, case_sensitive: bool = False) -> str:
-        if not query: return ""
-        lines = code.splitlines()
+    """Handles Text Finding and Replacing."""
+    def find_in_file(self, code, query, match_case):
+        if not query: return "Empty query."
+        flags = 0 if match_case else re.IGNORECASE
+        matches = list(re.finditer(re.escape(query), code, flags))
+        return f"Found {len(matches)} occurrences."
+
+    def find_in_all_files(self, file_list, query):
         results = []
-        for i, line in enumerate(lines, 1):
-            check_line = line if case_sensitive else line.lower()
-            check_query = query if case_sensitive else query.lower()
-            if check_query in check_line:
-                results.append(f"  L{i}: {line.strip()}")
-        if results:
-            return f"Found {len(results)} match(es):\n" + "\n".join(results)
-        return "No matches found."
-
-    @staticmethod
-    def find_in_all_files(files: Dict[str, str], query: str) -> str:
-        if not query: return ""
-        results = []
-        for fname, content in files.items():
-            lines = content.splitlines()
-            for i, line in enumerate(lines, 1):
-                if query.lower() in line.lower():
-                    results.append(f"  {fname}:L{i}: {line.strip()}")
-        if results:
-            return f"Found {len(results)} match(es) across all files:\n" + "\n".join(results)
-        return "No matches found."
-
-    @staticmethod
-    def replace_in_file(code: str, find: str, replace: str, case_sensitive: bool = False) -> tuple:
-        if not find: return code, "Nothing to find."
-        if case_sensitive:
-            count = code.count(find)
-            new_code = code.replace(find, replace)
-        else:
-            pattern = re.compile(re.escape(find), re.IGNORECASE)
-            matches = pattern.findall(code)
-            count = len(matches)
-            new_code = pattern.sub(replace, code)
-        return new_code, f"Replaced {count} occurrence(s)."
+        for filename in file_list:
+            try:
+                with open(filename, 'r') as f:
+                    content = f.read()
+                    if query in content:
+                        results.append(filename)
+            except: pass
+        return "\n".join(results) if results else "No matches in other files."
 
-finder = FindReplace()
+    def replace_in_file(self, code, find, replace, match_case):
+        if not find: return code, "Empty search."
+        flags = 0 if match_case else re.IGNORECASE
+        new_code, count = re.subn(re.escape(find), replace, code, flags=flags)
+        return new_code, f"Replaced {count} occurrences."
 
 # ═══════════════════════════════════════
-# PTY Terminal (with history)
+# 2. INSTANTIATE MANAGERS
 # ═══════════════════════════════════════
-class PTYTerminal:
-    STRIP_ANSI = re.compile(r'\x1b\[[0-9;]*[a-zA-Z]|\x1b\].*?\x07|\x1b\[.*?[@-~]|\r')
-    def __init__(self):
-        self.master_fd = None; self.pid = None
-        self.lock = threading.Lock(); self.log_lines = []; self.max_lines = 800
-        self.cmd_history: List[str] = []
-        self.max_history = 100
-        self._spawn_shell()
-        self._append("AXON TERMINAL v4.0")
-        self._append("══════════════════════════════════════")
-        self._append("Full PTY shell · Cross-file imports")
-        self._append("  pip/npm/apt-get/git/curl/make")
-        self._append("  GUI apps → Display · History → ↑↓")
-        self._append("══════════════════════════════════════")
-
-    def _spawn_shell(self):
-        try:
-            pid, fd = pty.openpty(); self.pid = os.fork()
-            if self.pid == 0:
-                os.close(pid); os.setsid(); fcntl.ioctl(fd, termios.TIOCSCTTY, 0)
-                os.dup2(fd,0); os.dup2(fd,1); os.dup2(fd,2)
-                if fd>2: os.close(fd)
-                env = os.environ.copy()
-                env.update({"TERM":"dumb","PS1":"$ ","DEBIAN_FRONTEND":"noninteractive","DISPLAY":DISPLAY_NUM})
-                os.execvpe("/bin/bash",["/bin/bash","--norc","--noprofile","-i"],env)
-            else:
-                os.close(fd); self.master_fd = pid
-                fl = fcntl.fcntl(self.master_fd, fcntl.F_GETFL)
-                fcntl.fcntl(self.master_fd, fcntl.F_SETFL, fl|os.O_NONBLOCK)
-                fcntl.ioctl(self.master_fd, termios.TIOCSWINSZ, struct.pack("HHHH",40,120,0,0))
-                time.sleep(0.3); self._read(0.2)
-        except Exception as e: self._append(f"[Shell Error] {e}"); self.master_fd=None
-
-    def _read(self, timeout=0.1):
-        if not self.master_fd: return ""
-        out=[]; dl=time.time()+timeout
-        while True:
-            r=dl-time.time()
-            if r<=0: break
-            try:
-                rd,_,_=select.select([self.master_fd],[],[],min(r,0.05))
-                if rd:
-                    c=os.read(self.master_fd,4096)
-                    if c: out.append(c.decode("utf-8",errors="replace")); dl=time.time()+0.15
-                    else: break
-                elif out: break
-            except OSError: break
-        return "".join(out)
-
-    def _clean(self, t):
-        c=self.STRIP_ANSI.sub("",t); return "".join(ch for ch in c if ch in "\n\t" or ord(ch)>=32)
-    def _append(self, t):
-        self.log_lines.append(t)
-        while len(self.log_lines)>self.max_lines: self.log_lines.pop(0)
-    def get_log(self): return "\n".join(self.log_lines)
-
-    def run_command(self, cmd):
-        cmd=cmd.strip()
-        if not cmd: return self.get_log()
-        # Save to history
-        if not self.cmd_history or self.cmd_history[-1] != cmd:
-            self.cmd_history.append(cmd)
-            if len(self.cmd_history) > self.max_history:
-                self.cmd_history.pop(0)
-        with self.lock:
-            if cmd.lower()=="clear": self.log_lines=[]; return ""
-            if cmd.lower()=="history":
-                self._append("Command history:")
-                for i, c in enumerate(self.cmd_history[-20:], 1):
-                    self._append(f"  {i:3d}  {c}")
-                return self.get_log()
-            if not self.master_fd:
-                self._append(f"$ {cmd}"); self._append("[Error] No shell."); return self.get_log()
-            self._read(0.05); self._append(f"$ {cmd}")
-            try: os.write(self.master_fd,(cmd+"\n").encode())
-            except OSError as e: self._append(f"[Write Error] {e}"); return self.get_log()
-            parts=cmd.split(); base=parts[0].lower() if parts else ""
-            lng=["pip","pip3","npm","npx","apt-get","apt","git","wget","curl","make","cmake","cargo","yarn","conda"]
-            wait=180 if base in lng else(60 if base in("python","python3","node")else 15)
-            chunks=[]; start=time.time(); idle=0
-            while time.time()-start<wait:
-                c=self._read(0.3)
-                if c: idle=0; chunks.append(c)
-                else:
-                    idle+=1
-                    if base not in lng and idle>=3: break
-                    if base in lng and idle>=10: break
-                if chunks and "".join(chunks).rstrip().endswith("$"): break
-            raw=self._clean("".join(chunks)); lines=raw.split("\n"); filtered=[]; skip=True
-            for l in lines:
-                s=l.strip()
-                if skip and s==cmd.strip(): skip=False; continue
-                if s in("$","$ "): continue
-                filtered.append(l)
-            result="\n".join(filtered).strip()
-            if result: self._append(result)
-            return self.get_log()
 
-    def get_history_choices(self):
-        return list(reversed(self.cmd_history[-30:]))
+vdisplay = VirtualDisplayManager()
+vdisplay.start()
 
-    def cleanup(self):
-        if self.pid and self.pid>0:
-            try: os.kill(self.pid, signal.SIGTERM)
-            except: pass
-        if self.master_fd:
-            try: os.close(self.master_fd)
-            except: pass
-
-terminal = PTYTerminal()
+fs = FileSystem()
+terminal = TerminalManager()
+gui_mgr = GUIManager()
+analyzer = CodeAnalyzer()
+snippets = SnippetManager()
+finder = FindReplace()
 
 # ═══════════════════════════════════════
-# Model
+# 3. UI THEME & CSS
 # ═══════════════════════════════════════
-@lru_cache(maxsize=1)
-def load_model():
-    print("Loading TinyLlama..."); t0=time.time()
-    try:
-        tok=AutoTokenizer.from_pretrained(MODEL_NAME)
-        mdl=AutoModelForCausalLM.from_pretrained(MODEL_NAME,dtype=DTYPE,device_map=DEVICE,low_cpu_mem_usage=True)
-        gc=GenerationConfig(max_new_tokens=256,temperature=0.2,top_p=0.95,do_sample=True)
-        pipe=pipeline("text-generation",model=mdl,tokenizer=tok,generation_config=gc)
-        print(f"Loaded in {time.time()-t0:.1f}s"); return pipe
-    except Exception as e: print(f"Model error: {e}"); return None
-
-def _llm(s,u,mt=512):
-    m=load_model()
-    if not m: return "Error: model not loaded"
-    try: return m(f"<|system|>{s}</s><|user|>\n{u}\n</s><|assistant|>",max_new_tokens=mt)[0]["generated_text"].strip()
-    except Exception as e: return str(e)
-
-def generate_completion(c): return _llm("Complete this Python code. No explanations.",c,256)
-def explain_code(c): return _llm("Explain this Python code concisely.",c)
-def refactor_code(c): return _llm("Refactor for PEP 8 and best practices.",c)
-def generate_code(p): return _llm(f"Write Python code for: {p}","")
-def create_diff_view(o,m):
-    import difflib; return "\n".join(difflib.unified_diff(o.splitlines(),m.splitlines(),lineterm=""))
-
-GUI_HINTS=["pygame","tkinter","turtle","pyglet","arcade","kivy","wx","Gtk","QApplication",
-           "display.set_mode","mainloop()"]
-def is_gui_code(c): return any(h in c for h in GUI_HINTS)
 
-# ═══════════════════════════════════════
-# Theme + CSS
-# ═══════════════════════════════════════
-theme = gr.themes.Default(
-    font=[gr.themes.GoogleFont("IBM Plex Mono"),"monospace"],
-    primary_hue="cyan", neutral_hue="gray",
-).set(
-    body_background_fill="#080808", body_text_color="#e0e0e0",
-    button_primary_background_fill="#0ff", button_primary_background_fill_hover="#0dd",
-    button_primary_text_color="#000",
-    button_secondary_background_fill="#1a1a1a", button_secondary_background_fill_hover="#2a2a2a",
-    button_secondary_text_color="#0ff", button_secondary_border_color="#0ff4",
-    block_background_fill="#0d0d0d", block_border_color="#1a1a1a",
-    block_label_text_color="#0f8",
-    input_background_fill="#0a0a0a", input_border_color="#1a3a2a", input_placeholder_color="#0f83",
+theme = gr.themes.Soft(
+    primary_hue="orange",
+    secondary_hue="slate",
+    neutral_hue="slate",
 )
 
 css = """
-* { border-radius: 0 !important; }
-.gradio-container { max-width: 100% !important; }
-.axon-header {
-    background: linear-gradient(90deg, #080808 0%, #0a1a1a 50%, #080808 100%);
-    border-bottom: 1px solid #0ff3; padding: 12px 20px !important; margin-bottom: 8px;
-}
-.axon-header h1 {
-    font-family: 'IBM Plex Mono', monospace !important;
-    background: linear-gradient(90deg, #0ff, #0f8);
-    -webkit-background-clip: text; -webkit-text-fill-color: transparent;
-    font-size: 1.4em !important; letter-spacing: 3px; margin: 0 !important;
-}
-.axon-header p { color: #0f84 !important; font-size: 0.75em; letter-spacing: 1px; }
-.terminal-box textarea {
-    background: #050505 !important; color: #00ff41 !important;
-    font-family: 'IBM Plex Mono', monospace !important;
-    font-size: 12px !important; line-height: 1.5 !important;
-    border: 1px solid #0f31 !important; text-shadow: 0 0 4px #0f32;
-}
-.terminal-input input {
-    background: #050505 !important; color: #00ff41 !important;
-    font-family: 'IBM Plex Mono', monospace !important;
-    border: 1px solid #0f32 !important; font-size: 12px !important;
-}
-.terminal-input input::placeholder { color: #0f83 !important; }
-.terminal-input input:focus { border-color: #0ff !important; box-shadow: 0 0 8px #0ff2 !important; }
-.code-editor textarea, .code-editor .cm-editor {
-    font-family: 'IBM Plex Mono', monospace !important; font-size: 13px !important;
-}
-.toolbar-btn {
-    font-size: 0.8em !important; padding: 6px 12px !important; letter-spacing: 0.5px;
-    text-transform: uppercase; font-weight: 600 !important;
-    border: 1px solid transparent !important; transition: all 0.2s ease !important;
-}
-.toolbar-btn:hover { border-color: #0ff !important; box-shadow: 0 0 12px #0ff2 !important; }
-.gr-accordion { border: 1px solid #1a2a2a !important; background: #0a0a0a !important; }
-.gr-accordion > .label-wrap {
-    background: #0d0d0d !important; border-bottom: 1px solid #1a2a2a !important;
-}
-.gr-accordion > .label-wrap:hover { background: #121a1a !important; }
-.gr-accordion > .label-wrap span {
-    color: #0ff !important; font-family: 'IBM Plex Mono', monospace !important;
-    letter-spacing: 1px; font-size: 0.85em;
-}
-.tabs > .tab-nav > button {
-    font-family: 'IBM Plex Mono', monospace !important;
-    letter-spacing: 1px; font-size: 0.8em; text-transform: uppercase;
-    color: #888 !important; border-bottom: 2px solid transparent !important;
-}
-.tabs > .tab-nav > button.selected {
-    color: #0ff !important; border-bottom-color: #0ff !important; text-shadow: 0 0 8px #0ff4;
-}
-.status-bar {
-    background: #0a0a0a !important; border-top: 1px solid #1a2a2a;
-    padding: 4px 12px !important; font-size: 11px !important; color: #0f84 !important;
-}
-.status-bar strong { color: #0ff !important; }
-::-webkit-scrollbar { width: 6px; height: 6px; }
-::-webkit-scrollbar-track { background: #0a0a0a; }
-::-webkit-scrollbar-thumb { background: #0f83; }
-input, textarea, select { color: #0f8 !important; }
-.stop-btn { background: #a00 !important; color: #fff !important; border: none !important; }
-.stop-btn:hover { background: #c00 !important; }
-/* Structure map */
-.structure-box textarea {
-    background: #050505 !important; color: #0ff !important;
-    font-family: 'IBM Plex Mono', monospace !important; font-size: 11px !important;
-}
-/* Snippet code */
-.snippet-code textarea {
-    background: #050505 !important; color: #c8e6c9 !important;
-    font-family: 'IBM Plex Mono', monospace !important; font-size: 12px !important;
-}
-/* Find results */
-.find-results textarea {
-    background: #050505 !important; color: #ffab40 !important;
-    font-family: 'IBM Plex Mono', monospace !important; font-size: 11px !important;
-}
-/* History dropdown */
-.history-dd { font-size: 12px !important; }
+.gradio-container { min-height: 100vh; }
+textarea { font-family: 'Consolas', 'Monaco', monospace !important; }
+.cm-scroller { overflow: auto !important; }
 """
 
 # ═══════════════════════════════════════
-# UI
+# 4. UI CONSTRUCTION
 # ═══════════════════════════════════════
-with gr.Blocks(title="Axon Pro — Python AI IDE") as demo:
-
-    with gr.Row(elem_classes="axon-header"):
-        gr.Markdown("# ⬡ AXON PRO\n\nPYTHON AI IDE — v4.0")
-
-    with gr.Row(equal_height=False):
-
-        # ══ LEFT SIDEBAR ══
-        with gr.Column(scale=1, min_width=200):
-
-            # File Explorer
-            with gr.Accordion("📁 EXPLORER", open=True):
-                file_list = gr.Dropdown(choices=fs.get_all_files(), value=fs.current_file,
-                                         label="", interactive=True, container=False)
-                new_file_name = gr.Textbox(placeholder="filename.py", label="", container=False)
-                with gr.Row():
-                    new_file_btn = gr.Button("+ NEW", variant="secondary", size="sm", elem_classes="toolbar-btn")
-                    save_btn = gr.Button("SAVE", variant="secondary", size="sm", elem_classes="toolbar-btn")
-                    del_file_btn = gr.Button("✗ DEL", variant="secondary", size="sm", elem_classes="toolbar-btn stop-btn")
-
-            # Code Structure
-            with gr.Accordion("🗺 STRUCTURE", open=True):
-                structure_view = gr.Textbox(
-                    value=analyzer.analyze(fs.get_current_file_content()),
-                    lines=10, max_lines=20, interactive=False,
-                    elem_classes="structure-box", label="", show_label=False)
-
-            # Snippets
-            with gr.Accordion("📋 SNIPPETS", open=False):
-                snippet_list = gr.Dropdown(choices=snippets.get_names(), label="Saved Snippets",
-                                            interactive=True, container=False, elem_classes="history-dd")
-                snippet_search = gr.Textbox(placeholder="Search snippets...", label="", container=False)
-                with gr.Row():
-                    snippet_insert_btn = gr.Button("INSERT", variant="primary", size="sm", elem_classes="toolbar-btn")
-                    snippet_del_btn = gr.Button("DEL", variant="secondary", size="sm", elem_classes="toolbar-btn stop-btn")
-                with gr.Accordion("Save New Snippet", open=False):
-                    snippet_name = gr.Textbox(placeholder="snippet name", label="", container=False)
-                    snippet_tags = gr.Textbox(placeholder="tags: api, util, loop", label="", container=False)
-                    snippet_save_btn = gr.Button("SAVE SNIPPET", variant="primary", size="sm", elem_classes="toolbar-btn")
-                    snippet_status = gr.Markdown("")
 
-        # ══ CENTER ══
-        with gr.Column(scale=4):
-            editor = gr.Code(value=fs.get_current_file_content(), label=f"  {fs.current_file}",
-                             language="python", lines=22, interactive=True, elem_classes="code-editor")
+with gr.Blocks(theme=theme, css=css, title="Python IDE") as demo:
+    
+    # State variables
+    # (Gradio 4.x usually handles state via component values, but we can use State objects if needed)
+    
+    with gr.Row():
+        # ─── LEFT SIDEBAR ───
+        with gr.Column(scale=1, min_width=250):
+            gr.Markdown("### 📂 Explorer")
+            
+            with gr.Row():
+                new_file_name = gr.Textbox(placeholder="NewFile.py", show_label=False, scale=3)
+                new_file_btn = gr.Button("➕", scale=1, min_width=10)
+            
+            file_list = gr.Dropdown(choices=fs.get_all_files(), label="Files", value=fs.current_file, interactive=True)
+            
+            with gr.Row():
+                refresh_btn = gr.Button("🔄 Refresh")
+                del_file_btn = gr.Button("🗑 Delete", variant="stop")
 
+            gr.Markdown("### 🌳 Structure")
+            structure_view = gr.Code(language="markdown", label="AST Outline", interactive=False, lines=15)
+
+        # ─── MAIN AREA ───
+        with gr.Column(scale=4):
+            
+            # Toolbar
             with gr.Row():
-                run_btn = gr.Button("▶ RUN", variant="primary", elem_classes="toolbar-btn")
-                run_gui_btn = gr.Button("🖥 GUI", variant="primary", elem_classes="toolbar-btn")
-                stop_gui_btn = gr.Button("■ STOP", variant="secondary", size="sm", elem_classes="toolbar-btn stop-btn")
-                complete_btn = gr.Button("✦ COMPLETE", variant="secondary", elem_classes="toolbar-btn")
-                explain_btn = gr.Button("◈ EXPLAIN", variant="secondary", elem_classes="toolbar-btn")
-                refactor_btn = gr.Button("⟲ REFACTOR", variant="secondary", elem_classes="toolbar-btn")
-
-            with gr.Tabs() as bottom_tabs:
-
-                with gr.Tab("⌘ TERMINAL", id="terminal-tab"):
-                    terminal_output = gr.Textbox(value=terminal.get_log(), lines=10, max_lines=25,
-                        interactive=False, elem_classes="terminal-box", label="", show_label=False)
+                run_btn = gr.Button("▶ RUN (Console)", variant="primary")
+                run_gui_btn = gr.Button("🖼 RUN (GUI)", variant="secondary")
+                stop_gui_btn = gr.Button("⏹ STOP", variant="stop")
+                save_btn = gr.Button("💾 SAVE")
+            
+            # Editor
+            editor = gr.Code(
+                value=fs.get_current_file_content(), 
+                language="python", 
+                label=f"Editor - {fs.current_file}", 
+                lines=20,
+                interactive=True
+            )
+
+            # AI/Tools Row (Placeholder for future expansion)
+            with gr.Row():
+                complete_btn = gr.Button("✨ Auto-Complete")
+                explain_btn = gr.Button("🧐 Explain")
+                refactor_btn = gr.Button("🛠 Refactor")
+
+            # Bottom Tabs
+            with gr.Tabs():
+                
+                # ─ Terminal Tab ─
+                with gr.Tab("💻 TERMINAL"):
+                    terminal_output = gr.Textbox(label="Output", lines=10, max_lines=20, elem_classes="term-out")
                     with gr.Row():
-                        terminal_input = gr.Textbox(
-                            placeholder="$ pip install pygame | npm i | git clone ...",
-                            scale=8, container=False, elem_classes="terminal-input")
-                        terminal_run_btn = gr.Button("⏎", variant="primary", scale=1, min_width=45)
+                        terminal_input = gr.Textbox(show_label=False, placeholder="Type command here...", scale=4)
+                        terminal_run_btn = gr.Button("Run", scale=1)
                     with gr.Row():
-                        clear_btn = gr.Button("CLEAR", variant="secondary", size="sm", elem_classes="toolbar-btn")
-                        history_dd = gr.Dropdown(choices=[], label="", container=False,
-                                                  interactive=True, scale=3, elem_classes="history-dd")
-                        history_refresh = gr.Button("↻ HISTORY", variant="secondary", size="sm",
-                                                     elem_classes="toolbar-btn")
+                        clear_btn = gr.Button("Clear Output")
+                        history_refresh = gr.Button("History")
+                        history_dd = gr.Dropdown(label="Command History", choices=[])
 
-                with gr.Tab("🔍 FIND", id="find-tab"):
+                # ─ Find/Replace Tab ─
+                with gr.Tab("🔍 FIND"):
+                    with gr.Row():
+                        find_query = gr.Textbox(label="Find", placeholder="Text to find...")
+                        replace_query = gr.Textbox(label="Replace", placeholder="Replacement text...")
+                    
+                    with gr.Row():
+                        find_case = gr.Checkbox(label="Match Case")
+                        
                     with gr.Row():
-                        find_query = gr.Textbox(placeholder="Search for...", label="Find", scale=4)
-                        replace_query = gr.Textbox(placeholder="Replace with...", label="Replace", scale=4)
-                        find_case = gr.Checkbox(label="Case sensitive", value=False)
+                        find_btn = gr.Button("FIND", variant="primary")
+                        find_all_btn = gr.Button("FIND ALL")
+                        replace_btn = gr.Button("REPLACE")
+                    
+                    find_results = gr.Textbox(label="Results", lines=5)
+
+                # ─ Snippets Tab ─
+                with gr.Tab("🧩 SNIPPETS"):
                     with gr.Row():
-                        find_btn = gr.Button("FIND", variant="primary", size="sm", elem_classes="toolbar-btn")
-                        find_all_btn = gr.Button("FIND ALL FILES", variant="secondary", size="sm", elem_classes="toolbar-btn")
-                        replace_btn = gr.Button("REPLACE ALL", variant="secondary", size="sm", elem_classes="toolbar-btn")
-                    find_results = gr.Textbox(value="", lines=6, interactive=False,
-                                               elem_classes="find-results", label="", show_label=False)
-
-                with gr.Tab("💬 AI CHAT", id="chat-tab"):
-                    chat_history = gr.Chatbot(label="", height=250)
+                        snippet_list = gr.Dropdown(choices=snippets.get_names(), label="Load Snippet")
+                        snippet_insert_btn = gr.Button("Insert into Editor")
+                    
+                    gr.Markdown("---")
                     with gr.Row():
-                        chat_input = gr.Textbox(placeholder="Describe code you want...",
-                                                 scale=8, container=False)
-                        send_btn = gr.Button("GEN", variant="primary", scale=1, elem_classes="toolbar-btn")
+                        snippet_name = gr.Textbox(label="Snippet Name")
+                        snippet_tags = gr.Textbox(label="Tags (comma sep)")
+                    snippet_save_btn = gr.Button("Save Selection as Snippet")
+                    snippet_status = gr.Markdown("")
 
-                with gr.Tab("⇄ DIFF", id="diff-tab"):
-                    diff_view = gr.Code(label="AI Changes", language="python", interactive=False, lines=15)
+                # ─ Display Tab (Xvfb) ─
+                with gr.Tab("🖥 DISPLAY"):
+                    display_image = gr.Image(label="Xvfb Capture", type="filepath", interactive=False, height=400)
                     with gr.Row():
-                        apply_btn = gr.Button("✓ APPLY", variant="primary", elem_classes="toolbar-btn")
-                        discard_btn = gr.Button("✗ DISCARD", variant="secondary", elem_classes="toolbar-btn")
-
-        # ══ RIGHT: DISPLAY ══
-        with gr.Column(scale=2, min_width=250):
-            with gr.Accordion("🖥 DISPLAY OUTPUT", open=True):
-                display_image = gr.Image(label="", type="filepath", interactive=False, height=420)
-                with gr.Row():
-                    capture_btn = gr.Button("📸 CAPTURE", variant="secondary", size="sm", elem_classes="toolbar-btn")
-                    auto_capture = gr.Checkbox(label="Auto (2s)", value=False)
-                gui_status = gr.Markdown(
-                    f"<small>Xvfb: {'● ON' if vdisplay.is_running else '○ OFF'} "
-                    f"| {SCREEN_W}x{SCREEN_H} | {gui_mgr.get_status()}</small>")
-
-    status_bar = gr.Markdown(
-        f"**AXON PRO v4.0** │ Python {sys.version.split()[0]} │ CPU │ "
-        f"TinyLlama-1.1B │ PTY + Xvfb │ Snippets + Structure + Find",
-        elem_classes="status-bar")
-
-    # State
-    diff_original_state = gr.State("")
-    diff_modified_state = gr.State("")
-    auto_timer = gr.Timer(2, active=False)
+                        capture_btn = gr.Button("📸 CAPTURE SCREEN", variant="primary")
+                        auto_capture = gr.Checkbox(label="Auto-refresh (2s)")
+                    
+                    auto_timer = gr.Timer(value=2.0, active=False)
 
     # ═══════════════════════════════════════
-    # HANDLERS
+    # 5. LOGIC FUNCTIONS
     # ═══════════════════════════════════════
+    
+    # --- File System Wrappers ---
+    def load_file_content(f):
+        fs.set_current_file(f)
+        content = fs.get_current_file_content()
+        return content, f"Editor - {f}", analyzer.analyze(content)
 
-    def update_file(content):
-        fs.save_file(content)
-        return analyzer.analyze(content)
+    def refresh_file_list():
+        files = fs.refresh()
+        return gr.update(choices=files, value=fs.current_file)
 
-    def load_file(filename):
-        fs.set_current_file(filename)
-        content = fs.get_current_file_content()
-        return content, gr.update(label=f"  {filename}"), analyzer.analyze(content)
+    def save_file_content(code):
+        fs.save_file(code)
+        return "File saved."
 
-    def create_file_fn(name):
-        if name and ("." in name or name.endswith("")):
-            fs.create_file(name)
-            content = fs.get_current_file_content()
-            return (gr.update(choices=fs.get_all_files(), value=name),
-                    content, gr.update(label=f"  {name}"), analyzer.analyze(content))
-        return gr.update(choices=fs.get_all_files()), fs.get_current_file_content(), gr.update(), gr.update()
+    def create_new_file(name):
+        if not name: return gr.update(), "Name required", ""
+        fs.create_file(name)
+        content = fs.get_current_file_content()
+        return gr.update(choices=fs.get_all_files(), value=name), f"Editor - {name}", analyzer.analyze(content)
 
-    def delete_file_fn():
-        name = fs.current_file
+    def delete_current_file(name):
         if fs.delete_file(name):
+            next_file = fs.current_file
             content = fs.get_current_file_content()
-            return (gr.update(choices=fs.get_all_files(), value=fs.current_file),
-                    content, gr.update(label=f"  {fs.current_file}"),
-                    analyzer.analyze(content))
-        return gr.update(), fs.get_current_file_content(), gr.update(), gr.update()
-
-    def terminal_cmd(cmd):
-        return terminal.run_command(cmd), ""
-
-    def clear_term(): terminal.log_lines=[]; return ""
-
-    def refresh_history():
-        return gr.update(choices=terminal.get_history_choices())
-
-    def select_history(cmd):
-        if cmd: return cmd
-        return gr.update()
-
-    # Run with cross-file import support
-    def run_editor(content):
-        fs.save_file(content)
-        if is_gui_code(content): return run_gui(content)[0]
-        terminal._append("$ python [editor]")
-        tmp = "/tmp/_axon_run.py"
-        with open(tmp, "w") as f: f.write(content)
+            return (
+                gr.update(choices=fs.get_all_files(), value=next_file), 
+                content, 
+                f"Editor - {next_file}", 
+                analyzer.analyze(content)
+            )
+        return gr.update(), gr.update(), gr.update(), gr.update()
+
+    # --- Execution Wrappers ---
+    def run_console_code(code):
+        # Simply run the code and return stdout
+        # Note: This runs in the same environment, simplistic.
+        fd, path = tempfile.mkstemp(suffix=".py")
+        with os.fdopen(fd, 'w') as f:
+            f.write(code)
         try:
-            env = os.environ.copy()
-            env["DISPLAY"] = DISPLAY_NUM
-            # Add workspace to PYTHONPATH so cross-file imports work
-            env["PYTHONPATH"] = fs._sync_dir + ":" + env.get("PYTHONPATH", "")
-            r = subprocess.run([sys.executable, tmp], capture_output=True, text=True,
-                               timeout=30, env=env, cwd=fs._sync_dir)
-            if r.stdout.strip(): terminal._append(r.stdout.rstrip())
-            if r.stderr.strip(): terminal._append(r.stderr.rstrip())
-            if not r.stdout.strip() and not r.stderr.strip(): terminal._append("(No output)")
-        except subprocess.TimeoutExpired: terminal._append("[Timed out 30s]")
-        except Exception as e: terminal._append(f"[Error] {e}")
-        return terminal.get_log()
-
-    def run_gui(content):
-        fs.save_file(content); msg=gui_mgr.launch(content); terminal._append(msg)
-        time.sleep(1.5); ss=vdisplay.capture()
-        st=f"<small>Xvfb: ● ON | {SCREEN_W}x{SCREEN_H} | {gui_mgr.get_status()}</small>"
-        return terminal.get_log(), ss, st
-
-    def stop_gui():
-        msg=gui_mgr.stop(); terminal._append(msg)
-        st=f"<small>Xvfb: {'● ON' if vdisplay.is_running else '○ OFF'} | {SCREEN_W}x{SCREEN_H} | {gui_mgr.get_status()}</small>"
-        return terminal.get_log(), st
-
-    def capture_display(): return vdisplay.capture()
-    def auto_capture_tick():
-        if gui_mgr.is_running: return vdisplay.capture()
-        return gr.update()
-    def toggle_auto(c): return gr.Timer(2, active=c)
-
-    def complete_fn(c):
-        result = c+"\n"+generate_completion(c)
-        return result, analyzer.analyze(result)
-    def explain_fn(c):
-        e=explain_code(c); fs.save_file(c); m=f'"""\nEXPLANATION:\n{e}\n"""\n\n{c}'
-        return create_diff_view(c,m),c,m
-    def refactor_fn(c):
-        r=refactor_code(c); fs.save_file(c); return create_diff_view(c,r),c,r
-    def generate_fn(p,h):
-        g=generate_code(p); d=create_diff_view("",g)
-        nh=h+[{"role":"user","content":p},{"role":"assistant","content":"Code generated → DIFF tab"}]
-        return d,"",g,nh,""
-
-    # Snippets
-    def save_snippet(name, tags, editor_content):
-        msg = snippets.add(name, editor_content, tags)
-        return msg, gr.update(choices=snippets.get_names())
-
-    def insert_snippet(name):
-        code = snippets.get(name)
-        return code if code else gr.update()
-
-    def delete_snippet(name):
-        msg = snippets.delete(name)
-        return msg, gr.update(choices=snippets.get_names())
-
-    def search_snippets(query):
-        results = snippets.search(query)
-        return gr.update(choices=results)
-
-    # Find & Replace
-    def do_find(code, query, case):
-        return finder.find_in_file(code, query, case)
-
-    def do_find_all(query):
-        return finder.find_in_all_files(fs.files, query)
-
-    def do_replace(code, find_q, replace_q, case):
-        new_code, msg = finder.replace_in_file(code, find_q, replace_q, case)
-        fs.save_file(new_code)
-        return new_code, msg, analyzer.analyze(new_code)
-
-    # ═══════════════════════════════════════
-    # WIRING
-    # ═══════════════════════════════════════
-
-    # File operations
-    editor.change(update_file, editor, structure_view)
-    file_list.change(load_file, file_list, [editor, editor, structure_view])
-    new_file_btn.click(create_file_fn, new_file_name, [file_list, editor, editor, structure_view])
-    del_file_btn.click(delete_file_fn, None, [file_list, editor, editor, structure_view])
+            res = subprocess.run([sys.executable, path], capture_output=True, text=True, timeout=5)
+            return f"OUTPUT:\n{res.stdout}\nERRORS:\n{res.stderr}"
+        except Exception as e:
+            return str(e)
 
-    # Terminal
-    terminal_input.submit(terminal_cmd, terminal_input, [terminal_output, terminal_input])
-    terminal_run_btn.click(terminal_cmd, terminal_input, [terminal_output, terminal_input])
-    clear_btn.click(clear_term, None, terminal_output)
-    history_refresh.click(refresh_history, None, history_dd)
-    history_dd.change(select_history, history_dd, terminal_input)
-
-    # Run
-    run_btn.click(run_editor, editor, terminal_output)
-    run_gui_btn.click(run_gui, editor, [terminal_output, display_image, gui_status])
-    stop_gui_btn.click(stop_gui, None, [terminal_output, gui_status])
-
-    # AI
-    complete_btn.click(complete_fn, editor, [editor, structure_view])
-    explain_btn.click(explain_fn, editor, [diff_view, diff_original_state, diff_modified_state]
-    ).then(lambda: gr.Tabs(selected="diff-tab"), None, bottom_tabs)
-    refactor_btn.click(refactor_fn, editor, [diff_view, diff_original_state, diff_modified_state]
-    ).then(lambda: gr.Tabs(selected="diff-tab"), None, bottom_tabs)
-    chat_input.submit(generate_fn, [chat_input, chat_history],
-        [diff_view, diff_original_state, diff_modified_state, chat_history, chat_input]
-    ).then(lambda: gr.Tabs(selected="diff-tab"), None, bottom_tabs)
-    send_btn.click(generate_fn, [chat_input, chat_history],
-        [diff_view, diff_original_state, diff_modified_state, chat_history, chat_input]
-    ).then(lambda: gr.Tabs(selected="diff-tab"), None, bottom_tabs)
-
-    # Diff
-    apply_btn.click(lambda mod: (mod, analyzer.analyze(mod)), diff_modified_state, [editor, structure_view]
-    ).then(lambda: gr.Tabs(selected="terminal-tab"), None, bottom_tabs)
-    discard_btn.click(lambda: (gr.update(), gr.update()), None, [editor, structure_view]
-    ).then(lambda: gr.Tabs(selected="terminal-tab"), None, bottom_tabs)
+    def run_gui_code(code):
+        return gui_mgr.launch(code)
 
-    # Snippets
-    snippet_save_btn.click(save_snippet, [snippet_name, snippet_tags, editor], [snippet_status, snippet_list])
-    snippet_insert_btn.click(insert_snippet, snippet_list, editor)
-    snippet_del_btn.click(delete_snippet, snippet_list, [snippet_status, snippet_list])
-    snippet_search.change(search_snippets, snippet_search, snippet_list)
+    def stop_gui_code():
+        return gui_mgr.stop()
 
-    # Find & Replace
-    find_btn.click(do_find, [editor, find_query, find_case], find_results)
-    find_all_btn.click(do_find_all, find_query, find_results)
-    replace_btn.click(do_replace, [editor, find_query, replace_query, find_case],
-                      [editor, find_results, structure_view])
-# ... [Keep your existing FIND tab code] ... 
-
-                # ═══ MISSING TAB: GUI DISPLAY ═══
-                with gr.Tab("🖥 DISPLAY", id="display-tab"):
-                    display_image = gr.Image(label="Virtual Display Output", interactive=False, type="filepath")
-                    with gr.Row():
-                        capture_btn = gr.Button("📸 CAPTURE SCREEN", variant="primary")
-                        auto_capture = gr.Checkbox(label="Auto-refresh (2s)")
-                    # Hidden timer for auto-refresh
-                    auto_timer = gr.Timer(value=2.0, active=False)
+    def terminal_submit(cmd):
+        return terminal.run_command(cmd), gr.update(value="")
 
-    # ═══════════════════════════════════════
-    # BRIDGE FUNCTIONS (These were missing)
-    # ═══════════════════════════════════════
-    
-    # Wrappers for Find/Replace
+    # --- Bridge Functions (Find/Display) ---
     def do_find(code, query, case):
         return finder.find_in_file(code, query, case)
 
@@ -958,16 +462,14 @@ with gr.Blocks(title="Axon Pro — Python AI IDE") as demo:
 
     def do_replace(code, find, replace, case):
         new_code, msg = finder.replace_in_file(code, find, replace, case)
-        # We need to update the editor, the results box, AND the structure view
+        # Update editor, results box, AND structure view
+        fs.save_file(new_code)  # Auto-save on replace
         new_structure = analyzer.analyze(new_code)
         return new_code, msg, new_structure
 
-    # Wrappers for Display
     def capture_display():
         path = vdisplay.capture()
-        if not path:
-            return None 
-        return path
+        return path if path else None
 
     def toggle_auto(checkbox_state):
         return gr.Timer(active=checkbox_state)
@@ -975,33 +477,73 @@ with gr.Blocks(title="Axon Pro — Python AI IDE") as demo:
     def auto_capture_tick():
         return capture_display()
 
+    # --- Snippet Wrappers ---
+    def insert_snippet(code, snippet_name):
+        s_code = snippets.get(snippet_name)
+        if not s_code: return code
+        return (code + "\n" + s_code) if code else s_code
+
+    def save_snippet_ui(name, code, tags):
+        return snippets.add(name, code, tags), gr.update(choices=snippets.get_names())
+
+    # --- AI Mock functions ---
+    def mock_ai_action(code, action):
+        return f"# {action} applied to code (AI Mock)\n" + code
+
     # ═══════════════════════════════════════
-    # EVENT WIRING
+    # 6. EVENT WIRING
     # ═══════════════════════════════════════
     
-    # Find & Replace Events
+    # File Explorer
+    file_list.change(fn=load_file_content, inputs=file_list, outputs=[editor, editor.label, structure_view])
+    refresh_btn.click(fn=refresh_file_list, inputs=None, outputs=file_list)
+    new_file_btn.click(fn=create_new_file, inputs=new_file_name, outputs=[file_list, editor.label, structure_view])
+    save_btn.click(fn=save_file_content, inputs=editor, outputs=None)
+    del_file_btn.click(fn=delete_current_file, inputs=file_list, outputs=[file_list, editor, editor.label, structure_view])
+
+    # Editor Actions
+    run_btn.click(fn=run_console_code, inputs=editor, outputs=terminal_output)
+    run_gui_btn.click(fn=run_gui_code, inputs=editor, outputs=terminal_output)
+    stop_gui_btn.click(fn=stop_gui_code, inputs=None, outputs=terminal_output)
+    
+    # AI Actions
+    complete_btn.click(fn=lambda c: mock_ai_action(c, "Completion"), inputs=editor, outputs=editor)
+    explain_btn.click(fn=lambda c: mock_ai_action(c, "Explanation"), inputs=editor, outputs=terminal_output)
+    refactor_btn.click(fn=lambda c: mock_ai_action(c, "Refactor"), inputs=editor, outputs=editor)
+
+    # Terminal
+    terminal_run_btn.click(fn=terminal_submit, inputs=terminal_input, outputs=[terminal_output, terminal_input])
+    terminal_input.submit(fn=terminal_submit, inputs=terminal_input, outputs=[terminal_output, terminal_input])
+    clear_btn.click(fn=lambda: terminal.run_command("clear"), inputs=None, outputs=terminal_output)
+    history_refresh.click(fn=lambda: gr.update(choices=terminal.get_history_choices()), inputs=None, outputs=history_dd)
+    history_dd.change(fn=lambda x: x, inputs=history_dd, outputs=terminal_input)
+
+    # Find & Replace
     find_btn.click(fn=do_find, inputs=[editor, find_query, find_case], outputs=find_results)
     find_all_btn.click(fn=do_find_all, inputs=[find_query], outputs=find_results)
-    
     replace_btn.click(fn=do_replace, 
                       inputs=[editor, find_query, replace_query, find_case],
                       outputs=[editor, find_results, structure_view])
 
-    # Display Events
+    # Display / Capture
     capture_btn.click(fn=capture_display, inputs=None, outputs=display_image)
     auto_capture.change(fn=toggle_auto, inputs=auto_capture, outputs=auto_timer)
     auto_timer.tick(fn=auto_capture_tick, inputs=None, outputs=display_image)
 
-    # Note: Ensure you have wired up the other buttons from the previous 
-    # part of the code (Run, Stop, Complete, etc.) or they won't work either!
+    # Snippets
+    snippet_insert_btn.click(fn=insert_snippet, inputs=[editor, snippet_list], outputs=editor)
+    snippet_save_btn.click(fn=save_snippet_ui, inputs=[snippet_name, editor, snippet_tags], outputs=[snippet_status, snippet_list])
 
 if __name__ == "__main__":
-    import atexit
+    # Register cleanup handlers
     atexit.register(terminal.cleanup)
     atexit.register(gui_mgr.stop)
     atexit.register(vdisplay.cleanup)
     
-    # Use queue=True for timers and long running tasks
-    demo.queue().launch(server_name="0.0.0.0", server_port=7860, 
-                        theme=theme, css=css, ssr_mode=False,
-                        allowed_paths=["/tmp"])
+    # Launch
+    demo.queue().launch(
+        server_name="0.0.0.0", 
+        server_port=7860, 
+        ssr_mode=False,
+        allowed_paths=["/tmp"]
+    )