Update app.py

app.py

@@ -14,10 +14,7 @@ import struct
 import termios
 import threading
 import re
-import base64
 import tempfile
-from io import StringIO
-from pathlib import Path
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, GenerationConfig
 from typing import List
 from functools import lru_cache
@@ -31,78 +28,62 @@ DTYPE = torch.float32
 DISPLAY_NUM = ":99"
 SCREEN_W, SCREEN_H = 800, 600
 
-# --- Virtual Display Setup ---
+# --- Virtual Display ---
 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
-            )
+                 f"{SCREEN_W}x{SCREEN_H}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...")
+                print("[Xvfb] Failed, installing...")
                 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
-                )
+                     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"[VirtualDisplay] Xvfb running on {self.display}")
+                print(f"[Xvfb] Running on {self.display}")
                 os.environ["DISPLAY"] = self.display
             else:
-                print("[VirtualDisplay] Xvfb could not start")
+                print("[Xvfb] Could not start")
                 self.xvfb_proc = None
-
         except FileNotFoundError:
-            print("[VirtualDisplay] Xvfb not found, installing...")
+            print("[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}")
+                print(f"[Xvfb] Install failed: {e}")
                 self.xvfb_proc = None
         except Exception as e:
-            print(f"[VirtualDisplay] Error: {e}")
+            print(f"[Xvfb] 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:
+        if not self.is_running:
             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
-            )
+                capture_output=True, timeout=5)
             if result.returncode == 0 and os.path.getsize(tmp.name) > 100:
                 return tmp.name
             os.unlink(tmp.name)
@@ -166,12 +147,12 @@ class PTYTerminal:
         self.log_lines: List[str] = []
         self.max_lines = 800
         self._spawn_shell()
-        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("══════════════════════════════════════")
+        self._append("AXON TERMINAL v3.0")
+        self._append("══════════════════════════════════════")
+        self._append("Full PTY shell — everything works:")
+        self._append("  pip/npm/apt-get/git/curl/wget/make")
+        self._append("  python/node — GUI apps render to Display")
+        self._append("══════════════════════════════════════")
 
     def _spawn_shell(self):
         try:
@@ -180,63 +161,50 @@ class PTYTerminal:
             if self.pid == 0:
                 os.close(pid)
                 os.setsid()
-                slave_fd = fd
-                fcntl.ioctl(slave_fd, termios.TIOCSCTTY, 0)
-                os.dup2(slave_fd, 0)
-                os.dup2(slave_fd, 1)
-                os.dup2(slave_fd, 2)
-                if slave_fd > 2:
-                    os.close(slave_fd)
+                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["TERM"] = "dumb"
-                env["PS1"] = "$ "
-                env["DEBIAN_FRONTEND"] = "noninteractive"
-                env["DISPLAY"] = DISPLAY_NUM
+                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
                 flags = fcntl.fcntl(self.master_fd, fcntl.F_GETFL)
                 fcntl.fcntl(self.master_fd, fcntl.F_SETFL, flags | os.O_NONBLOCK)
-                winsize = struct.pack("HHHH", 40, 120, 0, 0)
-                fcntl.ioctl(self.master_fd, termios.TIOCSWINSZ, winsize)
+                fcntl.ioctl(self.master_fd, termios.TIOCSWINSZ,
+                            struct.pack("HHHH", 40, 120, 0, 0))
                 time.sleep(0.3)
-                self._read_output()
+                self._read(0.2)
         except Exception as e:
-            self._append_line(f"[Shell Error] {e}")
+            self._append(f"[Shell Error] {e}")
             self.master_fd = None
 
-    def _read_output(self, timeout: float = 0.1) -> str:
-        if self.master_fd is None:
-            return ""
-        output = []
+    def _read(self, timeout=0.1) -> str:
+        if self.master_fd is None: return ""
+        out = []
         deadline = time.time() + timeout
         while True:
-            remaining = deadline - time.time()
-            if remaining <= 0:
-                break
+            rem = deadline - time.time()
+            if rem <= 0: break
             try:
