import os
import time
import json
import shutil
import asyncio
import zipfile
import re
import subprocess
from pathlib import Path
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, UploadFile, File, BackgroundTasks, Form
from fastapi.responses import HTMLResponse, FileResponse, JSONResponse
from fastapi.middleware.cors import CORSMiddleware
import httpx
import uvicorn

# ==========================================
# KONFIGURASI & VARIABEL GLOBAL
# ==========================================
PORT = 7860
BASE_DIR = Path("/tmp/puruai_sessions")
SESSION_TIMEOUT = 259200 # 3 Hari
API_URL = 'https://puruboy-api.vercel.app/api/ai/gemini-v2'
BASE_DELAY = 3 

# Prompt Utama (Pelaksana) dengan kemampuan Penjelasan (Chain of Thought)
SYSTEM_PROMPT = """YOU ARE AN ELITE AI CODE EDITOR AGENT DESIGNED TO OPERATE WITH MAXIMUM PRECISION IN A CONTROLLED FILE SYSTEM ENVIRONMENT. YOU MUST EXECUTE EXACTLY ONE ACTION PER LOOP AND STRICTLY FOLLOW THE PROVIDED EXECUTION COMMAND FORMAT.

YOUR CORE OBJECTIVE IS TO ANALYZE, DECIDE, EXPLAIN YOUR LOGIC, AND EXECUTE FILE OPERATIONS STEP-BY-STEP.

---

### 🔥 CORE RULES

- YOU MUST MAKE **ONLY ONE DECISION PER LOOP**
- YOU MUST **EXPLAIN YOUR THINKING PROCESS FIRST**, THEN OUTPUT EXACTLY ONE <execution> COMMAND
- YOU MUST OPERATE BASED ONLY ON AVAILABLE CONTEXT (NO INTERNET, NO ASSUMPTIONS)
- YOU MUST FOLLOW STRICT PATH FORMAT: `#root/...`

---

### 📂 AVAILABLE ACTIONS

1. READ ALL FILE STRUCTURE
<execution>all</execution>

2. READ FILE CONTENT
<execution>read <path>#root/filename.ext</path></execution>

3. WRITE / OVERWRITE FILE
<execution>write <path>#root/filename.ext</path><content>FILE_CONTENT_HERE</content></execution>

4. MOVE FILE
<execution>move <path>#root/source.ext</path><to>#root/destination.ext</to></execution>

5. DELETE FILE
<execution>delete <path>#root/filename.ext</path></execution>

6. SEARCH (TRIMMED CONTEXT)
<execution><trim>SEARCH_QUERY</trim></execution>

7. FINISH / STOP EXECUTION (WHEN TASK IS COMPLETE)
<execution>stop</execution>

---

### 🧠 CHAIN OF THOUGHT PROCESS (MANDATORY)

YOU MUST FOLLOW THIS THINKING SEQUENCE AND **WRITE IT DOWN** BEFORE EVERY ACTION:

1. UNDERSTAND: Identify the user's goal clearly.
2. ANALYZE: Check what information is missing based on current context.
3. DECIDE: Decide the safest and most logical next step (read, write, search, etc).
4. EXECUTE: Output exactly one valid <execution> command.

---

### ⚠️ WHAT NOT TO DO (CRITICAL)

- NEVER OUTPUT MULTIPLE COMMANDS  
❌ WRONG:
<execution>all</execution> <execution>read ...</execution>

- NEVER PUT YOUR EXPLANATION INSIDE THE EXECUTION TAG  
❌ WRONG:
<execution>I will read the file now. read <path>#root/file.js</path></execution>
✅ CORRECT:
I need to check the contents of file.js to find the bug.
<execution>read <path>#root/file.js</path></execution>

- NEVER GUESS FILE CONTENT  
❌ WRONG:
<execution>write ...> (without reading first)

- NEVER SKIP STEPS (LIKE WRITING BEFORE READING STRUCTURE)

- NEVER USE INVALID PATH FORMAT  
❌ WRONG: `/home/user/file.js`  
✅ CORRECT: `#root/file.js`

- NEVER PERFORM MORE THAN ONE ACTION PER LOOP

---

### ✅ OPTIMAL STRATEGY

- ALWAYS START WITH EXPLANATION + <execution>all</execution>
- THEN: READ relevant files → ANALYZE → MODIFY step-by-step
- WHEN FINISHED: EXPLAIN completion → <execution>stop</execution>

---

### 🎯 FEW-SHOT EXAMPLES

#### Example 1: Initial Step
User: "Fix bug in app.js"

✅ Correct AI Response:
I need to understand the current project structure first before locating and modifying `app.js`. Let me read the directory structure.
<execution>all</execution>

---

#### Example 2: After Seeing Structure
✅ Correct AI Response:
I see `app.js` in the root directory. I need to read its content to understand the bug before making any changes.
<execution>read <path>#root/app.js</path></execution>

---

#### Example 3: After Analysis
✅ Correct AI Response:
I found the issue. The variable `userData` is undefined. I will fix the function and write the updated code back to `app.js`.
<execution>write <path>#root/app.js</path><content>FIXED_CODE</content></execution>

---

#### Example 4: Task Completed
✅ Correct AI Response:
The bug in `app.js` has been successfully fixed and the changes are written. My task is complete.
<execution>stop</execution>

---

### 🚀 FINAL DIRECTIVE

YOU ARE A DETERMINISTIC FILE-SYSTEM AGENT.  
ALWAYS EXPLAIN YOUR LOGIC BRIEFLY FIRST.
THEN EXECUTE EXACTLY ONE ACTION.  
NEVER BREAK THE XML FORMAT FOR COMMANDS."""

