Update app.py

app.py

@@ -45,9 +45,9 @@ class VirtualDisplay:
             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
-            )
+                ["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}")
@@ -60,23 +60,23 @@ class VirtualDisplay:
     def capture(self):
         if not self.is_running: return None
         tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False); tmp.close()
-        # Try multiple capture methods
-        cmds = [
+        for cmd in [
             ["scrot", "--display", self.display, tmp.name],
-            f"import -window root -display {self.display} {tmp.name}"
-        ]
-        for cmd in cmds:
+            f"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=2, env={**os.environ, "DISPLAY": self.display}, **kw)
+                r = subprocess.run(cmd, capture_output=True, timeout=3,
+                                   env={**os.environ, "DISPLAY": self.display}, **kw)
                 if r.returncode == 0 and os.path.exists(tmp.name) and os.path.getsize(tmp.name) > 100:
                     return tmp.name
             except: pass
+        try: os.unlink(tmp.name)
+        except: pass
         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()
@@ -91,7 +91,7 @@ vdisplay = VirtualDisplay()
 class GUIProcessManager:
     def __init__(self):
         self.process = None
-    
+
     def launch(self, code):
         self.stop()
         tmp = "/tmp/_axon_gui_run.py"
@@ -127,8 +127,8 @@ gui_mgr = GUIProcessManager()
 class FileSystem:
     def __init__(self):
         self.files: Dict[str, str] = {
-            "main.py": '# Start coding here\nprint("Hello World")',
-            "utils.py": "def add(a, b): return a + b",
+            "main.py": '# Start coding here\nfrom utils import add\n\nresult = add(3, 7)\nprint(f"3 + 7 = {result}")',
+            "utils.py": "def add(a, b):\n    return a + b\n\ndef multiply(a, b):\n    return a * b",
         }
         self.current_file = "main.py"
         self._sync_dir = "/tmp/axon_workspace"
@@ -152,8 +152,7 @@ class FileSystem:
         if f in self.files: self.current_file = f
     def create_file(self, f, c=""):
         if f and f not in self.files:
-            self.files[f] = c; self.current_file = f
-            self._sync_to_disk()
+            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]
@@ -179,12 +178,17 @@ class SnippetLibrary:
                 with open(SNIPPETS_FILE) as f: self.snippets = json.load(f)
             except: self.snippets = {}
     def _save(self):
-        with open(SNIPPETS_FILE, "w") as f: json.dump(self.snippets, f)
+        try:
+            with open(SNIPPETS_FILE, "w") as f: json.dump(self.snippets, f)
+        except: pass
     def add(self, name, code, tags=""):
         if not name.strip(): return "Error: name required"
         self.snippets[name.strip()] = {"code": code, "tags": tags, "created": time.strftime("%Y-%m-%d")}
         self._save(); return f"Saved: {name}"
     def get(self, name): return self.snippets.get(name, {}).get("code", "")
+    def delete(self, name):
+        if name in self.snippets: del self.snippets[name]; self._save(); return f"Deleted: {name}"
+        return "Not found"
     def get_names(self): return list(self.snippets.keys())
 
 snippets = SnippetLibrary()
@@ -195,15 +199,23 @@ snippets = SnippetLibrary()
 class CodeAnalyzer:
     @staticmethod
     def analyze(code: str) -> str:
-        if not code.strip(): return "Empty file"
+        if not code or not code.strip(): return "Empty file"
         try: tree = ast.parse(code)
-        except SyntaxError as e: return f"⚠ Syntax error: line {e.lineno}: {e.msg}"
+        except SyntaxError as e: return f"⚠ Syntax error line {e.lineno}: {e.msg}"
         lines = [f"📊 {len(code.splitlines())} lines"]
+        imports = []
         for node in ast.walk(tree):
-            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
-                lines.append(f"⚡ def {node.name}(...)")
+            if isinstance(node, ast.Import):
+                for a in node.names: imports.append(a.name)
+            elif isinstance(node, ast.ImportFrom):
+                imports.append(f"from {node.module}")
+            elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
+                args = ", ".join(a.arg for a in node.args.args)
+                lines.append(f"  L{node.lineno} ⚡ def {node.name}({args})")
             elif isinstance(node, ast.ClassDef):
