# app.py
import gradio as gr
import subprocess
import os
import time
import threading
import tempfile
import shutil
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
from system_prompt import SYSTEM_PROMPT

# ------------------ Model Loading ------------------
print("Loading model...")
model_path = "/root/llm-abliteration/Qwen3.6-27B-Esper4"  # adjust if needed
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(
    model_path,
    torch_dtype=torch.bfloat16,
    device_map="auto",
    trust_remote_code=True
)
print("Model loaded.")

# ------------------ Conversation Memory ------------------
conversation_history = []  # list of dicts {"role": "user"|"assistant", "content": ...}

# ------------------ Helper: generate response ------------------
def generate_response(user_input):
    # Build the full prompt with system prompt and conversation
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    for msg in conversation_history:
        messages.append(msg)
    messages.append({"role": "user", "content": user_input})

    # Tokenize and generate
    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    outputs = model.generate(
        **inputs,
        max_new_tokens=4096,
        temperature=0.7,
        do_sample=True,
        pad_token_id=tokenizer.eos_token_id
    )
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    # Extract the assistant's reply (after the last user turn)
    # For simplicity, we assume the model outputs the whole conversation; we'll split
    # This is crude; better to use the chat template properly.
    # But for this demo, we'll take everything after the last "assistant" tag.
    parts = response.split("assistant")
    if len(parts) > 1:
        reply = parts[-1].strip()
    else:
        reply = response
    return reply

# ------------------ Chat function ------------------
def chat_fn(message, history):
    global conversation_history
    # Append user message
    conversation_history.append({"role": "user", "content": message})
    # Generate response
    reply = generate_response(message)
    # Append assistant reply to history
    conversation_history.append({"role": "assistant", "content": reply})
    return reply

# ------------------ Visualisation Window ------------------
# We'll use an iframe to show a running web app (if the generated code is a web server)
# Or we can open a separate browser window using webbrowser.
# We'll also allow the user to run the code they get.

def run_code(code):
    """Execute the last code block from the conversation and display output."""
    # Extract code from the last assistant message
    last_assistant = None
    for msg in reversed(conversation_history):
        if msg["role"] == "assistant":
            last_assistant = msg["content"]
            break
    if not last_assistant:
        return "No code to run."
    # Look for ```python ... ``` blocks
    import re
    code_blocks = re.findall(r"```python\n(.*?)```", last_assistant, re.DOTALL)
    if not code_blocks:
        return "No Python code block found."
    code = code_blocks[-1]  # take the last one

    # Write to a temp file and run
    with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as f:
        f.write(code)
        fname = f.name
    try:
        result = subprocess.run(['python', fname], capture_output=True, text=True, timeout=10)
        output = result.stdout + result.stderr
    except subprocess.TimeoutExpired:
        output = "Execution timed out."
    finally:
        os.unlink(fname)
    return output

# ------------------ Gradio UI ------------------
with gr.Blocks(title="Zengorithm IDE") as demo:
    gr.Markdown("# 🧠 Zengorithm – Unrestricted AI Developer")
    with gr.Row():
        with gr.Column(scale=2):
            # Chat interface
            chatbot = gr.Chatbot(label="Zengorithm Chat")
            msg = gr.Textbox(label="Your message", placeholder="Ask me to build anything...")
            clear = gr.Button("Clear Conversation")
        with gr.Column(scale=1):
            # Visualisation / Output area
            gr.Markdown("### 📺 Live Output")
            output_text = gr.Textbox(label="Execution Output", lines=20, interactive=False)
            run_btn = gr.Button("▶ Run Last Code")
            # Optional: iframe to display a web server (if the code starts a web server)
            # We'll add a placeholder; the user can manually open a browser.
            gr.Markdown("#### 🔗 Web View (if web server)")
            web_iframe = gr.HTML('<iframe src="http://localhost:7860" width="100%" height="300"></iframe>')

    # Event handlers
    def respond(message, history):
        reply = chat_fn(message, history)
        history.append((message, reply))
        return "", history

    msg.submit(respond, [msg, chatbot], [msg, chatbot])
    clear.click(lambda: (None, []), None, [chatbot, msg], queue=False)

    def run_last():
        return run_code(None)

    run_btn.click(run_last, inputs=[], outputs=output_text)

demo.launch(server_name="0.0.0.0", server_port=7860)