-                ready, _, _ = select.select([self.master_fd], [], [], min(remaining, 0.05))
-                if ready:
-                    chunk = os.read(self.master_fd, 4096)
-                    if chunk:
-                        output.append(chunk.decode("utf-8", errors="replace"))
+                r, _, _ = select.select([self.master_fd], [], [], min(rem, 0.05))
+                if r:
+                    c = os.read(self.master_fd, 4096)
+                    if c:
+                        out.append(c.decode("utf-8", errors="replace"))
                         deadline = time.time() + 0.15
-                    else:
-                        break
-                else:
-                    if output:
-                        break
-            except OSError:
-                break
-        return "".join(output)
-
-    def _clean_output(self, text: str) -> str:
-        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
-
-    def _append_line(self, text: str):
+                    else: break
+                elif out: break
+            except OSError: break
+        return "".join(out)
+
+    def _clean(self, text):
+        c = self.STRIP_ANSI.sub("", text)
+        return "".join(ch for ch in c if ch in "\n\t" or ord(ch) >= 32)
+
+    def _append(self, text):
         self.log_lines.append(text)
         while len(self.log_lines) > self.max_lines:
             self.log_lines.pop(0)
@@ -246,110 +214,92 @@ class PTYTerminal:
 
     def run_command(self, cmd: str) -> str:
         cmd = cmd.strip()
-        if not cmd:
-            return self.get_log()
+        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("[Error] No shell. Restart Space.")
+                self._append(f"$ {cmd}")
+                self._append("[Error] No shell. Restart Space.")
                 return self.get_log()
-            self._read_output(0.05)
-            self._append_line(f"$ {cmd}")
+            self._read(0.05)
+            self._append(f"$ {cmd}")
             try:
                 os.write(self.master_fd, (cmd + "\n").encode())
             except OSError as e:
-                self._append_line(f"[Write Error] {e}")
+                self._append(f"[Write Error] {e}")
                 return self.get_log()
+
             parts = cmd.split()
             base = parts[0].lower() if parts else ""
-            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 = []
+            long = ["pip","pip3","npm","npx","apt-get","apt","git",
+                    "wget","curl","make","cmake","cargo","yarn","conda"]
+            wait = 180 if base in long else (60 if base in ("python","python3","node") else 15)
+
+            chunks = []
             start = time.time()
             idle = 0
             while time.time() - start < wait:
-                chunk = self._read_output(0.3)
-                if chunk:
-                    idle = 0
-                    all_output.append(chunk)
-                    if "".join(all_output).rstrip().endswith("$"):
-                        break
+                c = self._read(0.3)
+                if c:
+                    idle = 0; chunks.append(c)
+                    if "".join(chunks).rstrip().endswith("$"): break
                 else:
                     idle += 1
-                    if base not in long_cmds and idle >= 3:
-                        break
-                    if base in long_cmds and idle >= 10:
-                        break
-            raw = "".join(all_output)
-            cleaned = self._clean_output(raw)
-            lines = cleaned.split("\n")
+                    if base not in long and idle >= 3: break
+                    if base in long and idle >= 10: break
+
+            raw = self._clean("".join(chunks))
+            lines = raw.split("\n")
             filtered = []
-            skip_echo = True
-            for line in lines:
-                s = line.strip()
-                if skip_echo and s == cmd.strip():
-                    skip_echo = False
-                    continue
-                if s in ("$", "$ "):
-                    continue
-                filtered.append(line)
+            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_line(result)
+            if result: self._append(result)
             return self.get_log()
 
     def run_editor_code(self, code: str) -> str:
         with self.lock:
-            self._append_line("$ python [editor]")
+            self._append("$ python [editor]")
         tmp = "/tmp/_axon_run.py"
         try:
-            with open(tmp, "w") as f:
-                f.write(code)
+            with open(tmp, "w") as f: f.write(code)
         except Exception as e:
-            with self.lock:
-                self._append_line(f"[Error] {e}")
+            with self.lock: self._append(f"[Error] {e}")
             return self.get_log()
         try:
             env = os.environ.copy()
             env["DISPLAY"] = DISPLAY_NUM