-                lines.append(f"🏗 class {node.name}")
+                lines.append(f"  L{node.lineno} 🏗 class {node.name}")
+        if imports:
+            lines.insert(1, f"📦 {', '.join(imports)}")
         return "\n".join(lines)
 
 analyzer = CodeAnalyzer()
@@ -214,56 +226,190 @@ class FindReplace:
         if not query: return ""
         res = []
         for i, line in enumerate(code.splitlines(), 1):
-            if (query in line) if case else (query.lower() in line.lower()):
-                res.append(f"L{i}: {line.strip()}")
-        return "\n".join(res) if res else "No matches."
+            match = (query in line) if case else (query.lower() in line.lower())
+            if match:
+                res.append(f"  L{i}: {line.strip()}")
+        return f"Found {len(res)} match(es):\n" + "\n".join(res) if res else "No matches."
+
+    @staticmethod
+    def find_all_files(files, query):
+        if not query: return ""
+        res = []
+        for fname, content in files.items():
+            for i, line in enumerate(content.splitlines(), 1):
+                if query.lower() in line.lower():
+                    res.append(f"  {fname}:L{i}: {line.strip()}")
+        return f"Found {len(res)} across all files:\n" + "\n".join(res) if res else "No matches."
 
     @staticmethod
     def replace_in_file(code, find, replace, case):
         if not find: return code, "Empty find."
         flags = 0 if case else re.IGNORECASE
         new_code, count = re.subn(re.escape(find), replace, code, flags=flags)
-        return new_code, f"Replaced {count}."
+        return new_code, f"Replaced {count} occurrence(s)."
 
 finder = FindReplace()
 
 # ═══════════════════════════════════════
-# Terminal
+# PTY Terminal (fixed)
 # ═══════════════════════════════════════
 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 = []
+        self.lock = threading.Lock()
+        self.log_lines: List[str] = []
+        self.max_lines = 800
+        self.cmd_history: List[str] = []
         self._spawn()
+        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 tab")
+        self._append("══════════════════════════════════════")
+
     def _spawn(self):
         try:
-            pid, fd = pty.openpty(); self.pid = os.fork()
+            pid, fd = pty.openpty()
+            self.pid = os.fork()
             if self.pid == 0:
-                os.setsid(); os.dup2(fd,0); os.dup2(fd,1); os.dup2(fd,2)
-                if fd>2: os.close(fd)
-                os.execvpe("/bin/bash", ["/bin/bash","--norc","-i"], {**os.environ, "PS1":"$ ", "TERM":"dumb"})
+                # Child — become the shell
+                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
-                fcntl.fcntl(self.master_fd, fcntl.F_SETFL, os.O_NONBLOCK)
-        except: pass
-    def read(self):
+                # Parent — hold the master fd
+                os.close(fd)
+                self.master_fd = pid
+                # Set non-blocking properly
+                flags = fcntl.fcntl(self.master_fd, fcntl.F_GETFL)
+                fcntl.fcntl(self.master_fd, fcntl.F_SETFL, flags | os.O_NONBLOCK)
+                # Set terminal size
+                fcntl.ioctl(self.master_fd, termios.TIOCSWINSZ,
+                            struct.pack("HHHH", 40, 120, 0, 0))
+                time.sleep(0.3)
+                self._read_raw(0.2)  # Drain initial output
+        except Exception as e:
+            self._append(f"[Shell Error] {e}")
+            self.master_fd = None
+
+    def _read_raw(self, timeout=0.1) -> str:
+        """Read raw bytes from PTY."""
         if not self.master_fd: return ""
