app.py

import streamlit as st
import requests
import json
import os
import psutil
import platform
import time
import pypdf
from concurrent.futures import ThreadPoolExecutor
from streamlit.runtime.scriptrunner import add_script_run_ctx, get_script_run_ctx
from agent import AgentWorkspace, CodeExecutor, extract_code_blocks, parse_filename_from_code

# -------------------- CONFIG --------------------
st.set_page_config(page_title="Neon AI", layout="wide", page_icon="⚡")

MEMORY_FILE = "/app/data/memory.json" if os.path.exists("/app/data") else "memory.json"

# -------------------- MEMORY SYSTEM --------------------
def load_memory():
    if os.path.exists(MEMORY_FILE):
        try:
            with open(MEMORY_FILE, "r") as f:
                return json.load(f)
        except:
            return []
    return []

def save_memory(messages):
    try:
        with open(MEMORY_FILE, "w") as f:
            json.dump(messages, f, indent=2)
    except:
        pass

if "messages" not in st.session_state:
    st.session_state.messages = load_memory()
if "file_context" not in st.session_state:
    st.session_state.file_context = ""

# Initialize Agent System
if "workspace" not in st.session_state:
    st.session_state.workspace = AgentWorkspace()
    st.session_state.executor = CodeExecutor(st.session_state.workspace)
if "terminal_log" not in st.session_state:
    st.session_state.terminal_log = "Welcome to Neon Terminal v1.0\n"
if "trigger_gen" not in st.session_state:
    st.session_state.trigger_gen = False

# -------------------- CSS --------------------
st.markdown("""
<style>
    /* Main Background */
    .stApp {
        background: radial-gradient(circle at top left, #1a0b2e, #000000);
        background-attachment: fixed;
    }
    
    /* Chat Bubbles */
    .user-msg {
        background: linear-gradient(135deg, #ff00cc 0%, #333 100%);
        color: white;
        padding: 12px 18px;
        border-radius: 18px 18px 2px 18px;
        max-width: 80%;
        float: right;
        clear: both;
        margin-bottom: 10px;
        box-shadow: 0 4px 10px rgba(0,0,0,0.3);
    }
    .bot-msg {
        background: rgba(255, 255, 255, 0.08);
        border-left: 3px solid #00ffe0;
        color: #e0e0e0;
        padding: 15px;
        border-radius: 18px 18px 18px 2px;
        max-width: 85%;
        float: left;
        clear: both;
        margin-bottom: 10px;
    }
    
    /* Input Field */
    .stTextInput input {
        background: rgba(20, 20, 20, 0.8) !important;
        color: white !important;
        border: 1px solid #333 !important;
    }

    /* Sidebar Gauge Container */
    .gauge-container {
        display: flex;
        justify-content: space-between;
        background: rgba(255,255,255,0.05);
        padding: 10px;
        border-radius: 10px;
        margin-bottom: 10px;
    }
    .gauge-box {
        width: 48%;
        text-align: center;
    }
    .gauge-label {
        font-size: 11px;
        color: #aaa;
        margin-top: 4px;
    }
    .gauge-val {
        font-size: 10px;
        color: #fff;
    }
    
    /* Clean up Streamlit UI */
    #MainMenu {visibility: hidden;}
    footer {visibility: hidden;}

    /* Terminal Style */
    .terminal-window {
        background-color: #0c0c0c;
        color: #00ff00;
        font-family: 'Courier New', Courier, monospace;
        padding: 15px;
        border-radius: 5px;
        height: 300px;
        overflow-y: auto;
        border: 1px solid #333;
        font-size: 12px;
        white-space: pre-wrap;
    }
</style>
""", unsafe_allow_html=True)

# -------------------- HELPERS --------------------
def get_size(bytes, suffix="B"):
    factor = 1024
    for unit in ["", "K", "M", "G", "T", "P"]:
        if bytes < factor:
            return f"{bytes:.1f}{unit}{suffix}"
        bytes /= factor

def get_system_stats():
    cpu = psutil.cpu_percent(interval=0.1)
    ram = psutil.virtual_memory()
    disk = psutil.disk_usage('/')
    return cpu, ram.percent, get_size(ram.used), get_size(ram.total), disk.percent, get_size(disk.free), get_size(disk.total)