-            result = subprocess.run(
-                [sys.executable, tmp], capture_output=True, text=True,
-                timeout=30, env=env
-            )
+            r = subprocess.run([sys.executable, tmp], capture_output=True,
+                               text=True, timeout=30, env=env)
             with self.lock:
-                if result.stdout.strip():
-                    self._append_line(result.stdout.rstrip())
-                if result.stderr.strip():
-                    self._append_line(result.stderr.rstrip())
-                if not result.stdout.strip() and not result.stderr.strip():
-                    self._append_line("(No output)")
+                if r.stdout.strip(): self._append(r.stdout.rstrip())
+                if r.stderr.strip(): self._append(r.stderr.rstrip())
+                if not r.stdout.strip() and not r.stderr.strip():
+                    self._append("(No output)")
         except subprocess.TimeoutExpired:
-            with self.lock:
-                self._append_line("[Timed out 30s]")
+            with self.lock: self._append("[Timed out 30s]")
         except Exception as e:
-            with self.lock:
-                self._append_line(f"[Error] {e}")
+            with self.lock: self._append(f"[Error] {e}")
         return self.get_log()
 
     def cleanup(self):
         if self.pid and self.pid > 0:
-            try:
-                os.kill(self.pid, signal.SIGTERM)
-            except ProcessLookupError:
-                pass
+            try: os.kill(self.pid, signal.SIGTERM)
+            except ProcessLookupError: pass
         if self.master_fd is not None:
-            try:
-                os.close(self.master_fd)
-            except OSError:
-                pass
+            try: os.close(self.master_fd)
+            except OSError: pass
 
 terminal = PTYTerminal()
 
-# --- Model Loading ---
+# --- Model ---
 @lru_cache(maxsize=1)
 def load_model():
     print("Loading TinyLlama...")
@@ -360,35 +310,32 @@ def load_model():
             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")
+        print(f"Loaded in {time.time()-t0:.1f}s")
         return pipe
     except Exception as e:
         print(f"Model error: {e}")
         return None
 
-# --- AI helpers ---
-def _llm(system: str, user: str, max_tok: int = 512) -> str:
+def _llm(system, user, max_tok=512):
     m = load_model()
-    if not m:
-        return "Error: model not loaded"
+    if not m: return "Error: model not loaded"
     prompt = f"<|system|>{system}</s><|user|>\n{user}\n</s><|assistant|>"
     try:
-        r = m(prompt, max_new_tokens=max_tok)
-        return r[0]["generated_text"].strip()
+        return m(prompt, max_new_tokens=max_tok)[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 generate_code(p): return _llm(f"Write Python code for: {p}", "")
 
-def create_diff_view(original: str, modified: str) -> str:
+def create_diff_view(orig, mod):
     import difflib
-    return "\n".join(difflib.unified_diff(original.splitlines(), modified.splitlines(), lineterm=""))
+    return "\n".join(difflib.unified_diff(orig.splitlines(), mod.splitlines(), lineterm=""))
 
 # ═══════════════════════════════════════════════════════════
-# GRADIO UI — Industrial Hacker Aesthetic
+# THEME — only valid params
 # ═══════════════════════════════════════════════════════════
 
 theme = gr.themes.Default(
@@ -404,21 +351,24 @@ theme = gr.themes.Default(
     button_secondary_background_fill="#1a1a1a",
     button_secondary_background_fill_hover="#2a2a2a",
     button_secondary_text_color="#0ff",
-    button_secondary_border_color="#0ff40",
+    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_text_color="#0f8",
+    input_placeholder_color="#0f83",
 )
 
+# ═══════════════════════════════════════════════════════════
+# CSS
+# ═══════════════════════════════════════════════════════════
+
 css = """
-/* ═══ GLOBAL ═══ */
 * { border-radius: 0 !important; }
 .gradio-container { max-width: 100% !important; }
 
-/* ═══ HEADER ═══ */
+/* ─── Header ─── */
 .axon-header {
     background: linear-gradient(90deg, #080808 0%, #0a1a1a 50%, #080808 100%);
     border-bottom: 1px solid #0ff3;
@@ -434,23 +384,9 @@ css = """
     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;
-}
+.axon-header p { color: #0f84 !important; font-size: 0.75em; letter-spacing: 1px; }
 
-/* ═══ TERMINAL ═══ */
+/* ─── Terminal ─── */
 .terminal-box textarea {
     background: #050505 !important;
     color: #00ff41 !important;
@@ -471,29 +407,20 @@ css = """
 .terminal-input input::placeholder { color: #0f83 !important; }
 .terminal-input input:focus { border-color: #0ff !important; box-shadow: 0 0 8px #0ff2 !important; }
 
-/* ═══ CODE EDITOR ═══ */
+/* ─── 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; }
+/* ─── File explorer ─── */
+.file-explorer { border-right: 1px solid #1a2a2a !important; }
 
-/* ═══ DISPLAY PANEL ═══ */
-.display-panel {
-    border: 1px solid #1a2a2a !important;
-    background: #050505 !important;
-    min-height: 200px;
-}
-.display-panel img {
-    image-rendering: pixelated;
-}
+/* ─── Display ─── */
+.display-panel { border: 1px solid #1a2a2a !important; background: #050505 !important; }
+.display-panel img { image-rendering: pixelated; }
 
-/* ═══ TOOLBAR ═══ */
+/* ─── Toolbar buttons ─── */
 .toolbar-btn {
     font-size: 0.8em !important;
     padding: 6px 12px !important;
@@ -503,12 +430,9 @@ css = """
     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;
-}
+.toolbar-btn:hover { border-color: #0ff !important; box-shadow: 0 0 12px #0ff2 !important; }
 
-/* ═══ COLLAPSIBLE SECTIONS (Accordion) ═══ */
+/* ─── Accordion ─── */
 .gr-accordion { border: 1px solid #1a2a2a !important; background: #0a0a0a !important; }
 .gr-accordion > .label-wrap {
     background: #0d0d0d !important;
@@ -523,112 +447,92 @@ css = """
     font-size: 0.85em;
 }
 
-/* ═══ TABS ═══ */
+/* ─── 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;
+    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;
+    color: #0ff !important; border-bottom-color: #0ff !important;
     text-shadow: 0 0 8px #0ff4;
 }
 
-/* ═══ STATUS BAR ═══ */
+/* ─── Status ─── */
 .status-bar {
-    background: #0a0a0a !important;
-    border-top: 1px solid #1a2a2a;
-    padding: 4px 12px !important;
-    font-size: 11px !important;
-    color: #0f84 !important;
+    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 ═══ */
+/* ─── 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; }
+/* ─── Generic input text color (covers what input_text_color would do) ─── */
+input, textarea, select { color: #0f8 !important; }
 """
 
+# ═══════════════════════════════════════════════════════════
+# UI LAYOUT
 # ═══════════════════════════════════════════════════════════
 
 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) ══
+        # ── LEFT: FILE EXPLORER ──
         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
-                )
+                    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 ══
+        # ── 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",
-            )
+                elem_classes="code-editor")
 
-            # 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,
-                    )
+                        label="", show_label=False)
                     with gr.Row():
                         terminal_input = gr.Textbox(
                             placeholder="$ pip install pygame  |  npm i  |  git clone  |  ls ...",
-                            scale=9, container=False, elem_classes="terminal-input",
-                        )
+                            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")
+                    chat_history = gr.Chatbot(label="", height=250)
                     with gr.Row():
-                        chat_input = gr.Textbox(
-                            placeholder="Describe the code you want...",
-                            scale=8, container=False,
-                        )
+                        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")
 
@@ -641,44 +545,34 @@ with gr.Blocks(title="Axon Pro — Python AI IDE") as demo:
                         discard_btn = gr.Button("✗ DISCARD", variant="secondary",
                                                   elem_classes="toolbar-btn")
 
-        # ══ RIGHT: DISPLAY PANEL (Collapsible) ══
+        # ── RIGHT: DISPLAY ──
         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>"
-                )
+                    "render here via virtual display</small>")
                 display_image = gr.Image(
                     label="", type="filepath", interactive=False,
-                    elem_classes="display-panel", height=400,
-                )
+                    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>"
-                )
+                    f"{'● running' if vdisplay.is_running else '○ stopped'} "
+                    f"| {SCREEN_W}x{SCREEN_H}</small>")
 