-        try:
-            out = b""
-            while True:
-                r,_,_ = select.select([self.master_fd],[],[],0.01)
-                if not r: break
-                out += os.read(self.master_fd, 1024)
-            if out:
-                txt = out.decode(errors="replace")
-                self.log.append(txt)
-                return "".join(self.log)[-2000:]
-        except: return "".join(self.log)[-2000:]
-        return "".join(self.log)[-2000:]
-    def write(self, cmd):
-        if self.master_fd: os.write(self.master_fd, (cmd+"\n").encode())
+        out = []; deadline = time.time() + timeout
+        while True:
+            rem = deadline - time.time()
+            if rem <= 0: break
+            try:
+                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
+                elif out: break
+            except OSError: break
+        return "".join(out)
+
+    def _clean(self, text):
+        """Strip ANSI escape codes."""
+        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)
+
+    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 history
+        if not self.cmd_history or self.cmd_history[-1] != cmd:
+            self.cmd_history.append(cmd)
+            if len(self.cmd_history) > 100: 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()
+
+            # Drain stale output
+            self._read_raw(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()
+
+            # Determine timeout based on command type
+            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)
+
+            # Collect output
+            chunks = []; start = time.time(); idle = 0
+            while time.time() - start < wait:
+                c = self._read_raw(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
+
+            # Clean and filter
+            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)
+
+            result = "\n".join(filtered).strip()
+            if result: self._append(result)
+            return self.get_log()
+
+    def get_history(self) -> List[str]:
+        return list(reversed(self.cmd_history[-30:]))
+
     def cleanup(self):
-        if self.pid: os.kill(self.pid, signal.SIGTERM)
+        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()
 
@@ -272,188 +418,467 @@ terminal = PTYTerminal()
 # ═══════════════════════════════════════
 @lru_cache(maxsize=1)
 def load_model():
-    print("Loading 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)
-        return pipeline("text-generation", model=mdl, tokenizer=tok)
-    except: return None
-
-def ai_gen(prompt, code):
+        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 ai_gen(system_prompt, code, max_tokens=256):
     pipe = load_model()
-    if not pipe: return "Model load failed."
-    formatted = f"<|system|>You are a Python coding assistant.</s><|user|>{prompt}\n{code}</s><|assistant|>"
-    try: return pipe(formatted, max_new_tokens=256, temperature=0.2)[0]["generated_text"].split("<|assistant|>")[-1]
-    except Exception as e: return str(e)
+    if not pipe: return "Error: model failed to load"
+    prompt = f"<|system|>{system_prompt}</s><|user|>\n{code}\n</s><|assistant|>"
+    try:
+        result = pipe(prompt, max_new_tokens=max_tokens, return_full_text=False)
+        return result[0]["generated_text"].strip()
+    except Exception as e:
+        return f"Error: {e}"
 
 # ═══════════════════════════════════════
-# UI Setup
+# GUI code detection
+# ═══════════════════════════════════════
+GUI_HINTS = ["pygame", "tkinter", "turtle", "pyglet", "arcade", "kivy",
+             "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(
     primary_hue="cyan", neutral_hue="gray",
     font=[gr.themes.GoogleFont("IBM Plex Mono"), "monospace"]
-).set(body_background_fill="#080808", body_text_color="#e0e0e0")
+).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",
+)
 
 css = """
+* { border-radius: 0 !important; }
 .gradio-container { max-width: 100% !important; }
-textarea { font-family: 'IBM Plex Mono', monospace !important; font-size: 12px !important; }
+
+/* Header */
+.axon-header {
+    background: linear-gradient(90deg, #080808, #0a1a1a, #080808);
+    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 */
+.term-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;
+}
+.term-input input {
+    background: #050505 !important; color: #00ff41 !important;
+    font-family: 'IBM Plex Mono', monospace !important;
+    border: 1px solid #0f32 !important; font-size: 12px !important;
+}
+.term-input input::placeholder { color: #0f83 !important; }
+.term-input input:focus { border-color: #0ff !important; box-shadow: 0 0 8px #0ff2 !important; }
+
+/* Editor */
+.code-editor textarea, .code-editor .cm-editor {
+    font-family: 'IBM Plex Mono', monospace !important; font-size: 13px !important;
+}
+
+/* Buttons */
+.tb {
+    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;
+}
+.tb:hover { border-color: #0ff !important; box-shadow: 0 0 12px #0ff2 !important; }
+
+/* Accordion */
+.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 */
+.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;
+}
+
+/* Structure + find */
+.struct-box textarea { background: #050505 !important; color: #0ff !important;
+    font-family: 'IBM Plex Mono', monospace !important; font-size: 11px !important; }
+.find-box textarea { background: #050505 !important; color: #ffab40 !important;
+    font-family: 'IBM Plex Mono', monospace !important; font-size: 11px !important; }
+
+/* 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; }
+
+/* Misc */
+::-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; }
 """
 