def process_uploaded_file(uploaded_file):
    if uploaded_file is None:
        return ""

    try:
        uploaded_file.seek(0)
        # PDF Processing
        if uploaded_file.type == "application/pdf":
            reader = pypdf.PdfReader(uploaded_file)
            text = ""
            for page in reader.pages:
                text += page.extract_text() + "\n"
            return text.strip()

        # Image Processing (Placeholder for now, just description)
        elif uploaded_file.type in ["image/png", "image/jpeg", "image/jpg"]:
            return f"[User uploaded an image: {uploaded_file.name}]"

        else:
            return f"[Unsupported file type: {uploaded_file.type}]"

    except Exception as e:
        return f"[Error processing file: {e}]"

# -------------------- SIDEBAR --------------------
with st.sidebar:
    st.markdown("### 🖥️ SYSTEM HUD")
    
    cpu, ram_p, ram_u, ram_t, disk_p, disk_f, disk_t = get_system_stats()
    
    cpu_c = "#ff4b4b" if cpu > 80 else "#00ffe0"
    ram_c = "#ff4b4b" if ram_p > 85 else "#ff00cc"
    
    # SVG Strings (Compact)
    cpu_svg = f'<svg viewBox="0 0 36 36" style="max-height:80px;"><path d="M18 2.0845 a 15.9155 15.9155 0 0 1 0 31.831 a 15.9155 15.9155 0 0 1 0 -31.831" fill="none" stroke="#444" stroke-width="3.8"/><path d="M18 2.0845 a 15.9155 15.9155 0 0 1 0 31.831 a 15.9155 15.9155 0 0 1 0 -31.831" fill="none" stroke="{cpu_c}" stroke-width="2.8" stroke-dasharray="{cpu}, 100"/><text x="18" y="20.35" fill="#fff" font-size="8px" text-anchor="middle" font-weight="bold">{cpu}%</text></svg>'
    ram_svg = f'<svg viewBox="0 0 36 36" style="max-height:80px;"><path d="M18 2.0845 a 15.9155 15.9155 0 0 1 0 31.831 a 15.9155 15.9155 0 0 1 0 -31.831" fill="none" stroke="#444" stroke-width="3.8"/><path d="M18 2.0845 a 15.9155 15.9155 0 0 1 0 31.831 a 15.9155 15.9155 0 0 1 0 -31.831" fill="none" stroke="{ram_c}" stroke-width="2.8" stroke-dasharray="{ram_p}, 100"/><text x="18" y="20.35" fill="#fff" font-size="8px" text-anchor="middle" font-weight="bold">{ram_p}%</text></svg>'
    
    st.markdown(f"""
    <div class="gauge-container">
        <div class="gauge-box">{cpu_svg}<div class="gauge-label">CPU Load</div></div>
        <div class="gauge-box">{ram_svg}<div class="gauge-label">RAM Usage</div><div class="gauge-val">{ram_u} / {ram_t}</div></div>
    </div>
    """, unsafe_allow_html=True)
    
    st.markdown(f"""
    <div style="background:rgba(255,255,255,0.05); padding:10px; border-radius:8px; font-size:12px;">
        <div style="display:flex; justify-content:space-between; margin-bottom:5px;">
            <span style="color:#aaa;">Storage</span><span style="color:#fff;">{disk_f} Free</span>
        </div>
        <div style="width:100%; height:6px; background:#444; border-radius:3px;">
            <div style="width:{disk_p}%; height:100%; background:linear-gradient(90deg, #00ffe0, #0077ff); border-radius:3px;"></div>
        </div>
        <div style="text-align:right; font-size:10px; color:#aaa; margin-top:3px;">{disk_p}% Full of {disk_t}</div>
    </div>
    <hr style="border-color:rgba(255,255,255,0.1);">
    """, unsafe_allow_html=True)
    
    # ---------------- MODEL SELECTION LOGIC ----------------
    st.markdown("### 🧠 Neural Core")
    
    # 1. Local Models (Ollama)
    local_models = []
    try:
        models_json = requests.get("http://localhost:11434/api/tags", timeout=1).json()
        local_models = [m["name"] for m in models_json["models"]]
    except:
        local_models = ["qwen2.5-coder:0.5b"]

    # 2. Cloud Dolphin Model
    cloud_dolphin = "☁️ Dolphin 24B (Free API)"

    # 3. Hugging Face Models (From ENV & Defaults)
    default_hf_models = [
        "dphn/Dolphin-Mistral-24B-Venice-Edition:featherless-ai",
        "google/gemma-7b",
        "moonshotai/Kimi-K2.5:novita",
        "qwen2.5-coder:0.5b"
    ]
    hf_models_env = os.environ.get("HF_MODELS", "")
    env_models = [m.strip() for m in hf_models_env.split(",") if m.strip()]
    hf_models = list(dict.fromkeys(env_models + default_hf_models))  # Remove duplicates while preserving order
    
    # Combine All
    all_models = [cloud_dolphin] + hf_models + local_models
    
    arena_mode = st.toggle("⚔️ Arena Mode (Multi-Model)", value=False)

    if arena_mode:
        selected_models = st.multiselect("Select Models", all_models, default=[all_models[0]] if all_models else [])
    else:
        selected_model = st.selectbox("Active Model", all_models)
        selected_models = [selected_model]
    
    # Brain Template (Visible for Cloud/HF Models)
    template_mode = st.selectbox(
        "Thinking Style", 
        ["creative", "logical", "code-advanced", "summary"],
        index=0
    )
    # Mode Toggle
    agent_mode = st.toggle("🛠️ Agent Mode (Code & Exec)", value=False)

    if agent_mode:
        auto_gpt_mode = st.toggle("🤖 Autonomous Agent (Auto-Fix)", value=False)
    else:
        auto_gpt_mode = False

    # File Upload
    # st.markdown("### 📂 Data Injection")  # Removed per user request

    # Collapsible File Upload
    with st.expander("📂 Upload File (PDF/Image)", expanded=False):
        uploaded_file = st.file_uploader("Drag and drop file here", type=["pdf", "png", "jpg", "jpeg"], label_visibility="collapsed")

    if uploaded_file:
        file_text = process_uploaded_file(uploaded_file)
        if file_text:
            st.session_state.file_context = file_text
            st.success("File processed & injected!")
            if uploaded_file.type != "application/pdf":
                st.image(uploaded_file, caption="Uploaded Image", use_container_width=True)
        else:
            st.session_state.file_context = ""
    else:
        st.session_state.file_context = ""

    col1, col2 = st.columns(2)
    if col1.button("♻ Refresh"): st.rerun()
    if col2.button("🧹 Clear"): 
        st.session_state.messages = []
        save_memory([])
        st.rerun()