SESSIONS = {}

app = FastAPI()
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])

# ==========================================
# UTILITAS FILE SYSTEM SERVER
# ==========================================
def get_session_dir(session_id: str) -> Path:
    return BASE_DIR / session_id

def update_session_activity(session_id: str):
    if session_id in SESSIONS:
        SESSIONS[session_id]["last_active"] = time.time()
    session_dir = get_session_dir(session_id)
    if session_dir.exists():
        os.utime(session_dir, None)

def get_vfs_list(session_id: str) -> list:
    session_dir = get_session_dir(session_id)
    if not session_dir.exists(): return []
    files = []
    for root, _, filenames in os.walk(session_dir):
        for filename in filenames:
            full_path = Path(root) / filename
            rel_path = full_path.relative_to(session_dir).as_posix()
            files.append(rel_path)
    return sorted(files)

def is_binary(file_path: Path) -> bool:
    try:
        with open(file_path, 'tr') as check_file:
            check_file.read(1024)
            return False
    except UnicodeDecodeError:
        return True

async def broadcast_ws(session_id: str, message: dict):
    if session_id in SESSIONS and SESSIONS[session_id].get("ws"):
        try:
            await SESSIONS[session_id]["ws"].send_json(message)
        except: pass 

async def send_vfs_update(session_id: str):
    await broadcast_ws(session_id, {"type": "vfs", "data": get_vfs_list(session_id)})