+# ═══════════════════════════════════════
+# UI
+# ═══════════════════════════════════════
 with gr.Blocks(title="Axon Pro", theme=theme, css=css) as demo:
-    
-    with gr.Row():
-        gr.Markdown("# ⬡ AXON PRO v4")
 
-    with gr.Row():
-        # Left Sidebar
+    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):
             with gr.Accordion("📁 EXPLORER", open=True):
-                file_list = gr.Dropdown(choices=fs.get_all_files(), value=fs.current_file, label="Files")
+                file_list = gr.Dropdown(choices=fs.get_all_files(), value=fs.current_file,
+                                         label="Files", interactive=True)
                 new_file_txt = gr.Textbox(placeholder="new_file.py", show_label=False)
                 with gr.Row():
-                    new_btn = gr.Button("NEW", size="sm")
-                    save_btn = gr.Button("SAVE", size="sm")
-                    del_btn = gr.Button("DEL", size="sm", variant="stop")
-            
+                    new_btn = gr.Button("NEW", size="sm", elem_classes="tb")
+                    save_btn = gr.Button("SAVE", size="sm", elem_classes="tb")
+                    del_btn = gr.Button("DEL", size="sm", elem_classes="tb stop-btn")
+
             with gr.Accordion("🗺 STRUCTURE", open=True):
-                structure_view = gr.Textbox(label="", lines=10, interactive=False)
-            
+                structure_view = gr.Textbox(
+                    value=analyzer.analyze(fs.get_current_file_content()),
+                    label="", lines=10, interactive=False, elem_classes="struct-box",
+                    show_label=False)
+
             with gr.Accordion("📋 SNIPPETS", open=False):
-                snip_list = gr.Dropdown(choices=snippets.get_names(), label="Saved")
-                snip_name = gr.Textbox(placeholder="Name", show_label=False)
+                snip_list = gr.Dropdown(choices=snippets.get_names(), label="Saved",
+                                         interactive=True)
+                snip_name = gr.Textbox(placeholder="Snippet name", show_label=False)
+                snip_tags = gr.Textbox(placeholder="Tags: api, util, loop", show_label=False)
                 with gr.Row():
-                    snip_save_btn = gr.Button("SAVE", size="sm")
-                    snip_load_btn = gr.Button("INSERT", size="sm")
-                status_msg = gr.Markdown("")
+                    snip_save = gr.Button("SAVE", size="sm", elem_classes="tb")
+                    snip_insert = gr.Button("INSERT", size="sm", elem_classes="tb")
+                    snip_del = gr.Button("DEL", size="sm", elem_classes="tb stop-btn")
+                snip_status = gr.Markdown("")
 
-        # Main Editor Area
+        # ══ CENTER ══
         with gr.Column(scale=4):
-            editor = gr.Code(value=fs.get_current_file_content(), language="python", lines=20)
-            
+            editor = gr.Code(value=fs.get_current_file_content(),
+                             label=f"  {fs.current_file}", language="python",
+                             lines=20, interactive=True, elem_classes="code-editor")
+
             with gr.Row():