# -------------------- MAIN CHAT --------------------
st.markdown("<h1 style='text-align:center;'>⚡ NEON <span style='color:#00ffe0; text-shadow:0 0 15px rgba(0,255,224,0.6);'>AI</span></h1>", unsafe_allow_html=True)

# Layout Setup
main_chat_container = st.container()
workspace_container = None

if agent_mode:
    col1, col2 = st.columns([1.2, 1])
    with col1:
        main_chat_container = st.container()
    with col2:
        workspace_container = st.container()

with main_chat_container:
    # Persistent Model Header
    if not arena_mode and selected_models:
        st.markdown(f"<h3 style='text-align:center; color:#aaa; margin-top:-20px;'>Active: <span style='color:#00ffe0;'>{selected_models[0]}</span></h3>", unsafe_allow_html=True)

    if arena_mode and selected_models:
        cols = st.columns(len(selected_models))
        for idx, model_name in enumerate(selected_models):
            with cols[idx]:
                # Sticky Header
                st.markdown(f"""
                <div style='
                    position: sticky;
                    top: 0;
                    z-index: 100;
                    background: rgba(26, 11, 46, 0.95);
                    backdrop-filter: blur(10px);
                    padding: 10px;
                    border-radius: 0 0 10px 10px;
                    text-align: center;
                    color: #00ffe0;
                    font-weight: bold;
                    border-bottom: 2px solid #00ffe0;
                    margin-bottom: 15px;
                    box-shadow: 0 4px 6px rgba(0,0,0,0.3);
                '>🧠 {model_name}</div>
                """, unsafe_allow_html=True)

                for msg in st.session_state.messages:
                    if msg["role"] == "user":
                        st.markdown(f"<div class='user-msg' style='float:none; margin-left:auto; margin-right:0;'>👤 {msg['content']}</div>", unsafe_allow_html=True)
                    elif msg["role"] == "assistant":
                        if msg.get("model") == model_name or msg.get("model") is None:
                            st.markdown(f"<div class='bot-msg' style='float:none; margin-right:auto; margin-left:0;'>🤖 {msg['content']}</div>", unsafe_allow_html=True)
    else:
        for msg in st.session_state.messages:
            if msg["role"] == "user":
                st.markdown(f"<div class='user-msg'>👤 {msg['content']}</div>", unsafe_allow_html=True)
            else:
                content = msg['content']
                tag = f"<br><span style='font-size:0.7em; color:#888;'>({msg.get('model')})</span>" if msg.get("model") else ""
                st.markdown(f"<div class='bot-msg'>🤖 {content}{tag}</div>", unsafe_allow_html=True)

    st.markdown("<div style='clear:both;'></div>", unsafe_allow_html=True)