-    # ── STATUS BAR ──
     status_bar = gr.Markdown(
         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",
-    )
+        f"TinyLlama-1.1B │ PTY Shell │ Xvfb {'ON' if vdisplay.is_running else 'OFF'}",
+        elem_classes="status-bar")
 
-    # ── STATE ──
+    # ── State ──
     diff_original_state = gr.State("")
     diff_modified_state = gr.State("")
 
-    # ═══════════════════════���══════════════════
-    # HANDLERS
-    # ══════════════════════════════════════════
-
+    # ── Handlers ──
     def update_file(content):
         fs.save_file(content)
 
@@ -690,15 +584,11 @@ with gr.Blocks(title="Axon Pro — Python AI IDE") as demo:
         if name and "." in name:
             fs.create_file(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(),
-                gr.update())
+                    fs.get_current_file_content(), gr.update(label=f"  {name}"))
+        return gr.update(choices=fs.get_all_files()), fs.get_current_file_content(), gr.update()
 
     def terminal_cmd(cmd):
-        log = terminal.run_command(cmd)
-        return log, ""
+        return terminal.run_command(cmd), ""
 
     def clear_term():
         terminal.log_lines = []
@@ -711,89 +601,67 @@ with gr.Blocks(title="Axon Pro — Python AI IDE") as demo:
         return content + "\n" + generate_completion(content)
 
     def explain_fn(content):
-        explanation = explain_code(content)
+        exp = explain_code(content)
         fs.save_file(content)
-        orig = content
-        mod = f'"""\nEXPLANATION:\n{explanation}\n"""\n\n{content}'
-        return create_diff_view(orig, mod), orig, mod
+        mod = f'"""\nEXPLANATION:\n{exp}\n"""\n\n{content}'
+        return create_diff_view(content, mod), content, mod
 
     def refactor_fn(content):
-        refactored = refactor_code(content)
+        ref = refactor_code(content)
         fs.save_file(content)
-        return create_diff_view(content, refactored), content, refactored
+        return create_diff_view(content, ref), content, ref
 
     def generate_fn(prompt_text, history):
-        generated = generate_code(prompt_text)
-        diff = create_diff_view("", generated)
-        new_history = history + [
+        gen = generate_code(prompt_text)
+        diff = create_diff_view("", gen)
+        new_h = history + [
             {"role": "user", "content": prompt_text},
             {"role": "assistant", "content": "Code generated → check DIFF tab"},
         ]
-        return diff, "", generated, new_history, ""
+        return diff, "", gen, new_h, ""
 
     def capture_display():
-        path = vdisplay.capture_screenshot()
-        if path:
-            return path
-        return None
-
-    # ══════════════════════════════════════════
-    # WIRING
-    # ══════════════════════════════════════════
+        return vdisplay.capture_screenshot()
 
+    # ── Wiring ──
     editor.change(update_file, editor, None)
     file_list.change(load_file, file_list, [editor, editor])
     new_file_btn.click(create_file_fn, new_file_name, [file_list, editor, editor])
 
-    # 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)
 
-    # Editor actions
     run_btn.click(run_editor, editor, terminal_output)
     complete_btn.click(complete_fn, editor, editor)
 
-    explain_btn.click(
-        explain_fn, editor,
+    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,
+    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
-    chat_input.submit(
-        generate_fn, [chat_input, chat_history],
-        [diff_view, diff_original_state, diff_modified_state, chat_history, chat_input],
+    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],
+    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 actions
-    apply_btn.click(
-        lambda mod: mod, diff_modified_state, editor
+    apply_btn.click(lambda mod: mod, diff_modified_state, editor
     ).then(lambda: gr.Tabs(selected="terminal-tab"), None, bottom_tabs)
 
-    discard_btn.click(
-        lambda: gr.update(), None, editor
+    discard_btn.click(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,
-    )
+    demo.launch(server_name="0.0.0.0", server_port=7860,
+                theme=theme, css=css, ssr_mode=False)