Update app.py

app.py

@@ -14,10 +14,12 @@ import struct
 import termios
 import threading
 import re
+import base64
+import tempfile
 from io import StringIO
-import contextlib
+from pathlib import Path
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, GenerationConfig
-from typing import List, Tuple
+from typing import List
 from functools import lru_cache
 
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
@@ -26,6 +28,101 @@ 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
+
+# --- Virtual Display Setup ---
+class VirtualDisplay:
+    """Manages Xvfb virtual framebuffer for capturing GUI app output."""
+
+    def __init__(self):
+        self.xvfb_proc = None
+        self.display = DISPLAY_NUM
+        self.width = SCREEN_W
+        self.height = SCREEN_H
+        self._start_xvfb()
+
+    def _start_xvfb(self):
+        """Start Xvfb virtual framebuffer."""
+        try:
+            # Kill any existing Xvfb on this display
+            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"{self.width}x{self.height}x24", "-ac", "-nolisten", "tcp"],
+                stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
+            )
+            time.sleep(0.5)
+
+            # Verify it started
+            if self.xvfb_proc.poll() is not None:
+                print("[VirtualDisplay] Xvfb failed to start, trying to install...")
+                subprocess.run(["apt-get", "update", "-qq"], capture_output=True, timeout=60)
+                subprocess.run(["apt-get", "install", "-y", "-qq", "xvfb", "x11-utils",
+                                "imagemagick"], capture_output=True, timeout=120)
+                self.xvfb_proc = subprocess.Popen(
+                    ["Xvfb", self.display, "-screen", "0",
+                     f"{self.width}x{self.height}x24", "-ac", "-nolisten", "tcp"],
+                    stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
+                )
+                time.sleep(0.5)
+
+            if self.xvfb_proc.poll() is None:
+                print(f"[VirtualDisplay] Xvfb running on {self.display}")
+                os.environ["DISPLAY"] = self.display
+            else:
+                print("[VirtualDisplay] Xvfb could not start")
+                self.xvfb_proc = None
+
+        except FileNotFoundError:
+            print("[VirtualDisplay] Xvfb not found, installing...")
+            try:
+                subprocess.run(["apt-get", "update", "-qq"], capture_output=True, timeout=60)
+                subprocess.run(["apt-get", "install", "-y", "-qq", "xvfb", "x11-utils",
+                                "imagemagick"], capture_output=True, timeout=120)
+                self._start_xvfb()
+            except Exception as e:
+                print(f"[VirtualDisplay] Install failed: {e}")
+                self.xvfb_proc = None
+        except Exception as e:
+            print(f"[VirtualDisplay] Error: {e}")
+            self.xvfb_proc = None
+
+    def capture_screenshot(self) -> str | None:
+        """Capture current display as base64 PNG. Returns None on failure."""
+        if self.xvfb_proc is None or self.xvfb_proc.poll() is not None:
+            return None
+        try:
+            tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
+            tmp.close()
+            # Try xdotool + import (ImageMagick)
+            result = subprocess.run(
+                ["import", "-window", "root", "-display", self.display, tmp.name],
+                capture_output=True, timeout=5
+            )
+            if result.returncode == 0 and os.path.getsize(tmp.name) > 100:
+                return tmp.name
+            os.unlink(tmp.name)
+        except Exception:
+            pass
+        return None
+
+    @property
+    def is_running(self) -> bool:
+        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 subprocess.TimeoutExpired:
+                self.xvfb_proc.kill()
+
+vdisplay = VirtualDisplay()
 
 # --- File System ---
 class FileSystem:
@@ -58,14 +155,8 @@ class FileSystem:
 
 fs = FileSystem()
 
-# --- PTY Terminal Engine ---
+# --- PTY Terminal ---
 class PTYTerminal:
-    """
-    Real PTY-based terminal. Spawns an actual bash shell so every command
-    that works in a normal Debian/Ubuntu terminal works here — pip, npm,
-    apt-get, git, curl, python, node, you name it.
-    """
-
     STRIP_ANSI = re.compile(r'\x1b\[[0-9;]*[a-zA-Z]|\x1b\].*?\x07|\x1b\[.*?[@-~]|\r')
 
     def __init__(self):