# -------------------- WORKSPACE UI (Right Column) --------------------
if agent_mode and workspace_container:
    with workspace_container:
        st.markdown("### 🛠️ Workspace")

        # Files Tab
        files = st.session_state.workspace.list_files()
        selected_file = st.selectbox("📂 Files", ["(New File)"] + files)

        file_content = ""
        if selected_file and selected_file != "(New File)":
            file_content = st.session_state.workspace.read_file(selected_file)

        # Editor (Read Only for now)
        st.code(file_content if file_content else "# No file selected", language="python", line_numbers=True)

        # Actions
        c1, c2 = st.columns(2)
        if c1.button("▶ Run File", disabled=not selected_file):
            if selected_file.endswith(".py"):
                output = st.session_state.executor.run_python(selected_file)
                st.session_state.terminal_log += f"\n$ python {selected_file}\n{output}\n"
            elif selected_file.endswith(".sh"):
                output = st.session_state.executor.run_command(f"bash {selected_file}")
                st.session_state.terminal_log += f"\n$ bash {selected_file}\n{output}\n"
            else:
                st.session_state.terminal_log += f"\n$ {selected_file} is not executable.\n"
            st.rerun()

        if c2.button("💾 Save to Memory"):
             # Placeholder for saving to RAG or Context
             pass

        # Terminal
        st.markdown("#### 📟 Terminal Output")
        st.markdown(f"<div class='terminal-window'>{st.session_state.terminal_log}</div>", unsafe_allow_html=True)

        if st.button("Clear Terminal"):
            st.session_state.terminal_log = "$ \n"
            st.rerun()

prompt = st.chat_input("Inject code or query...")

should_run = False
if prompt:
    # 1. Append User Message
    st.session_state.messages.append({"role": "user", "content": prompt})
    save_memory(st.session_state.messages)
    should_run = True

is_triggered_run = False
if st.session_state.trigger_gen:
    st.session_state.trigger_gen = False
    should_run = True
    is_triggered_run = True
    # Recover prompt from last message for display purposes
    if st.session_state.messages and st.session_state.messages[-1]["role"] == "user":
        prompt = st.session_state.messages[-1]["content"]