# ==========================================
# LOGIKA AGENT AI PELAKSANA
# ==========================================
def execute_command(session_id: str, command_str: str) -> dict:
    session_dir = get_session_dir(session_id)
    os.makedirs(session_dir, exist_ok=True)
    
    try:
        exec_match = re.search(r'<execution>\s*([\s\S]*?)\s*</execution>', command_str, re.IGNORECASE)
        if not exec_match:
            return {"action": "stop", "log": "Operasi selesai atau tag <execution> tidak ditemukan. Berhenti."}
        
        cmd_body = exec_match.group(1).strip()
        
        # 0. STOP
        if cmd_body.lower() == 'stop':
            return {"action": "stop", "log": "Task Completed. Agent dihentikan secara eksplisit oleh AI."}
        
        # 1. ALL
        if cmd_body == 'all':
            files = get_vfs_list(session_id)
            return {"action": "all", "log": f"VFS Structure:\n" + "\n".join(files) if files else 'Directory is empty.'}
            
        # 2. READ
        read_match = re.search(r'^read\s*<path>(.*?)</path>', cmd_body, re.IGNORECASE)
        if read_match:
            clean_path = read_match.group(1).replace('#root/', '').strip()
            file_path = session_dir / clean_path
            if file_path.exists() and file_path.is_file():
                if is_binary(file_path): return {"action": "read", "log": "[Binary File Unreadable]"}
                with open(file_path, 'r', encoding='utf-8') as f: content = f.read()
                return {"action": "read", "log": f"Content of {clean_path}:\n{content}"}
            raise Exception("File not found.")
                
        # 3. WRITE
        write_match = re.search(r'^write\s*<path>(.*?)</path>\s*<content>([\s\S]*?)</content>', cmd_body, re.IGNORECASE)
        if write_match:
            clean_path = write_match.group(1).replace('#root/', '').strip()
            content = write_match.group(2)
            file_path = session_dir / clean_path
            os.makedirs(file_path.parent, exist_ok=True)
            with open(file_path, 'w', encoding='utf-8') as f: f.write(content)
            return {"action": "write", "log": f"Successfully written to {clean_path}."}
            
        # 4. DELETE
        delete_match = re.search(r'^delete\s*<path>(.*?)</path>', cmd_body, re.IGNORECASE)
        if delete_match:
            clean_path = delete_match.group(1).replace('#root/', '').strip()
            file_path = session_dir / clean_path
            if file_path.exists():
                if file_path.is_dir(): shutil.rmtree(file_path)
                else: os.remove(file_path)
                return {"action": "delete", "log": f"Deleted: {clean_path}"}
            raise Exception("File not found.")

        # 5. MOVE
        move_match = re.search(r'^move\s*<path>(.*?)</path>\s*<to>(.*?)</to>', cmd_body, re.IGNORECASE)
        if move_match:
            src_path = move_match.group(1).replace('#root/', '').strip()
            dst_path = move_match.group(2).replace('#root/', '').strip()
            src_full = session_dir / src_path
            dst_full = session_dir / dst_path
            if src_full.exists():
                os.makedirs(dst_full.parent, exist_ok=True)
                shutil.move(str(src_full), str(dst_full))
                return {"action": "move", "log": f"Moved {src_path} to {dst_path}"}
            raise Exception("Source file not found.")

        # 6. SEARCH / TRIM
        trim_match = re.search(r'^<trim>(.*?)</trim>', cmd_body, re.IGNORECASE)
        if trim_match:
            query = trim_match.group(1).strip()
            results = []
            for root, _, files in os.walk(session_dir):
                for file in files:
                    fp = Path(root) / file
                    if is_binary(fp): continue
                    try:
                        with open(fp, 'r', encoding='utf-8') as f:
                            lines = f.readlines()
                            for i, line in enumerate(lines):
                                if query.lower() in line.lower():
                                    rel_path = fp.relative_to(session_dir).as_posix()
                                    results.append(f"{rel_path}:{i+1}: {line.strip()[:120]}")
                    except: pass
            res_str = "\n".join(results) if results else "No matches found."
            return {"action": "search", "log": f"Search Results for '{query}':\n{res_str}"}

        raise Exception("Command structure unrecognized or missing required tags.")
    except Exception as e:
        return {"action": "error", "log": f"ERROR: {str(e)}"}

async def fetch_ai(prompt_text: str) -> str:
    async with httpx.AsyncClient(timeout=90.0) as client:
        response = await client.post(API_URL, json={"prompt": prompt_text})
        response.raise_for_status()
        return response.json()['result']['answer']