@@ -75,22 +166,18 @@ class PTYTerminal:
         self.log_lines: List[str] = []
         self.max_lines = 800
         self._spawn_shell()
-        self._append_line("Axon Pro Terminal v2.0 — Full PTY Shell")
-        self._append_line("────────────────────────────────────────")
-        self._append_line("This is a real bash shell. Everything works:")
-        self._append_line("  pip install <pkg>     npm install <pkg>")
-        self._append_line("  apt-get install <pkg> git clone <url>")
-        self._append_line("  python script.py      node app.js")
-        self._append_line("  curl, wget, ls, cat, grep, etc.")
-        self._append_line("────────────────────────────────────────")
+        self._append_line("AXON TERMINAL v3.0")
+        self._append_line("══════════════════════════════════════")
+        self._append_line("Full PTY shell — everything works:")
+        self._append_line("  pip/npm/apt-get/git/curl/wget/make")
+        self._append_line("  python/node — GUI apps render to Display")
+        self._append_line("══════════════════════════════════════")
 
     def _spawn_shell(self):
-        """Fork a real bash process attached to a PTY."""
         try:
             pid, fd = pty.openpty()
             self.pid = os.fork()
             if self.pid == 0:
-                # Child process — become the shell
                 os.close(pid)
                 os.setsid()
                 slave_fd = fd
@@ -104,9 +191,9 @@ class PTYTerminal:
                 env["TERM"] = "dumb"
                 env["PS1"] = "$ "
                 env["DEBIAN_FRONTEND"] = "noninteractive"
+                env["DISPLAY"] = DISPLAY_NUM
                 os.execvpe("/bin/bash", ["/bin/bash", "--norc", "--noprofile", "-i"], env)
             else:
-                # Parent process
                 os.close(fd)
                 self.master_fd = pid
                 flags = fcntl.fcntl(self.master_fd, fcntl.F_GETFL)
@@ -116,11 +203,10 @@ class PTYTerminal:
                 time.sleep(0.3)
                 self._read_output()
         except Exception as e:
-            self._append_line(f"[Terminal Error] Failed to spawn shell: {e}")
+            self._append_line(f"[Shell Error] {e}")
             self.master_fd = None
 
     def _read_output(self, timeout: float = 0.1) -> str:
-        """Read all available output from the PTY."""
         if self.master_fd is None:
             return ""
         output = []
@@ -146,7 +232,6 @@ class PTYTerminal:
         return "".join(output)
 
     def _clean_output(self, text: str) -> str:
-        """Strip ANSI escape codes and control characters."""
         cleaned = self.STRIP_ANSI.sub("", text)
         cleaned = "".join(ch for ch in cleaned if ch == "\n" or ch == "\t" or (ord(ch) >= 32))
         return cleaned
@@ -160,102 +245,80 @@ class PTYTerminal:
         return "\n".join(self.log_lines)
 
     def run_command(self, cmd: str) -> str:
-        """Send a command to the real bash shell and capture output."""
         cmd = cmd.strip()
         if not cmd:
             return self.get_log()
-
         with self.lock:
             if cmd.lower() == "clear":
                 self.log_lines = []
                 return ""
-
             if self.master_fd is None:
                 self._append_line(f"$ {cmd}")
-                self._append_line("[Terminal Error] No shell process. Restart the Space.")
+                self._append_line("[Error] No shell. Restart Space.")
                 return self.get_log()
-
-            # Drain any leftover output
             self._read_output(0.05)
-
             self._append_line(f"$ {cmd}")
-
             try:
                 os.write(self.master_fd, (cmd + "\n").encode())
             except OSError as e:
                 self._append_line(f"[Write Error] {e}")
                 return self.get_log()
-
             parts = cmd.split()
             base = parts[0].lower() if parts else ""
