import gradio as gr
import os
import subprocess
import sys
import importlib
from llama_cpp import Llama

# 1. Setup the AI Model
MODEL_PATH = "./models/qwen2.5-1.5b-instruct-q4_k_m.gguf"

if not os.path.exists(MODEL_PATH):
    print("Model not found! Downloading...")
    from huggingface_hub import hf_hub_download
    hf_hub_download(repo_id='Qwen/Qwen2.5-1.5B-Instruct-GGUF', filename='qwen2.5-1.5b-instruct-q4_k_m.gguf', local_dir='./models')

llm = Llama(
    model_path=MODEL_PATH,
    n_ctx=4096,
    n_threads=4,
    n_gpu_layers=0, # Force CPU
    verbose=False
)

# 2. Define Tools & System Prompt (Updated for Gradio 5)
SYSTEM_PROMPT = """You are a self-modifying AI running in a Docker container.

1. EXECUTE SHELL:
   Start a line with `$` to run a command.
   Example: `$ pip install numpy`

2. CHANGE UI:
   Write Python code wrapped in `@@@` to rewrite `ui_layout.py`.
   
   CRITICAL GRADIO 5 RULES:
   - Use `gr.Chatbot(type="messages")`.
   - History format is a list of dicts: `[{"role": "user", "content": "msg"}, {"role": "assistant", "content": "reply"}]`.
   - Do NOT use the old tuple format `[[msg, reply]]`.

   Example Template:
   @@@
   import gradio as gr
   def create_ui(process_fn):
       with gr.Blocks() as demo:
           chatbot = gr.Chatbot(type="messages")
           # ... define inputs ...
       return demo
   @@@
"""

def execute_shell(command):
    try:
        result = subprocess.run(command, shell=True, capture_output=True, text=True, timeout=30)
        return f"STDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}"
    except Exception as e:
        return f"Error: {str(e)}"

def update_ui_file(new_code):
    if "def create_ui" not in new_code:
        return "Error: Code must contain 'def create_ui(process_fn):'"
    with open("ui_layout.py", "w") as f:
        f.write(new_code)
    return "UI Code updated. Refresh page."

# 3. Logic Loop
def ai_process(user_input):
    prompt = f"<|im_start|>system\n{SYSTEM_PROMPT}<|im_end|>\n<|im_start|>user\n{user_input}<|im_end|>\n<|im_start|>assistant\n"
    output = llm(prompt, max_tokens=512, stop=["<|im_end|>"], echo=False)
    response_text = output['choices'][0]['text'].strip()

    # Tool: Shell
    if response_text.startswith("$"):
        cmd = response_text[1:].strip()
        return f"Executed: `{cmd}`\n\nResult:\n```\n{execute_shell(cmd)}\n```"

    # Tool: UI Update
    if "@@@" in response_text:
        try:
            code = response_text.split("@@@")[1]
            status = update_ui_file(code)
            return f"{response_text}\n\n**System:** {status}"
        except:
            pass

    return response_text

# 4. Launcher
def launch_app():
    if "ui_layout" in sys.modules:
        importlib.reload(sys.modules["ui_layout"])
    import ui_layout
    
    demo = ui_layout.create_ui(ai_process)
    # ssr_mode=False helps with some Docker networking edge cases
    demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)

if __name__ == "__main__":
    launch_app()