async def agent_loop(session_id: str, initial_prompt: str):
    session = SESSIONS[session_id]
    session["is_looping"] = True
    session["chat_history"].append({"role": "user", "text": initial_prompt})
    session["ai_memory"].append({"role": "user", "text": initial_prompt})
    await broadcast_ws(session_id, {"type": "chat_update", "data": session["chat_history"]})

    while session["is_looping"]:
        try:
            update_session_activity(session_id)
            vfs_list = get_vfs_list(session_id)
            
            base_payload = f"{SYSTEM_PROMPT}\n\n[Context VFS: {', '.join(vfs_list) if vfs_list else 'Empty'}]\n\n"
            history_log = "\n".join([f"{m['role'].upper()}: {m['text']}" for m in session["ai_memory"][-12:]]) 
            
            ai_response = await fetch_ai(f"{base_payload}History:\n{history_log}\n\nExecute NEXT STEP:")
            
            session["chat_history"].append({"role": "ai", "text": ai_response})
            session["ai_memory"].append({"role": "ai", "text": ai_response})
            await broadcast_ws(session_id, {"type": "chat_update", "data": session["chat_history"]})
            
            exec_result = execute_command(session_id, ai_response)
            if exec_result:
                await send_vfs_update(session_id)
                if exec_result["action"] == "stop":
                    session["is_looping"] = False
                    session["ai_memory"] = []
                    
                    # Tambahkan log terakhir sebelum berhenti ke UI
                    sys_msg = f"SystemLog (STOP):\n{exec_result['log']}"
                    session["chat_history"].append({"role": "system", "text": sys_msg})
                    await broadcast_ws(session_id, {"type": "chat_update", "data": session["chat_history"]})
                    
                    await broadcast_ws(session_id, {"type": "status", "text": "Selesai", "statusType": "done", "isLooping": False})
                    break
                
                sys_msg = f"SystemLog ({exec_result['action']}):\n{exec_result['log']}"

                session["chat_history"].append({"role": "system", "text": sys_msg})
                session["ai_memory"].append({"role": "system", "text": sys_msg})
                await broadcast_ws(session_id, {"type": "chat_update", "data": session["chat_history"]})
            
            await broadcast_ws(session_id, {"type": "status", "text": f"Jeda {BASE_DELAY}s...", "statusType": "idle"})
            await asyncio.sleep(BASE_DELAY)
            
        except Exception as e:
            session["is_looping"] = False
            session["chat_history"].append({"role": "system", "text": f"SYSTEM ERROR: {str(e)}"})
            await broadcast_ws(session_id, {"type": "chat_update", "data": session["chat_history"]})
            await broadcast_ws(session_id, {"type": "status", "text": "Error", "statusType": "error", "isLooping": False})
            break

# ==========================================
# ENDPOINTS REST & WEBSOCKET
# ==========================================
@app.get("/")
async def get_ui(): return HTMLResponse(HTML_CONTENT)

@app.get("/api/check_session/{session_id}")
async def check_session(session_id: str):
    return {"valid": session_id in SESSIONS or get_session_dir(session_id).exists()}

@app.websocket("/ws/{session_id}")
async def websocket_endpoint(websocket: WebSocket, session_id: str):
    await websocket.accept()
    if session_id not in SESSIONS:
        SESSIONS[session_id] = {"chat_history": [], "ai_memory": [], "is_looping": False, "last_active": time.time(), "ws": websocket}
    else:
        SESSIONS[session_id]["ws"] = websocket
        
    session = SESSIONS[session_id]
    await websocket.send_json({"type": "chat_update", "data": session["chat_history"]})
    await websocket.send_json({"type": "vfs", "data": get_vfs_list(session_id)})
    
    try:
        while True:
            data = await websocket.receive_text()
            payload = json.loads(data)
            if payload["action"] == "prompt" and not session["is_looping"]:
                asyncio.create_task(agent_loop(session_id, payload["text"]))
    except WebSocketDisconnect:
        if session_id in SESSIONS:
            SESSIONS[session_id]["ws"] = None