-            long_commands = ["pip", "pip3", "npm", "npx", "apt-get", "apt", "git", "wget",
-                             "curl", "make", "cmake", "cargo", "yarn", "conda"]
-
-            if base in long_commands:
-                wait_time = 180
-            elif base in ("python", "python3", "node"):
-                wait_time = 60
-            else:
-                wait_time = 15
-
+            long_cmds = ["pip", "pip3", "npm", "npx", "apt-get", "apt", "git",
+                         "wget", "curl", "make", "cmake", "cargo", "yarn", "conda"]
+            wait = 180 if base in long_cmds else (60 if base in ("python", "python3", "node") else 15)
             all_output = []
             start = time.time()
-            idle_count = 0
-
-            while time.time() - start < wait_time:
+            idle = 0
+            while time.time() - start < wait:
                 chunk = self._read_output(0.3)
                 if chunk:
-                    idle_count = 0
+                    idle = 0
                     all_output.append(chunk)
-                    combined = "".join(all_output)
-                    if combined.rstrip().endswith("$"):
+                    if "".join(all_output).rstrip().endswith("$"):
                         break
                 else:
-                    idle_count += 1
-                    if base not in long_commands and idle_count >= 3:
+                    idle += 1
+                    if base not in long_cmds and idle >= 3:
                         break
-                    if base in long_commands and idle_count >= 10:
+                    if base in long_cmds and idle >= 10:
                         break
-
             raw = "".join(all_output)
             cleaned = self._clean_output(raw)
-
             lines = cleaned.split("\n")
             filtered = []
             skip_echo = True
             for line in lines:
-                stripped = line.strip()
-                if skip_echo and stripped == cmd.strip():
+                s = line.strip()
+                if skip_echo and s == cmd.strip():
                     skip_echo = False
                     continue
-                if stripped == "$" or stripped == "$ ":
+                if s in ("$", "$ "):
                     continue
                 filtered.append(line)
-
             result = "\n".join(filtered).strip()
             if result:
                 self._append_line(result)
-
             return self.get_log()
 
     def run_editor_code(self, code: str) -> str:
-        """Execute code from the editor by writing to a temp file."""
         with self.lock:
             self._append_line("$ python [editor]")
-
-        tmp_path = "/tmp/_axon_editor_run.py"
+        tmp = "/tmp/_axon_run.py"
         try:
-            with open(tmp_path, "w") as f:
+            with open(tmp, "w") as f:
                 f.write(code)
         except Exception as e:
             with self.lock:
-                self._append_line(f"[Error] Could not write temp file: {e}")
+                self._append_line(f"[Error] {e}")
             return self.get_log()
-
         try:
+            env = os.environ.copy()
+            env["DISPLAY"] = DISPLAY_NUM
             result = subprocess.run(
-                [sys.executable, tmp_path],
-                capture_output=True, text=True, timeout=30
+                [sys.executable, tmp], capture_output=True, text=True,
+                timeout=30, env=env
             )
             with self.lock:
                 if result.stdout.strip():
@@ -266,11 +329,10 @@ class PTYTerminal:
                     self._append_line("(No output)")
         except subprocess.TimeoutExpired:
             with self.lock:
-                self._append_line("[Timed out after 30s]")
+                self._append_line("[Timed out 30s]")
         except Exception as e:
             with self.lock:
                 self._append_line(f"[Error] {e}")
-
         return self.get_log()
 
     def cleanup(self):
@@ -290,335 +352,448 @@ terminal = PTYTerminal()
 # --- Model Loading ---
 @lru_cache(maxsize=1)
 def load_model():
-    print("Loading TinyLlama model for CPU...")
-    start_time = time.time()
+    print("Loading TinyLlama...")
+    t0 = time.time()
     try:
-        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-        model = AutoModelForCausalLM.from_pretrained(
-            MODEL_NAME,
-            dtype=DTYPE,
-            device_map=DEVICE,
-            low_cpu_mem_usage=True
-        )
-        gen_config = GenerationConfig(
-            max_new_tokens=256,
-            temperature=0.2,
-            top_p=0.95,
-            do_sample=True,
-        )
-        pipe = pipeline(
-            "text-generation",
-            model=model,
-            tokenizer=tokenizer,
-            generation_config=gen_config,
-        )
-        print(f"Model loaded in {time.time() - start_time:.2f}s")
+        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"Model loaded in {time.time()-t0:.1f}s")
         return pipe
     except Exception as e:
