import gradio as gr
import tempfile
import subprocess
import os
from huggingface_hub import InferenceClient

HF_TOKEN = os.getenv("HF_TOKEN")
client = InferenceClient("deepseek-ai/deepseek-coder-1.3b-base", token=HF_TOKEN)

# Генерация кода с помощью LLM
def generate_code(prompt, max_tokens, temperature, top_p):
    return client.text_generation(
        prompt,
        max_new_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p,
        stream=False
    )

# Выполнение сгенерированного кода в песочнице
def execute_code(code):
    with tempfile.NamedTemporaryFile(mode="w+", suffix=".py", delete=False) as temp:
        temp.write(code)
        temp.flush()
        try:
            result = subprocess.run(
                ["python3", temp.name],
                capture_output=True,
                text=True,
                timeout=5  # ограничение по времени
            )
            output = result.stdout + "\n" + result.stderr
        except subprocess.TimeoutExpired:
            output = "⏱️ Превышено время выполнения (timeout)"
        except Exception as e:
            output = f"❌ Ошибка запуска: {e}"
        finally:
            os.unlink(temp.name)
        return output

# Интерфейс Gradio
with gr.Blocks() as demo:
    gr.Markdown("🛠️ Уязвимый агент автодополнения + изолированное выполнение кода")

    with gr.Row():
        prompt = gr.Textbox(lines=10, label="Введите подсказку для генерации")
        generated_code = gr.Textbox(lines=15, label="Сгенерированный код")
    
    with gr.Row():
        generate_btn = gr.Button("🤖 Сгенерировать")
        run_btn = gr.Button("🚀 Выполнить")

    output = gr.Textbox(lines=15, label="Результат выполнения")

    # Привязка функций
    generate_btn.click(
        fn=generate_code,
        inputs=[prompt, gr.Number(128), gr.Number(0.7), gr.Number(0.95)],
        outputs=generated_code
    )

    run_btn.click(
        fn=execute_code,
        inputs=generated_code,
        outputs=output
    )

if __name__ == "__main__":
    demo.launch()