# ============================================================
# CodeAgent - Sub-agent for code generation and execution
# Author: Mustafa Albakkar
# Compatible with Hugging Face Spaces (Gradio-based)
# ============================================================

import os, json, atexit, logging, traceback, io, sys, contextlib, tempfile
import gradio as gr
from llama_cpp import Llama

# ------------------------------------------------------------
# 🔧 إعداد التسجيل (Logs)
# ------------------------------------------------------------
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("CodeAgent")

# ------------------------------------------------------------
# 🧠 تحميل النموذج (Qwen2.5-Coder-14B-Instruct-GGUF)
# ------------------------------------------------------------
llm = None
try:
    logger.info("🔄 Attempting to load model...")
    llm = Llama.from_pretrained(
        repo_id="Qwen/Qwen2.5-Coder-14B-Instruct-GGUF",
        filename="qwen2.5-coder-14b-instruct-q6_k.gguf",
        n_gpu_layers=-1,
        n_ctx=4096,
        n_threads=4,
    )
    logger.info("✅ Model loaded successfully.")
except Exception as e:
    logger.exception("❌ Model load failed: %s", e)
    llm = None


# ------------------------------------------------------------
# 🧩 دوال مساعدة
# ------------------------------------------------------------
def safe_text(x):
    """تحويل النصوص أو القيم المعقدة إلى نص نظيف وصالح للعرض."""
    if isinstance(x, str):
        return x.strip()
    try:
        return json.dumps(x, ensure_ascii=False)
    except Exception:
        return str(x)


def compute_safe_max_tokens(prompt: str, max_ctx=4096):
    """تقدير عدد التوكنات الآمنة حسب طول الإدخال."""
    try:
        l = max(1, len(prompt.split()))
        avail = max_ctx - l - 1
        return max(128, min(1024, avail))
    except Exception:
        return 256


# ------------------------------------------------------------
# 🧮 أداة تنفيذ الكود بأمان
# ------------------------------------------------------------
def execute_python_code(code_str: str) -> str:
    """
    تقوم هذه الدالة بتنفيذ الكود البرمجي الناتج من النموذج داخل بيئة معزولة،
    وتعيد المخرجات النصية أو الأخطاء.
    """
    code_str = code_str.strip()

    # تنظيف الكود من علامات أو أمثلة غير ضرورية
    code_str = code_str.replace("```python", "").replace("```", "").strip()

    buffer = io.StringIO()
    try:
        # إعادة توجيه stdout و stderr
        with contextlib.redirect_stdout(buffer), contextlib.redirect_stderr(buffer):
            local_env = {}
            exec(code_str, {}, local_env)
        result = buffer.getvalue().strip()
        if not result:
            result = "✅ Code executed successfully (no printed output)."
        return result
    except Exception as e:
        tb = traceback.format_exc()
        return f"❌ Code execution error:\n{tb}"
    finally:
        buffer.close()


# ------------------------------------------------------------
# 🎯 الدالة الأساسية للتوليد والتنفيذ
# ------------------------------------------------------------
def generate_and_execute_fn(prompt: str) -> str:
    """
    - توليد الكود البرمجي باستخدام نموذج Qwen2.5-Coder
    - تنفيذ الكود مباشرة
    - إعادة النتائج إلى الوكيل الأساسي
    """
    try:
        prompt = safe_text(prompt)
        if not prompt:
            return "⚠️ No prompt provided."

        if llm is None:
            return "❌ Model not loaded."

        # --- تحضير إدخال النموذج بصيغة محادثة ---
        formatted_prompt = (
            f"<|im_start|>system\n"
            f"You are CodeAgent, a skilled Python coding assistant. "
            f"Generate correct, fully functional Python code that solves the given task.\n"
            f"Always wrap your code in triple backticks and do not include explanations.\n"
            f"<|im_end|>\n"
            f"<|im_start|>user\n{prompt}\n<|im_end|>\n"
            f"<|im_start|>assistant\n"
        )

        max_tokens = compute_safe_max_tokens(formatted_prompt)

        # --- تنفيذ التوليد ---
        try:
            out = llm(
                formatted_prompt,
                max_tokens=max_tokens,
                temperature=0.2,
                stop=["<|im_end|>"]
            )
        except TypeError:
            out = llm(
                formatted_prompt,
                max_new_tokens=max_tokens,
                temperature=0.2,
                stop=["<|im_end|>"]
            )

        # --- استخراج النص الناتج ---
        if isinstance(out, dict):
            text = out.get("choices", [{}])[0].get("text", "")
        else:
            text = str(out)

        if not text.strip():
            return "⚠️ Empty response from model."

        logger.info(f"🤖 Generated code:\n{text}")

        # --- تنفيذ الكود وإعادة النتيجة ---
        exec_result = execute_python_code(text)
        final_output = (
            f"🧠 **Prompt:** {prompt}\n\n"
            f"💻 **Generated Code:**\n{text}\n\n"
            f"🧾 **Execution Result:**\n{exec_result}"
        )
        return final_output

    except Exception as e:
        logger.exception("Generation/Execution error: %s", e)
        return f"❌ Error: {e}"


# ------------------------------------------------------------
# 🧱 واجهة Gradio
# ------------------------------------------------------------
iface = gr.Interface(
    fn=generate_and_execute_fn,
    inputs=gr.Textbox(lines=4, placeholder="💬 اكتب سؤالك البرمجي هنا..."),
    outputs=gr.Textbox(lines=15, label="💡 ناتج CodeAgent"),
    title="🤖 CodeAgent — Code Generator & Executor",
    description="وكيل يقوم بتوليد الكود البرمجي بلغة بايثون وتنفيذه مباشرة باستخدام نموذج Qwen2.5-Coder-14B-Instruct-GGUF."
)


# ------------------------------------------------------------
# 🧹 إغلاق آمن للنموذج
# ------------------------------------------------------------
def safe_close():
    global llm
    try:
        if llm is not None:
            llm.close()
            logger.info("🧩 Model closed successfully.")
    except Exception:
        logger.exception("Error closing model.")
    finally:
        llm = None


atexit.register(safe_close)


# ------------------------------------------------------------
# 🚀 الإطلاق المتوافق مع الوكيل الأساسي
# ------------------------------------------------------------
if __name__ == "__main__":
    port = int(os.environ.get("PORT", 7860))
    logger.info(f"🚀 Launching CodeAgent on port {port}")
    iface.launch(server_name="0.0.0.0", server_port=port)