-        print(f"Error loading model: {e}")
+        print(f"Model error: {e}")
         return None
 
-# --- AI Functions ---
-
-def generate_completion(code: str) -> str:
-    model = load_model()
-    if not model:
-        return "# Error: Model failed to load."
-    prompt = f"<|system|>You are a Python expert. Complete the code. No explanations.</s><|user|>\n{code}\n</s><|assistant|>"
+# --- AI helpers ---
+def _llm(system: str, user: str, max_tok: int = 512) -> str:
+    m = load_model()
+    if not m:
+        return "Error: model not loaded"
+    prompt = f"<|system|>{system}</s><|user|>\n{user}\n</s><|assistant|>"
     try:
-        result = model(prompt, num_return_sequences=1, return_full_text=False)
-        return result[0]["generated_text"].strip()
-    except Exception as e:
-        return f"# Error: {str(e)}"
-
-def explain_code(code: str) -> str:
-    model = load_model()
-    if not model:
-        return "Error loading model."
-    prompt = f"<|system|>Explain this Python code concisely.</s><|user|>\n{code}\n</s><|assistant|>"
-    try:
-        result = model(prompt, max_new_tokens=512)
-        return result[0]["generated_text"].strip()
+        r = m(prompt, max_new_tokens=max_tok)
+        return r[0]["generated_text"].strip()
     except Exception as e:
         return str(e)
 
-def refactor_code(code: str) -> str:
-    model = load_model()
-    if not model:
-        return "Error loading model."
-    prompt = f"<|system|>Refactor this Python code for PEP 8 and best practices.</s><|user|>\n{code}\n</s><|assistant|>"
-    try:
-        result = model(prompt, max_new_tokens=512)
-        return result[0]["generated_text"].strip()
-    except Exception as e:
-        return str(e)
-
-def generate_code(prompt_text: str) -> str:
-    model = load_model()
-    if not model:
-        return "Error loading model."
-    prompt = f"<|system|>Write Python code for: {prompt_text}</s><|assistant|>"
-    try:
-        result = model(prompt, max_new_tokens=512)
-        return result[0]["generated_text"].strip()
-    except Exception as e:
-        return str(e)
+def generate_completion(code): return _llm("Complete this Python code. No explanations.", code, 256)
+def explain_code(code): return _llm("Explain this Python code concisely.", code)
+def refactor_code(code): return _llm("Refactor for PEP 8 and best practices.", code)
+def generate_code(prompt): return _llm(f"Write Python code for: {prompt}", "")
 
 def create_diff_view(original: str, modified: str) -> str:
     import difflib
-    diff = difflib.unified_diff(original.splitlines(), modified.splitlines(), lineterm="")
-    return "\n".join(diff)
+    return "\n".join(difflib.unified_diff(original.splitlines(), modified.splitlines(), lineterm=""))
 