-                run_btn = gr.Button("▶ RUN CLI", variant="primary")
-                run_gui_btn = gr.Button("🖥 RUN GUI", variant="primary")
-                stop_btn = gr.Button("■ STOP", variant="stop")
-                ai_complete = gr.Button("✨ COMPLETE")
-                ai_explain = gr.Button("🤔 EXPLAIN")
-
-            with gr.Tabs():
-                with gr.Tab("TERMINAL"):
-                    term_out = gr.Textbox(lines=12, label="Output", interactive=False)
-                    term_in = gr.Textbox(placeholder="$ Command...", show_label=False)
-                    term_run = gr.Button("Enter")
-
-                with gr.Tab("FIND"):
+                run_btn = gr.Button("▶ RUN", variant="primary", elem_classes="tb")
+                stop_btn = gr.Button("■ STOP", size="sm", elem_classes="tb stop-btn")
+                ai_complete = gr.Button("✦ COMPLETE", elem_classes="tb")
+                ai_explain = gr.Button("◈ EXPLAIN", elem_classes="tb")
+                ai_refactor = gr.Button("⟲ REFACTOR", elem_classes="tb")
+
+            with gr.Tabs() as bottom_tabs:
+
+                with gr.Tab("⌘ TERMINAL", id="term-tab"):
+                    term_out = gr.Textbox(value=terminal.get_log(), lines=12, max_lines=25,
+                                           interactive=False, elem_classes="term-box",
+                                           label="", show_label=False)
+                    with gr.Row():
+                        term_in = gr.Textbox(placeholder="$ pip install, git clone, ls ...",
+                                              show_label=False, elem_classes="term-input")
+                    with gr.Row():
+                        clear_btn = gr.Button("CLEAR", size="sm", elem_classes="tb")
+                        hist_dd = gr.Dropdown(choices=[], label="", interactive=True,
+                                               scale=3, container=False)
+                        hist_btn = gr.Button("↻ HISTORY", size="sm", elem_classes="tb")
+
+                with gr.Tab("🔍 FIND", id="find-tab"):
+                    with gr.Row():
+                        find_query = gr.Textbox(label="Find", scale=4)
+                        replace_query = gr.Textbox(label="Replace", scale=4)
+                        find_case = gr.Checkbox(label="Case", value=False)
                     with gr.Row():
-                        find_query = gr.Textbox(label="Find")
-                        replace_query = gr.Textbox(label="Replace")
-                        find_case = gr.Checkbox(label="Case Sensitive")
+                        find_btn = gr.Button("FIND", size="sm", elem_classes="tb")
+                        find_all_btn = gr.Button("FIND ALL FILES", size="sm", elem_classes="tb")
+                        replace_btn = gr.Button("REPLACE ALL", size="sm", elem_classes="tb")
+                    find_results = gr.Textbox(value="", lines=6, interactive=False,
+                                               elem_classes="find-box", label="", show_label=False)
+
+                with gr.Tab("💬 AI CHAT", id="chat-tab"):
+                    chat_history = gr.Chatbot(label="", height=250)
                     with gr.Row():
-                        find_btn = gr.Button("Find")
-                        replace_btn = gr.Button("Replace All")
-                    find_results = gr.Textbox(label="Results", lines=5)
+                        chat_input = gr.Textbox(placeholder="Describe code you want...",
+                                                 scale=8, container=False)
+                        send_btn = gr.Button("GEN", variant="primary", scale=1, elem_classes="tb")
 
-                with gr.Tab("DISPLAY"):
-                    display_image = gr.Image(label="Virtual Display", interactive=False, type="filepath")
+                with gr.Tab("⇄ DIFF", id="diff-tab"):
+                    diff_view = gr.Code(label="AI Changes", language="python",
+                                         interactive=False, lines=15)
                     with gr.Row():
-                        capture_btn = gr.Button("📸 Refresh Screen")
-                        auto_capture = gr.Checkbox(label="Auto Refresh (2s)")
-                    auto_timer = gr.Timer(value=2.0, active=False)
+                        apply_btn = gr.Button("✓ APPLY", variant="primary", elem_classes="tb")
+                        discard_btn = gr.Button("✗ DISCARD", elem_classes="tb")
+
+        # ══ 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", size="sm", elem_classes="tb")
+                    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}</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 = gr.State("")
+    diff_modified = gr.State("")
+    auto_timer = gr.Timer(2, active=False)
 
     # ═══════════════════════════════════════
