import gradio as gr
import subprocess
import threading

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct"

print("Loading model...")

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype="auto",
    device_map="cpu"
)

print("Model loaded")


def ask_ai(prompt):

    try:
        messages = [
            {
                "role": "user",
                "content": prompt
            }
        ]

        text = tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True
        )

        inputs = tokenizer(
            text,
            return_tensors="pt"
        )

        output = model.generate(
            **inputs,
            max_new_tokens=512,
            temperature=0.7,
            do_sample=True
        )

        result = tokenizer.decode(
            output[0],
            skip_special_tokens=True
        )

        return result

    except Exception as e:
        return str(e)


def execute_code(code):

    try:

        proc = subprocess.run(
            ["python3", "-c", code],
            capture_output=True,
            text=True,
            timeout=30
        )

        if proc.returncode == 0:
            return proc.stdout or "Done"

        return proc.stderr

    except Exception as e:
        return str(e)


def ai_and_run(prompt):

    ai_response = ask_ai(
        prompt +
        "\nReturn only executable Python code."
    )

    result = execute_code(ai_response)

    return (
        "AI RESPONSE:\n\n"
        + ai_response
        + "\n\nOUTPUT:\n\n"
        + result
    )


with gr.Blocks(theme=gr.themes.Soft()) as demo:

    gr.Markdown("# Qwen AI + Sandbox")

    with gr.Tab("Chat"):

        prompt = gr.Textbox(
            label="Message",
            lines=4
        )

        output = gr.Textbox(
            label="Response",
            lines=15
        )

        btn = gr.Button("Send")

        btn.click(
            ask_ai,
            prompt,
            output
        )

    with gr.Tab("Python Sandbox"):

        code = gr.Textbox(
            label="Python Code",
            lines=12
        )

        result = gr.Textbox(
            label="Output",
            lines=12
        )

        run = gr.Button("Run")

        run.click(
            execute_code,
            code,
            result
        )

    with gr.Tab("AI Generate & Run"):

        p = gr.Textbox(
            label="Instruction",
            lines=4
        )

        r = gr.Textbox(
            label="Result",
            lines=20
        )

        b = gr.Button("Generate & Run")

        b.click(
            ai_and_run,
            p,
            r
        )

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