-# --- Gradio UI ---
+# ═══════════════════════════════════════════════════════════
+# GRADIO UI — Industrial Hacker Aesthetic
+# ═══════════════════════════════════════════════════════════
 
 theme = gr.themes.Default(
-    font=[gr.themes.GoogleFont("JetBrains Mono"), "monospace"],
-    primary_hue="blue",
+    font=[gr.themes.GoogleFont("IBM Plex Mono"), "monospace"],
+    primary_hue="cyan",
     neutral_hue="gray",
 ).set(
-    button_primary_background_fill="*primary_500",
-    button_primary_background_fill_hover="*primary_600",
-    body_background_fill="*neutral_900",
-    body_text_color="*neutral_50",
+    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="#0ff40",
+    block_background_fill="#0d0d0d",
+    block_border_color="#1a1a1a",
+    block_label_text_color="#0f8",
+    input_background_fill="#0a0a0a",
+    input_border_color="#1a3a2a",
+    input_text_color="#0f8",
 )
 
 css = """
-    * { border-radius: 0 !important; }
-    .terminal-box textarea {
-        background-color: #0a0a0a !important;
-        color: #00ff41 !important;
-        font-family: 'JetBrains Mono', monospace !important;
-        font-size: 13px !important;
-        line-height: 1.4 !important;
-    }
-    .terminal-input input {
-        background-color: #0a0a0a !important;
-        color: #00ff41 !important;
-        font-family: 'JetBrains Mono', monospace !important;
-        border: 1px solid #1a3a1a !important;
-        font-size: 13px !important;
-    }
-    .terminal-input input::placeholder {
-        color: #2a6a2a !important;
-    }
-    .diff-view { height: 40vh; }
-    .code-wrap { white-space: pre-wrap !important; }
-    .status-bar { font-size: 12px; opacity: 0.7; }
+/* ═══ GLOBAL ═══ */
+* { border-radius: 0 !important; }
+.gradio-container { max-width: 100% !important; }
+
+/* ═══ HEADER ═══ */
+.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;
+}
+
+/* ═══ PANELS ═══ */
+.panel-section {
+    border: 1px solid #1a2a2a !important;
+    background: #0a0a0a !important;
+    transition: all 0.3s cubic-bezier(0.4, 0, 0.2, 1);
+}
+.panel-section:hover {
+    border-color: #0ff3 !important;
+}
+
+/* ═══ TERMINAL ═══ */
+.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;
+    caret-color: #0f8 !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 ═══ */
+.code-editor textarea, .code-editor .cm-editor {
+    font-family: 'IBM Plex Mono', monospace !important;
+    font-size: 13px !important;
+}
+
+/* ═══ FILE EXPLORER ═══ */
+.file-explorer {
+    border-right: 1px solid #1a2a2a !important;
+}
+.file-explorer .wrap { border-color: #0f32 !important; }
+
+/* ═══ DISPLAY PANEL ═══ */
+.display-panel {
+    border: 1px solid #1a2a2a !important;
+    background: #050505 !important;
+    min-height: 200px;
+}
+.display-panel img {
+    image-rendering: pixelated;
+}
+
+/* ═══ TOOLBAR ═══ */
+.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;
+}
+
+/* ═══ COLLAPSIBLE SECTIONS (Accordion) ═══ */
+.gr-accordion { border: 1px solid #1a2a2a !important; background: #0a0a0a !important; }
+.gr-accordion > .label-wrap {
+    background: #0d0d0d !important;
+    border-bottom: 1px solid #1a2a2a !important;
+    transition: background 0.2s ease;
+}
+.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 ═══ */
+.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;
+    transition: all 0.2s ease;
+}
+.tabs > .tab-nav > button.selected {
+    color: #0ff !important;
+    border-bottom-color: #0ff !important;
+    text-shadow: 0 0 8px #0ff4;
+}
+
+/* ═══ STATUS BAR ═══ */
+.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; }
+
+/* ═══ SCROLLBAR ═══ */
+::-webkit-scrollbar { width: 6px; height: 6px; }
+::-webkit-scrollbar-track { background: #0a0a0a; }
+::-webkit-scrollbar-thumb { background: #0f83; }
+::-webkit-scrollbar-thumb:hover { background: #0ff; }
+
+/* ═══ CHATBOT ═══ */
+.chatbot-container .message { border-radius: 0 !important; font-family: 'IBM Plex Mono', monospace !important; }
 """
 
