import gradio as gr
import torch
import sys
from io import StringIO
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

print("Loading Base Model...")
base_model_name = "Qwen/Qwen2.5-Coder-1.5B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(base_model_name)
base_model = AutoModelForCausalLM.from_pretrained(
    base_model_name,
    torch_dtype=torch.float16,
    device_map="auto"
)

print("Loading Trained LoRA Weights...")
try:
    model = PeftModel.from_pretrained(base_model, "MaxKio/lammo-coder-Max")
except Exception as e:
    print(f"Warning: Could not load custom weights, using base model: {e}")
    model = base_model

def run_python_code(code_str):
    """بيئة معزولة بسيطة لتشغيل كود البايثون المولد وعرض النتيجة"""
    old_stdout = sys.stdout
    redirected_output = sys.stdout = StringIO()
    try:
        exec(code_str, {})
        sys.stdout = old_stdout
        return redirected_output.getvalue() or "تم تنفيذ الكود بنجاح دون مخرجات (stdout)."
    except Exception as e:
        sys.stdout = old_stdout
        return f"خطأ أثناء تشغيل الكود: {e}"

def respond(message, chat_history):
    messages = []
    for val in chat_history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})
            
    messages.append({"role": "user", "content": message})
    
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
    
    generated_ids = model.generate(
        **model_inputs,
        max_new_tokens=512,
        do_sample=True,
        temperature=0.7,
        top_p=0.9
    )
    
    generated_ids = [output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)]
    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
    return response

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("""# <center> 🤖 AgentM - Coding Assistant 🤖 </center>
    ### <center>مساعدك البرمجي الشخصي القوي مع إمكانية تشغيل الأكواد البرمجية مباشرة ومحاكاة بيئة المطورين</center>""")
    
    with gr.Tab("الدردشة البرمجية"):
        chatbot = gr.ChatInterface(fn=respond)
        
    with gr.Tab("منصة تشغيل الأكواد (Coding Sandbox)"):
        gr.Markdown("### اكتب كود بايثون هنا لتشغيله ورؤية النتيجة فوراً")
        code_input = gr.Code(language="python", label="كود البايثون")
        run_btn = gr.Button("تشغيل الكود ⚡", variant="primary")
        code_output = gr.Textbox(label="مخرجات التنفيذ")
        
        run_btn.click(fn=run_python_code, inputs=code_input, outputs=code_output)

demo.launch()