-    # LOGIC & WIRING
+    # HANDLERS
     # ═══════════════════════════════════════
 
-    # --- File System Handlers ---
-    def load_file(f):
-        fs.set_current_file(f)
-        code = fs.get_current_file_content()
-        return code, analyzer.analyze(code)
+    # --- File ops ---
+    def on_editor_change(content):
+        fs.save_file(content)
+        return analyzer.analyze(content)
 
-    def create_file(name):
+    def on_file_select(filename):
+        fs.set_current_file(filename)
+        c = fs.get_current_file_content()
+        return c, gr.update(label=f"  {filename}"), analyzer.analyze(c)
+
+    def on_new_file(name):
+        if not name or not name.strip():
+            return gr.update(), gr.update(), gr.update(), gr.update()
+        name = name.strip()
         fs.create_file(name)
-        return gr.Dropdown(choices=fs.get_all_files(), value=name), "", ""
+        c = fs.get_current_file_content()
+        return (gr.update(choices=fs.get_all_files(), value=name),
+                c, gr.update(label=f"  {name}"), analyzer.analyze(c))
 
-    def save_file(code):
+    def on_save(code):
         fs.save_file(code)
-        return "Saved."
-
-    def delete_file(f):
-        fs.delete_file(f)
-        new_f = fs.current_file
-        new_c = fs.get_current_file_content()
-        return gr.Dropdown(choices=fs.get_all_files(), value=new_f), new_c
-
-    file_list.change(load_file, file_list, [editor, structure_view])
-    new_btn.click(create_file, new_file_txt, [file_list, new_file_txt, editor])
-    save_btn.click(save_file, editor, status_msg)
-    del_btn.click(delete_file, file_list, [file_list, editor])
-    editor.change(lambda c: analyzer.analyze(c), editor, structure_view)
-
-    # --- Execution Handlers ---
-    def run_cli(code):
+        return (f"**AXON PRO v4.0** │ Python {sys.version.split()[0]} │ CPU │ "
+                f"TinyLlama-1.1B │ ✓ Saved {fs.current_file}")
+
+    def on_delete():
+        name = fs.current_file
+        if fs.delete_file(name):
+            c = fs.get_current_file_content()
+            return (gr.update(choices=fs.get_all_files(), value=fs.current_file),
+                    c, gr.update(label=f"  {fs.current_file}"), analyzer.analyze(c))
+        return gr.update(), gr.update(), gr.update(), gr.update()
+
+    # --- Run (smart: auto-detects GUI vs CLI) ---
+    def on_run(code):
         fs.save_file(code)
-        cmd = f"python3 {fs._sync_dir}/{fs.current_file}"
-        terminal.write(cmd)
-        time.sleep(0.5)
-        return terminal.read()
 
-    def run_gui(code):
+        if is_gui_code(code):
+            # GUI app — launch as background process, capture display
+            msg = gui_mgr.launch(code)
+            terminal._append(msg)
+            time.sleep(1.5)
+            ss = vdisplay.capture()
+            return terminal.get_log(), ss, f"<small>Xvfb: ● ON | {gui_mgr.get_status()}</small>"
+
+        # Normal script — run with subprocess, capture output
+        terminal._append("$ python [editor]")
+        tmp = "/tmp/_axon_run.py"
+        with open(tmp, "w") as f: f.write(code)
+        try:
+            env = os.environ.copy()
+            env["DISPLAY"] = DISPLAY_NUM
+            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(), gr.update(), gr.update()
+
+    def on_stop():
+        msg = gui_mgr.stop(); terminal._append(msg)
+        return terminal.get_log(), f"<small>Xvfb: {'● ON' if vdisplay.is_running else '○ OFF'}</small>"
+
+    # --- Terminal ---
+    def on_term_cmd(cmd):
+        return terminal.run_command(cmd), ""
+
+    def on_clear():
+        terminal.log_lines = []; return ""
+
+    def on_refresh_hist():
+        return gr.update(choices=terminal.get_history())
+
+    def on_select_hist(cmd):
+        return cmd if cmd else gr.update()
+
+    # --- AI ---
+    def on_complete(code):
+        result = ai_gen("Complete this Python code. Only output the completion, no explanations.", code)
+        new_code = code + "\n" + result
+        return new_code, analyzer.analyze(new_code)
+
+    def on_explain(code):
+        explanation = ai_gen("Explain this Python code concisely.", code, 512)
+        terminal._append(f"[AI Explanation]\n{explanation}")
+        return terminal.get_log()
+
+    def on_refactor(code):
+        refactored = ai_gen("Refactor this Python code for PEP 8 and best practices. Output only code.", code, 512)
         fs.save_file(code)