-with gr.Blocks(title="Axon Pro - Python IDE") as demo:
-
-    gr.Markdown("# 🐍 Axon Pro — Python AI IDE")
-
-    with gr.Tabs() as main_tabs:
-        with gr.Tab("Editor", id="editor-tab"):
-            with gr.Row(equal_height=True):
-                with gr.Column(scale=1, min_width=200):
-                    gr.Markdown("### 📁 Explorer")
-                    file_list = gr.Dropdown(
-                        choices=fs.get_all_files(),
-                        value=fs.current_file,
-                        label="Files",
-                        interactive=True,
-                        container=False,
-                    )
-                    with gr.Row():
-                        new_file_btn = gr.Button("➕ New", variant="secondary")
-                        save_btn = gr.Button("💾 Save", variant="secondary")
-                    new_file_name = gr.Textbox(
-                        placeholder="script.py", label="New File", container=False
-                    )
+# ═══════════════════════════════════════════════════════════
+
+with gr.Blocks(title="Axon Pro — Python AI IDE") as demo:
+
+    # ── HEADER ──
+    with gr.Row(elem_classes="axon-header"):
+        gr.Markdown("# ⬡ AXON PRO\n\nPYTHON AI IDE — v3.0")
+
+    # ── MAIN LAYOUT ──
+    with gr.Row(equal_height=False):
+
+        # ══ LEFT: FILE EXPLORER (Collapsible) ══
+        with gr.Column(scale=1, min_width=180, elem_classes="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")
+
+        # ══ CENTER: EDITOR + TERMINAL ══
+        with gr.Column(scale=4):
+
+            # Code Editor
+            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.Column(scale=4):
-                    editor = gr.Code(
-                        value=fs.get_current_file_content(),
-                        label="Code Editor",
-                        language="python",
-                        lines=20,
-                        interactive=True,
-                        elem_classes="code-wrap",
+            # Toolbar
+            with gr.Row():
+                run_btn = gr.Button("▶ RUN", variant="primary", elem_classes="toolbar-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")
+
+            # Bottom panels (Collapsible)
+            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,
                     )
-
                     with gr.Row():
-                        run_btn = gr.Button("▶ Run", variant="primary")
-                        complete_btn = gr.Button("✨ Complete", variant="secondary")
-                        explain_btn = gr.Button("📝 Explain", variant="secondary")
-                        refactor_btn = gr.Button("🔧 Refactor", variant="secondary")
-
-                    with gr.Tabs():
-                        with gr.Tab("Terminal", id="terminal-tab"):
-                            terminal_output = gr.Textbox(
-                                value=terminal.get_log(),
-                                lines=12,
-                                max_lines=30,
-                                interactive=False,
-                                elem_classes="terminal-box",
-                                label="Terminal",
-                                show_label=False,
-                            )
-                            with gr.Row():
-                                terminal_input = gr.Textbox(
-                                    placeholder="$ type any command: pip install pygame, npm i, git clone, ls ...",
-                                    scale=8,
-                                    container=False,
-                                    elem_classes="terminal-input",
-                                )
-                                terminal_run_btn = gr.Button(
-                                    "⏎", variant="primary", scale=1, min_width=50
-                                )
-                            clear_btn = gr.Button(
-                                "Clear", variant="secondary", size="sm"
-                            )
-
-                        with gr.Tab("AI Chat", id="chat-tab"):
-                            chat_history = gr.Chatbot(
-                                label="Axon AI", height=300
-                            )
-                            with gr.Row():
-                                chat_input = gr.Textbox(
-                                    placeholder="Ask to generate Python code...",
-                                    scale=7,
-                                    container=False,
-                                )
-                                send_btn = gr.Button(
-                                    "Generate", variant="primary", scale=1
-                                )
-
-        with gr.Tab("Diff View", id="diff-tab"):
-            diff_view = gr.Code(
-                label="AI Changes",
-                language="python",
-                elem_classes="diff-view code-wrap",
-                interactive=False,
-            )
-            with gr.Row():
-                apply_btn = gr.Button("Apply Changes", variant="primary")
-                discard_btn = gr.Button("Discard Changes", variant="secondary")
-
+                        terminal_input = gr.Textbox(
+                            placeholder="$ pip install pygame  |  npm i  |  git clone  |  ls ...",
+                            scale=9, container=False, elem_classes="terminal-input",
+                        )
+                        terminal_run_btn = gr.Button("⏎", variant="primary", scale=1, min_width=45)
+                    clear_btn = gr.Button("CLEAR", variant="secondary", size="sm",
+                                           elem_classes="toolbar-btn")
+
+                with gr.Tab("💬 AI CHAT", id="chat-tab"):
+                    chat_history = gr.Chatbot(label="", height=250,
+                                               elem_classes="chatbot-container")
+                    with gr.Row():
+                        chat_input = gr.Textbox(
+                            placeholder="Describe the code you want...",
+                            scale=8, container=False,
+                        )
+                        send_btn = gr.Button("GEN", variant="primary", scale=1,
+                                              elem_classes="toolbar-btn")
+
+                with gr.Tab("⇄ DIFF", id="diff-tab"):
+                    diff_view = gr.Code(label="AI Changes", language="python",
+                                         interactive=False, lines=15)
+                    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 PANEL (Collapsible) ══
+        with gr.Column(scale=2, min_width=250):
+            with gr.Accordion("🖥 DISPLAY OUTPUT", open=True):
+                gr.Markdown(
+                    "<small style='color:#0f84'>GUI apps (pygame, tkinter, etc.) "
+                    "render here via virtual display</small>"
+                )
+                display_image = gr.Image(
+                    label="", type="filepath", interactive=False,
+                    elem_classes="display-panel", height=400,
+                )
+                with gr.Row():
+                    capture_btn = gr.Button("📸 CAPTURE", variant="secondary", size="sm",
+                                             elem_classes="toolbar-btn")
+                    auto_refresh = gr.Checkbox(label="Auto-refresh", value=False)
+
+                display_status = gr.Markdown(
+                    f"<small style='color:#0f84'>Xvfb: "
+                    f"{'●  running' if vdisplay.is_running else '○  stopped'} "
+                    f"| {SCREEN_W}x{SCREEN_H}</small>"
+                )
+
+    # ── STATUS BAR ──
     status_bar = gr.Markdown(
-        "**AXON PRO v2.0** | Python 3.x | CPU Mode | TinyLlama-1.1B | Full PTY Terminal",
+        f"**AXON PRO v3.0** │ Python {sys.version.split()[0]} │ CPU │ "
+        f"TinyLlama-1.1B │ PTY Shell │ "
+        f"Xvfb {'ON' if vdisplay.is_running else 'OFF'}",
         elem_classes="status-bar",
     )
 