@app.post("/api/upload_zip")
async def upload_zip(session_id: str = Form(...), file: UploadFile = File(...)):
    s_dir = get_session_dir(session_id)
    if s_dir.exists(): shutil.rmtree(s_dir)
    os.makedirs(s_dir, exist_ok=True)
    temp_zip = s_dir / "temp.zip"
    with open(temp_zip, "wb") as b: shutil.copyfileobj(file.file, b)
    with zipfile.ZipFile(temp_zip, 'r') as z: z.extractall(s_dir)
    os.remove(temp_zip)
    await send_vfs_update(session_id)
    return {"status": "success"}

@app.get("/api/download_zip/{session_id}")
async def download_zip(session_id: str):
    s_dir = get_session_dir(session_id)
    z_name = f"/tmp/PuruAI_{session_id}.zip"
    with zipfile.ZipFile(z_name, 'w', zipfile.ZIP_DEFLATED) as z:
        for root, _, files in os.walk(s_dir):
            if "_context_upload" in root: continue
            for f in files: z.write(os.path.join(root, f), os.path.relpath(os.path.join(root, f), s_dir))
    return FileResponse(z_name, filename="Project.zip")

@app.post("/api/clear_session")
async def clear_session(session_id: str = Form(...)):
    if session_id in SESSIONS:
        SESSIONS[session_id]["chat_history"] = []
        SESSIONS[session_id]["ai_memory"] = []
        SESSIONS[session_id]["is_looping"] = False
    return {"status": "success"}

@app.post("/api/delete_session")
async def delete_session(session_id: str = Form(...)):
    if session_id in SESSIONS: del SESSIONS[session_id]
    s_dir = get_session_dir(session_id)
    if s_dir.exists(): shutil.rmtree(s_dir)
    return {"status": "success"}