-        return gui_mgr.launch(code)
-
-    run_btn.click(run_cli, editor, term_out)
-    run_gui_btn.click(run_gui, editor, term_out)
-    stop_btn.click(gui_mgr.stop, None, term_out)
-
-    # --- AI Handlers ---
-    ai_complete.click(lambda c: ai_gen("Complete this code:", c), editor, editor)
-    ai_explain.click(lambda c: ai_gen("Explain this code:", c), editor, term_out)
-
-    # --- Terminal Handlers ---
-    def term_cmd(cmd):
-        terminal.write(cmd)
-        time.sleep(0.2)
-        return "", terminal.read()
-    
-    term_run.click(term_cmd, term_in, [term_in, term_out])
-    term_in.submit(term_cmd, term_in, [term_in, term_out])
-
-    # --- Find/Replace Handlers ---
-    find_btn.click(finder.find_in_file, [editor, find_query, find_case], find_results)
-    
-    def replace_h(c, f, r, case):
-        new_c, msg = finder.replace_in_file(c, f, r, case)
-        return new_c, msg
-    
-    replace_btn.click(replace_h, [editor, find_query, replace_query, find_case], [editor, find_results])
-
-    # --- Snippet Handlers ---
-    snip_save_btn.click(snippets.add, [snip_name, editor], status_msg)
-    
-    def insert_snip(name, code):
-        return code + "\n" + snippets.get(name)
-        
-    snip_load_btn.click(insert_snip, [snip_list, editor], editor)
-
-    # --- Display Handlers ---
-    def capture():
-        path = vdisplay.capture()
-        return path if path else None
-
-    capture_btn.click(capture, None, display_image)
-    auto_capture.change(lambda x: gr.Timer(active=x), auto_capture, auto_timer)
-    auto_timer.tick(capture, None, display_image)
+        import difflib
+        diff = "\n".join(difflib.unified_diff(code.splitlines(), refactored.splitlines(), lineterm=""))
+        return diff, code, refactored
+
+    def on_generate(prompt_text, history):
+        generated = ai_gen(f"Write Python code for: {prompt_text}", "", 512)
+        import difflib
+        diff = "\n".join(difflib.unified_diff([], generated.splitlines(), lineterm=""))
+        new_h = history + [
+            {"role": "user", "content": prompt_text},
+            {"role": "assistant", "content": "Code generated → check DIFF tab"},
+        ]
+        return diff, "", generated, new_h, ""
+
+    def on_apply(modified):
+        return modified, analyzer.analyze(modified)
+
+    # --- Snippets ---
+    def on_snip_save(name, tags, code):
+        msg = snippets.add(name, code, tags)
+        return msg, gr.update(choices=snippets.get_names())
+
+    def on_snip_insert(name, code):
+        snip_code = snippets.get(name)
+        if snip_code:
+            # Append snippet at cursor position (end of file)
+            return code + "\n\n" + snip_code if code.strip() else snip_code
+        return code
+
+    def on_snip_del(name):
+        msg = snippets.delete(name)
+        return msg, gr.update(choices=snippets.get_names())
+
+    # --- Find ---
+    def on_find(code, query, case):
+        return finder.find_in_file(code, query, case)
+
+    def on_find_all(query):
+        return finder.find_all_files(fs.files, query)
+
+    def on_replace(code, fq, rq, case):
+        new_code, msg = finder.replace_in_file(code, fq, rq, case)
+        fs.save_file(new_code)
+        return new_code, msg, analyzer.analyze(new_code)
+
+    # --- Display ---
+    def on_capture():
+        return vdisplay.capture()
+
+    def on_auto_tick():
+        if gui_mgr.process and gui_mgr.process.poll() is None:
+            return vdisplay.capture()
+        return gr.update()
+
+    def on_auto_toggle(checked):
+        return gr.Timer(2, active=checked)
+
+    # ═══════════════════════════════════════
+    # WIRING
+    # ═══════════════════════════════════════
+
+    # Files
+    editor.change(on_editor_change, editor, structure_view)
+    file_list.change(on_file_select, file_list, [editor, editor, structure_view])
+    new_btn.click(on_new_file, new_file_txt, [file_list, editor, editor, structure_view])
+    save_btn.click(on_save, editor, status_bar)
+    del_btn.click(on_delete, None, [file_list, editor, editor, structure_view])
+
+    # Run
+    run_btn.click(on_run, editor, [term_out, display_image, gui_status])
+    stop_btn.click(on_stop, None, [term_out, gui_status])
+
+    # Terminal
+    term_in.submit(on_term_cmd, term_in, [term_out, term_in])
+    clear_btn.click(on_clear, None, term_out)
+    hist_btn.click(on_refresh_hist, None, hist_dd)
+    hist_dd.change(on_select_hist, hist_dd, term_in)
+
+    # AI
+    ai_complete.click(on_complete, editor, [editor, structure_view])
+    ai_explain.click(on_explain, editor, term_out)
+    ai_refactor.click(on_refactor, editor, [diff_view, diff_original, diff_modified]
+    ).then(lambda: gr.Tabs(selected="diff-tab"), None, bottom_tabs)
+
+    # Chat
+    chat_input.submit(on_generate, [chat_input, chat_history],
+        [diff_view, diff_original, diff_modified, chat_history, chat_input]
+    ).then(lambda: gr.Tabs(selected="diff-tab"), None, bottom_tabs)
+    send_btn.click(on_generate, [chat_input, chat_history],
+        [diff_view, diff_original, diff_modified, chat_history, chat_input]
+    ).then(lambda: gr.Tabs(selected="diff-tab"), None, bottom_tabs)
+
+    # Diff
+    apply_btn.click(on_apply, diff_modified, [editor, structure_view]
+    ).then(lambda: gr.Tabs(selected="term-tab"), None, bottom_tabs)
+    discard_btn.click(lambda: (gr.update(), gr.update()), None, [editor, structure_view]
+    ).then(lambda: gr.Tabs(selected="term-tab"), None, bottom_tabs)
+
+    # Snippets
+    snip_save.click(on_snip_save, [snip_name, snip_tags, editor], [snip_status, snip_list])
+    snip_insert.click(on_snip_insert, [snip_list, editor], editor)
+    snip_del.click(on_snip_del, snip_list, [snip_status, snip_list])
+
+    # Find
+    find_btn.click(on_find, [editor, find_query, find_case], find_results)
+    find_all_btn.click(on_find_all, find_query, find_results)
+    replace_btn.click(on_replace, [editor, find_query, replace_query, find_case],
+                      [editor, find_results, structure_view])
+
+    # Display
+    capture_btn.click(on_capture, None, display_image)
+    auto_capture.change(on_auto_toggle, auto_capture, auto_timer)
+    auto_timer.tick(on_auto_tick, None, display_image)
+
 
 if __name__ == "__main__":
     import atexit
     atexit.register(terminal.cleanup)
     atexit.register(gui_mgr.stop)
     atexit.register(vdisplay.cleanup)
-    
-    demo.queue().launch(
-        server_name="0.0.0.0", 
-        server_port=7860, 
-        ssr_mode=False,
-        allowed_paths=["/tmp"]
-    )
+    demo.queue().launch(server_name="0.0.0.0", server_port=7860,
+                        ssr_mode=False, allowed_paths=["/tmp"])