-    # State
-    current_file_state = gr.State(fs.current_file)
+    # ── STATE ──
     diff_original_state = gr.State("")
     diff_modified_state = gr.State("")
-    diff_mode_state = gr.State(False)
 
-    # --- Handlers ---
+    # ══════════════════════════════════════════
+    # HANDLERS
+    # ══════════════════════════════════════════
+
     def update_file(content):
         fs.save_file(content)
-        return fs.get_current_file_content()
 
     def load_file(filename):
         fs.set_current_file(filename)
-        return fs.get_current_file_content(), filename
+        return fs.get_current_file_content(), gr.update(label=f"  {filename}")
 
-    def create_file(name):
+    def create_file_fn(name):
         if name and "." in name:
             fs.create_file(name)
-            return (
-                gr.update(choices=fs.get_all_files(), value=name),
+            return (gr.update(choices=fs.get_all_files(), value=name),
+                    fs.get_current_file_content(),
+                    gr.update(label=f"  {name}"))
+        return (gr.update(choices=fs.get_all_files()),
                 fs.get_current_file_content(),
-            )
-        return gr.update(choices=fs.get_all_files()), fs.get_current_file_content()
+                gr.update())
 
-    def terminal_command(cmd):
+    def terminal_cmd(cmd):
         log = terminal.run_command(cmd)
         return log, ""
 
-    def clear_terminal_fn():
+    def clear_term():
         terminal.log_lines = []
         return ""
 
-    def run_editor_wrapper(content):
+    def run_editor(content):
         return terminal.run_editor_code(content)
 
-    def complete_wrapper(content):
-        comp = generate_completion(content)
-        return content + "\n" + comp
+    def complete_fn(content):
+        return content + "\n" + generate_completion(content)
 
-    def explain_wrapper(content):
+    def explain_fn(content):
         explanation = explain_code(content)
         fs.save_file(content)
-        diff_orig = content
-        diff_mod = f'"""\nEXPLANATION:\n{explanation}\n"""\n\n{content}'
-        diff = create_diff_view(diff_orig, diff_mod)
-        return diff, diff_orig, diff_mod, True
+        orig = content
+        mod = f'"""\nEXPLANATION:\n{explanation}\n"""\n\n{content}'
+        return create_diff_view(orig, mod), orig, mod
 
-    def refactor_wrapper(content):
+    def refactor_fn(content):
         refactored = refactor_code(content)
         fs.save_file(content)
-        diff = create_diff_view(content, refactored)
-        return diff, content, refactored, True
+        return create_diff_view(content, refactored), content, refactored
 