if should_run:
    # Determine containers for streaming
    if arena_mode and selected_models:
        with main_chat_container:
            live_containers = st.columns(len(selected_models))
    elif auto_gpt_mode:
        # Full width for autonomous mode visualization
        with main_chat_container:
            if not is_triggered_run:
                 st.markdown(f"<div class='user-msg'>👤 {prompt}</div>", unsafe_allow_html=True)
            live_containers = [st.container()]
    else:
        with main_chat_container:
            if not is_triggered_run:
                st.markdown(f"<div class='user-msg'>👤 {prompt}</div>", unsafe_allow_html=True)
            live_containers = [st.container()]

    # Capture the current context to pass to threads
    main_ctx = get_script_run_ctx()

    def run_autonomous_loop(model_name, prompt, container, ctx):
        """Runs an Auto-GPT style loop: Generate -> Run -> Fix -> Repeat."""
        add_script_run_ctx(ctx=ctx)

        max_retries = 5
        current_prompt = prompt
        history = [
            {"role": "system", "content": "You are an autonomous coding agent. Write complete, runnable Python code to solve the user's problem. If you encounter errors, analyze them and fix your code. Always output the full corrected code block."}
        ]
        if st.session_state.file_context:
            history.append({"role": "system", "content": f"CONTEXT FILE:\n{st.session_state.file_context}"})

        history.append({"role": "user", "content": current_prompt})

        with container:
            st.markdown(f"### 🤖 Autonomous Agent: {model_name}")
            status_container = st.container()

            for attempt in range(max_retries):
                with status_container.status(f"Attempt {attempt+1}/{max_retries}", expanded=True) as status:

                    # 1. Generate Code
                    st.write("🧠 Thinking & Writing Code...")
                    full_response = ""

                    # --- MODEL CALL (Simplified for Loop) ---
                    try:
                        # Prepare Context
                        msgs = history.copy()

                        # API Call (Generic Adapter)
                        if model_name in hf_models:
                            hf_token = os.environ.get("HF_TOKEN")
                            api_url = "https://router.huggingface.co/v1/chat/completions"
                            headers = {"Authorization": f"Bearer {hf_token}"}
                            # Increased max_tokens to prevent code truncation
                            payload = {"model": model_name, "messages": msgs, "stream": False, "max_tokens": 4096}
                            resp = requests.post(api_url, headers=headers, json=payload).json()
                            full_response = resp["choices"][0]["message"]["content"]

                        elif model_name == cloud_dolphin:
                            api_url = "https://chat.dphn.ai/api/chat"
                            payload = {"messages": msgs, "model": "dolphinserver:24B", "template": "code-advanced"}
                            resp = requests.post(api_url, json=payload, headers={"Content-Type": "application/json"}).json()
                            full_response = resp["choices"][0]["message"]["content"]

                        else: # Local Ollama
                            resp = requests.post("http://localhost:11434/api/chat", json={"model": model_name, "messages": msgs, "stream": False}).json()
                            full_response = resp["message"]["content"]

                    except Exception as e:
                        status.update(label="❌ Generation Failed", state="error")
                        st.error(f"Model Error: {e}")
                        break
                    # ----------------------------------------

                    st.markdown(full_response)
                    history.append({"role": "assistant", "content": full_response})

                    # 2. Extract & Save
                    blocks = extract_code_blocks(full_response)
                    if not blocks:
                        status.update(label="⚠️ No Code Generated", state="complete")
                        st.warning("Model didn't generate any code.")
                        return # Exit loop if no code

                    # Assume last block is the main script
                    code_block = blocks[-1]
                    code = code_block['code']
                    lang = code_block['language']

                    if lang not in ["python", "py", "bash", "sh"]:
                        status.update(label="⚠️ Non-Executable Code", state="complete")
                        st.info(f"Generated {lang} code, skipping execution.")
                        return

                    filename = parse_filename_from_code(code)
                    if not filename:
                        ext = "py" if lang in ["python", "py"] else "sh"
                        filename = f"agent_script_{int(time.time())}.{ext}"

                    st.session_state.workspace.save_file(filename, code)
                    st.write(f"💾 Saved: `{filename}`")

                    # 3. Execute
                    st.write("⚙️ Executing...")
                    stdout, stderr, exit_code = "", "", 0

                    if filename.endswith(".py"):
                        stdout, stderr, exit_code = st.session_state.executor.run_python_safe(filename)
                    elif filename.endswith(".sh"):
                        stdout, stderr, exit_code = st.session_state.executor.run_command_safe(f"bash {filename}")

                    # 4. Check Results
                    if exit_code == 0 and not stderr:
                        status.update(label="✅ Success!", state="complete")
                        st.success("Execution Successful!")
                        st.code(stdout if stdout else "# No Output")
                        st.session_state.terminal_log += f"\n[AUTO-AGENT SUCCESS] {filename}\n{stdout}\n"
                        return # DONE!
                    else:
                        status.update(label=f"❌ Failed (Exit: {exit_code})", state="error")
                        error_msg = f"Standard Output:\n{stdout}\n\nStandard Error:\n{stderr}"
                        st.error(f"Execution Failed:\n{stderr}")
                        st.session_state.terminal_log += f"\n[AUTO-AGENT ERROR] {filename}\n{stderr}\n"

                        # 5. Loop Back
                        feedback = f"The code you wrote in {filename} failed with exit code {exit_code}.\n\nOUTPUT:\n{stdout}\n\nERROR:\n{stderr}\n\nPlease fix the code and output the full corrected script."
                        history.append({"role": "user", "content": feedback})
                        st.write("🔄 Looping for Fix...")
                        time.sleep(1)

            st.error("❌ Max retries reached. Agent could not solve the problem.")

    def run_chat_thread(model_name, container, ctx):
        # Manually attach context to this thread
        add_script_run_ctx(ctx=ctx)

        with container:
            if arena_mode and not is_triggered_run:
                st.markdown(f"<div class='user-msg' style='float:none; margin-left:auto; margin-right:0;'>👤 {prompt}</div>", unsafe_allow_html=True)
            
            msg_placeholder = st.empty()
            full_response = ""
            bot_style = "style='float:none; margin-right:auto; margin-left:0;'" if arena_mode else ""

            try:
                # Prepare Context
                model_history = [
                    m for m in st.session_state.messages
                    if m["role"] == "user" or (m["role"] == "assistant" and (m.get("model") == model_name or m.get("model") is None))
                ]

                # INJECT FILE CONTEXT IF AVAILABLE
                if st.session_state.file_context:
                    # We inject it as a system message at the start, or append to the last user message if no system support?
                    # System message is cleaner.
                    context_msg = f"CONTEXT FROM UPLOADED FILE:\n{st.session_state.file_context}\n\nUSER QUERY:\n"
                    # For RAG, it's often better to prepend to the latest user prompt or add as a system instruction.
                    # Let's add as a system instruction if the model supports it, or prepend to the last user message.

                    # Simpler approach: Prepend to the first message if it's user, or insert system message.
                    # We will insert a system message at index 0.
                    model_history.insert(0, {"role": "system", "content": f"You have access to the following file content. Use it to answer questions if relevant:\n\n{st.session_state.file_context}"})

                # ROUTE 1: HUGGING FACE
                if model_name in hf_models:
                    hf_token = os.environ.get("HF_TOKEN")
                    if not hf_token: raise Exception("HF_TOKEN not found.")
                    api_url = "https://router.huggingface.co/v1/chat/completions"
                    headers = {"Authorization": f"Bearer {hf_token}"}
                    payload = {"model": model_name, "messages": model_history, "stream": True, "max_tokens": 1024}
                    response = requests.post(api_url, headers=headers, json=payload, stream=True)
                    for line in response.iter_lines():
                        if not line: continue
                        if line.startswith(b"data: "):
                            line_data = line.decode("utf-8").lstrip("data: ").strip()
                            if line_data == "[DONE]": break
                            try:
                                chunk = json.loads(line_data)
                                content = chunk["choices"][0]["delta"].get("content", "")
                                if content:
                                    full_response += content
                                    msg_placeholder.markdown(f"<div class='bot-msg' {bot_style}>🤖 {full_response}▌</div>", unsafe_allow_html=True)
                            except: pass

                # ROUTE 2: CLOUD DOLPHIN
                elif model_name == cloud_dolphin:
                    api_url = "https://chat.dphn.ai/api/chat"
                    payload = {"messages": model_history, "model": "dolphinserver:24B", "template": template_mode}
                    headers = {"Content-Type": "application/json"}
                    response = requests.post(api_url, json=payload, headers=headers, stream=True)
                    for line in response.iter_lines():
                        if not line: continue
                        line_text = line.decode("utf-8")
                        if line_text.startswith("data: "):
                            json_str = line_text[6:]
                            if json_str.strip() == "[DONE]": break
                            try:
                                chunk = json.loads(json_str)
                                if "choices" in chunk and len(chunk["choices"]) > 0:
                                    content = chunk["choices"][0]["delta"].get("content", "")
                                    if content:
                                        full_response += content
                                        msg_placeholder.markdown(f"<div class='bot-msg' {bot_style}>🤖 {full_response}▌</div>", unsafe_allow_html=True)
                            except: pass

                # ROUTE 3: LOCAL OLLAMA
                else:
                    current_msgs = model_history.copy()
                    # If file context was already inserted into model_history, we are good.
                    # But if template_mode is not creative, we might have conflicting system prompts.
                    # Let's just append the template mode system prompt if needed.
                    if template_mode != "creative":
                        # Check if we already have a system prompt (from file)
                        if current_msgs and current_msgs[0]["role"] == "system":
                            current_msgs[0]["content"] += f"\n\nAlso, act as a {template_mode} assistant."
                        else:
                            current_msgs.insert(0, {"role": "system", "content": f"You are a {template_mode} assistant."})

                    response = requests.post(
                        "http://localhost:11434/api/chat",
                        json={"model": model_name, "messages": current_msgs, "stream": True},
                        stream=True
                    )
                    for line in response.iter_lines():
                        if line:
                            data = json.loads(line.decode("utf-8"))
                            if "message" in data:
                                content = data["message"]["content"]
                                full_response += content
                                msg_placeholder.markdown(f"<div class='bot-msg' {bot_style}>🤖 {full_response}▌</div>", unsafe_allow_html=True)

                # Final Render & Return
                msg_placeholder.markdown(f"<div class='bot-msg' {bot_style}>🤖 {full_response}</div>", unsafe_allow_html=True)

                # --- AGENT: Auto-Save Code ---
                if agent_mode:
                    blocks = extract_code_blocks(full_response)
                    for block in blocks:
                        code = block['code']
                        lang = block['language']

                        # Guess Filename
                        filename = parse_filename_from_code(code)
                        if not filename:
                            ext = "py" if lang == "python" else "sh" if lang in ["bash", "sh"] else "txt"
                            filename = f"script_{int(time.time())}.{ext}"

                        st.session_state.workspace.save_file(filename, code)
                        try:
                            st.toast(f"💾 Saved: {filename}", icon="✅")
                        except: pass
                # -----------------------------

                return {"role": "assistant", "content": full_response, "model": model_name}

            except Exception as e:
                msg_placeholder.error(f"⚠️ Error: {e}")
                return None

    # Run Threads
    with ThreadPoolExecutor(max_workers=len(selected_models)) as executor:
        futures = []
        for model, container in zip(selected_models, live_containers):
            if agent_mode and auto_gpt_mode:
                 # Autonomous Loop
                 future = executor.submit(run_autonomous_loop, model, prompt, container, main_ctx)
            else:
                 # Standard Chat
                 future = executor.submit(run_chat_thread, model, container, main_ctx)
            futures.append(future)

        # Wait for all to complete
        for future in futures:
            try:
                result = future.result()
                # Autonomous loop handles its own history/display, so we might return None or handle differently.
                # Standard chat returns a dict to append.
                if result and not auto_gpt_mode:
                    st.session_state.messages.append(result)
            except Exception as e:
                st.error(f"Thread failed: {e}")

        if not auto_gpt_mode:
            save_memory(st.session_state.messages)

        # --- AGENT: Cross-Review Trigger ---
        if agent_mode and not auto_gpt_mode and len(selected_models) > 1:
            with main_chat_container:
                if st.button("🔍 Peer Review"):
                    # Trigger a review by the other model
                    # We simply append a user message requesting review and rerun
                    # Ideally, we target the *other* model, but our current system broadcasts.
                    # We can craft a prompt: "Model [Name], please review the above."

                    review_prompt = "Please review the code and response generated above. Check for errors, security issues, and suggest improvements."
                    st.session_state.messages.append({"role": "user", "content": review_prompt})
                    save_memory(st.session_state.messages)
                    st.session_state.trigger_gen = True
                    st.rerun()