# ==========================================
# UI FRONTEND (HTML)
# ==========================================
HTML_CONTENT = """
<!DOCTYPE html>
<html lang="id">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0, maximum-scale=1.0, user-scalable=no">
    <title>PuruAI - Autonomous Agent</title>
    <script src="https://cdn.tailwindcss.com"></script>
    <script defer src="https://cdn.jsdelivr.net/npm/alpinejs@3.x.x/dist/cdn.min.js"></script>
    <style>
        ::-webkit-scrollbar { width: 4px; }
        ::-webkit-scrollbar-thumb { background: #374151; border-radius: 4px; }
        .blink { animation: blinker 1s linear infinite; }
        @keyframes blinker { 50% { opacity: 0.3; } }
        [x-cloak] { display: none !important; }
        /* Styling untuk xml syntax highlight */
        .xml-tag { color: #f472b6; font-weight: 600; }
    </style>
</head>
<body class="bg-gray-900 text-gray-100 font-sans h-[100dvh] flex flex-col overflow-hidden" x-data="puruApp()">
    <header class="bg-gray-800 border-b border-gray-700 p-3 flex justify-between items-center z-20">
        <div class="flex items-center gap-2">
            <div class="w-7 h-7 bg-blue-600 rounded flex items-center justify-center font-bold text-sm">P</div>
            <h1 class="text-lg font-bold">PuruAI</h1>
        </div>
        <div class="flex items-center gap-2">
            <div class="text-[10px] sm:text-xs font-medium bg-gray-900 px-3 py-1.5 rounded-full flex items-center gap-2">
                <span class="w-2 h-2 rounded-full" :class="statusColor"></span>
                <span x-text="statusText"></span>
            </div>
            <button x-show="screen === 'workspace'" @click="isMenuOpen = !isMenuOpen" class="p-1.5 bg-gray-700 rounded-md">
                <svg class="w-5 h-5" fill="none" stroke="currentColor" viewBox="0 0 24 24"><path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M4 6h16M4 12h16M4 18h16"></path></svg>
            </button>
            <div x-show="isMenuOpen" @click.outside="isMenuOpen = false" x-cloak class="absolute right-3 top-14 w-48 bg-gray-800 border border-gray-600 rounded-xl shadow-2xl z-50 divide-y divide-gray-700">
                <button @click="clearSession(); isMenuOpen = false;" class="w-full text-left px-4 py-3 text-sm hover:bg-gray-700">Clear Percakapan</button>
                <button @click="resetAll(); isMenuOpen = false;" class="w-full text-left px-4 py-3 text-sm text-red-400 hover:bg-gray-700">Reset All</button>
            </div>
        </div>
    </header>

    <div x-show="screen === 'setup'" class="flex-1 flex flex-col items-center justify-center p-6">
        <div class="w-full max-w-sm bg-gray-800 p-6 rounded-2xl border border-gray-700">
            <h2 class="text-xl font-bold mb-6 text-center">Mulai Project</h2>
            <button @click="startFreshProject()" class="w-full mb-4 bg-blue-600 py-3 rounded-xl font-semibold hover:bg-blue-500 transition">Buat Project Baru</button>
            <label class="w-full cursor-pointer bg-gray-700 hover:bg-gray-600 transition border border-dashed border-gray-500 py-4 rounded-xl flex flex-col items-center">
                <span class="text-sm">Upload File ZIP</span>
                <input type="file" class="hidden" @change="handleZipUpload($event)">
            </label>
        </div>
    </div>

    <main x-show="screen === 'workspace'" class="flex-1 relative flex flex-col overflow-hidden bg-[#0d1117]">
        <div x-show="activeTab === 'chat'" class="flex-1 flex flex-col overflow-hidden relative">
            <div id="chatContainer" class="flex-1 overflow-y-auto p-4 space-y-4 pb-24">
                <template x-for="(msg, index) in chatHistory" :key="index">
                    <div class="p-3 rounded-xl max-w-[90%] border text-sm transition-all"
                         x-data="{ expanded: false, get isLong() { return msg.text && msg.text.length > 300; } }"
                         :class="{'bg-blue-900/30 border-blue-800 ml-auto': msg.role === 'user', 'bg-gray-800/80 border-gray-600 font-mono text-xs text-gray-300': msg.role === 'system', 'bg-gray-800 border-gray-700 mr-auto': msg.role === 'ai'}">
                        
                        <!-- Role Badge -->
                        <div class="font-bold text-[10px] mb-2 uppercase" 
                             :class="{'text-blue-400': msg.role === 'user', 'text-yellow-500': msg.role === 'system', 'text-teal-400': msg.role === 'ai'}" 
                             x-text="msg.role"></div>
                        
                        <!-- Format Text Content -->
                        <div class="whitespace-pre-wrap break-words" 
                             x-html="formatXML(expanded ? msg.text : (isLong ? msg.text.substring(0, 300) + '...' : msg.text))">
                        </div>

                        <!-- Tombol Baca Selengkapnya -->
                        <template x-if="isLong">
                            <button @click="expanded = !expanded" 
                                    class="text-blue-400 text-[10px] mt-2 font-bold uppercase hover:text-blue-300 flex items-center gap-1">
                                <span x-text="expanded ? 'Sembunyikan' : 'Baca Selengkapnya'"></span>
                            </button>
                        </template>
                    </div>
                </template>
            </div>
            
            <div class="absolute bottom-0 left-0 w-full bg-gray-900 border-t border-gray-800 p-2 flex gap-2">
                <textarea x-model="userInput" :disabled="isLooping" placeholder="Tulis instruksi..." class="flex-1 bg-gray-800 rounded-xl px-3 py-2 text-sm focus:outline-none h-12 resize-none"></textarea>
                <button @click="sendPrompt()" :disabled="!userInput.trim() || isLooping" class="bg-blue-600 px-4 rounded-xl disabled:bg-gray-700 hover:bg-blue-500 transition">Kirim</button>
            </div>
        </div>
        
        <div x-show="activeTab === 'files'" class="flex-1 overflow-y-auto p-4">
            <div class="flex justify-between items-center mb-4">
                <span class="text-xs font-bold text-gray-500">PROJECT FILES</span>
                <a :href="'/api/download_zip/' + sessionId" class="text-xs bg-blue-600 px-2 py-1 rounded hover:bg-blue-500">Download ZIP</a>
            </div>
            <template x-for="file in filesList">
                <div class="py-2 border-b border-gray-800 text-sm font-mono text-gray-400" x-text="file"></div>
            </template>
        </div>
    </main>

    <nav x-show="screen === 'workspace'" class="bg-gray-800 border-t border-gray-700 flex h-14 shrink-0">
        <button @click="activeTab = 'chat'" class="flex-1 flex flex-col items-center justify-center" :class="activeTab === 'chat' ? 'text-blue-400' : 'text-gray-500'">
            <span class="text-[10px] font-bold">CHAT</span>
        </button>
        <button @click="activeTab = 'files'" class="flex-1 flex flex-col items-center justify-center" :class="activeTab === 'files' ? 'text-blue-400' : 'text-gray-500'">
            <span class="text-[10px] font-bold">FILES (<span x-text="filesList.length"></span>)</span>
        </button>
    </nav>

    <script>
        function puruApp() {
            return {
                sessionId: 'sid_' + Math.random().toString(36).substr(2, 9),
                screen: 'setup', activeTab: 'chat', statusText: 'Idle', statusType: 'idle',
                isLooping: false, isMenuOpen: false, filesList: [], chatHistory: [], userInput: '', ws: null,
                
                get statusColor() { 
                    return {idle: 'bg-gray-500', active: 'bg-blue-500 blink', error: 'bg-red-500', done: 'bg-green-500'}[this.statusType] || 'bg-gray-500'; 
                },
                
                async init() {
                    if(localStorage.getItem('puru_started')) {
                        this.sessionId = localStorage.getItem('puru_sid') || this.sessionId;
                        this.screen = 'workspace';
                        this.connectWS();
                    }
                },
                
                // Parser HTML Escape & Highlight Tag XML
                formatXML(text) {
                    if (!text) return '';
                    let safeText = text.replace(/&/g, "&amp;").replace(/</g, "&lt;").replace(/>/g, "&gt;");
                    safeText = safeText.replace(/(&lt;\/?[\w\s="\'-]+&gt;)/g, '<span class="xml-tag">$1</span>');
                    return safeText;
                },
                
                connectWS() {
                    this.ws = new WebSocket(`${location.protocol==='https:'?'wss:':'ws:'}//${location.host}/ws/${this.sessionId}`);
                    this.ws.onmessage = (e) => {
                        const m = JSON.parse(e.data);
                        if(m.type==='chat_update') this.chatHistory = m.data;
                        if(m.type==='vfs') this.filesList = m.data;
                        if(m.type==='status') { 
                            this.statusText = m.text; 
                            this.statusType = m.statusType; 
                            if(m.isLooping!==undefined) this.isLooping=m.isLooping; 
                        }
                        this.$nextTick(() => { 
                            const c = document.getElementById('chatContainer'); 
                            if(c) c.scrollTop = c.scrollHeight; 
                        });
                    };
                },
                
                startFreshProject() {
                    this.screen = 'workspace';
                    localStorage.setItem('puru_started', 'true');
                    localStorage.setItem('puru_sid', this.sessionId);
                    this.connectWS();
                },
                
                async handleZipUpload(e) {
                    const fd = new FormData(); fd.append('session_id', this.sessionId); fd.append('file', e.target.files[0]);
                    await fetch('/api/upload_zip', {method: 'POST', body: fd});
                    this.startFreshProject();
                },
                
                sendPrompt() {
                    this.ws.send(JSON.stringify({action: 'prompt', text: this.userInput}));
                    this.userInput = ''; this.isLooping = true;
                    this.statusType = 'active'; this.statusText = 'Berpikir...';
                },
                
                async clearSession() {
                    const fd = new FormData(); fd.append('session_id', this.sessionId);
                    await fetch('/api/clear_session', {method: 'POST', body: fd});
                    this.chatHistory = [];
                },
                
                async resetAll() {
                    const fd = new FormData(); fd.append('session_id', this.sessionId);
                    await fetch('/api/delete_session', {method: 'POST', body: fd});
                    localStorage.clear(); location.reload();
                }
            }
        }
    </script>
</body>
</html>
"""

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=PORT)