-    def generate_wrapper(prompt_text, history):
+    def generate_fn(prompt_text, history):
         generated = generate_code(prompt_text)
         diff = create_diff_view("", generated)
         new_history = history + [
             {"role": "user", "content": prompt_text},
-            {"role": "assistant", "content": "Generated code is available in the Diff View tab."},
+            {"role": "assistant", "content": "Code generated → check DIFF tab"},
         ]
-        return diff, "", generated, True, new_history, ""
+        return diff, "", generated, new_history, ""
+
+    def capture_display():
+        path = vdisplay.capture_screenshot()
+        if path:
+            return path
+        return None
+
+    # ══════════════════════════════════════════
+    # WIRING
+    # ══════════════════════════════════════════
 
-    # --- Wiring ---
     editor.change(update_file, editor, None)
-    file_list.change(load_file, file_list, [editor, current_file_state])
-    new_file_btn.click(create_file, new_file_name, [file_list, editor])
+    file_list.change(load_file, file_list, [editor, editor])
+    new_file_btn.click(create_file_fn, new_file_name, [file_list, editor, editor])
 
-    terminal_input.submit(
-        terminal_command, terminal_input, [terminal_output, terminal_input]
-    )
-    terminal_run_btn.click(
-        terminal_command, terminal_input, [terminal_output, terminal_input]
-    )
-    clear_btn.click(clear_terminal_fn, None, terminal_output)
+    # 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)
 
-    run_btn.click(run_editor_wrapper, editor, terminal_output)
-    complete_btn.click(complete_wrapper, editor, editor)
+    # Editor actions
+    run_btn.click(run_editor, editor, terminal_output)
+    complete_btn.click(complete_fn, editor, editor)
 
     explain_btn.click(
-        explain_wrapper,
-        editor,
-        [diff_view, diff_original_state, diff_modified_state, diff_mode_state],
-    ).then(lambda: gr.Tabs(selected="diff-tab"), None, main_tabs)
+        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_wrapper,
-        editor,
-        [diff_view, diff_original_state, diff_modified_state, diff_mode_state],
-    ).then(lambda: gr.Tabs(selected="diff-tab"), None, main_tabs)
+        refactor_fn, editor,
+        [diff_view, diff_original_state, diff_modified_state]
+    ).then(lambda: gr.Tabs(selected="diff-tab"), None, bottom_tabs)
 
+    # Chat
     chat_input.submit(
-        generate_wrapper,
-        [chat_input, chat_history],
-        [diff_view, diff_original_state, diff_modified_state, diff_mode_state, chat_history, chat_input],
-    ).then(lambda: gr.Tabs(selected="diff-tab"), None, main_tabs)
+        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_wrapper,
-        [chat_input, chat_history],
-        [diff_view, diff_original_state, diff_modified_state, diff_mode_state, chat_history, chat_input],
-    ).then(lambda: gr.Tabs(selected="diff-tab"), None, main_tabs)
+        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 actions
     apply_btn.click(
-        lambda mod: (mod, False), diff_modified_state, [editor, diff_mode_state]
-    ).then(lambda: gr.Tabs(selected="editor-tab"), None, main_tabs)
+        lambda mod: mod, diff_modified_state, editor
+    ).then(lambda: gr.Tabs(selected="terminal-tab"), None, bottom_tabs)
 
     discard_btn.click(
-        lambda: (gr.update(), False), None, [editor, diff_mode_state]
-    ).then(lambda: gr.Tabs(selected="editor-tab"), None, main_tabs)
+        lambda: gr.update(), None, editor
+    ).then(lambda: gr.Tabs(selected="terminal-tab"), None, bottom_tabs)
+
+    # Display capture
+    capture_btn.click(capture_display, None, display_image)
 
 
 if __name__ == "__main__":
     import atexit
     atexit.register(terminal.cleanup)
+    atexit.register(vdisplay.cleanup)
 
     demo.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        theme=theme,
-        css=css,
-        ssr_mode=False,
+        server_name="0.0.0.0", server_port=7860,
+        theme=theme, css=css, ssr_mode=False,
     )