import json
import gradio as gr

# تحميل البيانات
with open("dataset.json", "r", encoding="utf-8") as f:
    DATA = json.load(f)

CATEGORY_MAP = {entry["category"].lower(): entry for entry in DATA}

ALIASES = {
    "xss": "XSS Testing",
    "cross site scripting": "XSS Testing",
    "cross-site scripting": "XSS Testing",
    "lfi": "LFI Testing",
    "local file inclusion": "LFI Testing",
    "open redirect": "Open Redirect",
    "redirect": "Open Redirect",
    "ssrf": "SSRF Testing",
    "server side request forgery": "SSRF Testing",
    "git": "Git Repository Disclosure",
    "git disclosure": "Git Repository Disclosure",
    "git repository": "Git Repository Disclosure",
    "subdomain takeover": "Subdomain Takeover",
    "takeover": "Subdomain Takeover",
    "cors": "CORS Testing",
    "cross origin": "CORS Testing",
    "wordpress": "WordPress Security Testing",
    "wp": "WordPress Security Testing",
    "directory bruteforce": "Directory & File Bruteforcing",
    "dir bruteforce": "Directory & File Bruteforcing",
    "bruteforce": "Directory & File Bruteforcing",
    "hidden parameter": "Hidden Parameter Discovery",
    "parameter discovery": "Hidden Parameter Discovery",
    "arjun": "Hidden Parameter Discovery",
    "sensitive file": "Sensitive File Discovery",
    "file discovery": "Sensitive File Discovery",
    "vulnerability scanning": "Vulnerability Scanning",
    "nuclei": "Vulnerability Scanning",
    "scanning": "Vulnerability Scanning",
    "url collection": "URL Collection & Analysis",
    "url analysis": "URL Collection & Analysis",
    "gau": "URL Collection & Analysis",
    "live host": "Live Host Discovery",
    "host discovery": "Live Host Discovery",
    "httpx": "Live Host Discovery",
    "subdomain enumeration": "Subdomain Enumeration",
    "subdomain": "Subdomain Enumeration",
    "subfinder": "Subdomain Enumeration",
    "additional tools": "Additional Tools",
    "tools": "Additional Tools",
}

def find_category(query: str):
    q = query.strip().lower()
    for alias, cat_name in ALIASES.items():
        if alias in q:
            key = cat_name.lower()
            if key in CATEGORY_MAP:
                return CATEGORY_MAP[key], None
    for key, entry in CATEGORY_MAP.items():
        if key in q or q in key:
            return entry, None
    return None, (
        "❌ لم يتم التعرف على الثغرة. الثغرات المتاحة:\n\n"
        + "\n".join(f"• {e['category']}" for e in DATA)
    )

def wants_commands(query: str) -> bool:
    keywords = [
        "command", "commands", "أوامر", "امر", "أمر",
        "how", "كيف", "tool", "أداة", "run", "تشغيل",
        "exploit", "اختبار", "test",
    ]
    return any(k in query.lower() for k in keywords)

def format_description_only(entry) -> str:
    return (
        f"## 🛡️ {entry['category']}\n\n"
        f"{entry['description']}\n\n"
        f"---\n"
        f"💡 اكتب **`{entry['category']} commands`** لعرض الأوامر"
    )

def format_with_commands(entry) -> str:
    lines = [
        f"## 🛡️ {entry['category']}\n\n{entry['description']}\n\n---\n\n## ⚙️ الأوامر\n"
    ]
    for cmd in entry.get("commands", []):
        lines.append(f"### {cmd['id']}. {cmd['description']}\n")
        lines.append(f"```bash\n{cmd['command']}\n```\n")
    return "\n".join(lines)

def respond(message: str, history: list):
    if not message.strip():
        return history, ""

    entry, err = find_category(message)
    answer = err if err else (
        format_with_commands(entry) if wants_commands(message)
        else format_description_only(entry)
    )

    # بناء التاريخ بالصيغة الجديدة
    history = history or []
    history.append({"role": "user", "content": message})
    history.append({"role": "assistant", "content": answer})
    return history, ""

# واجهة Gradio
CATEGORIES_MD = "\n".join(f"- {e['category']}" for e in DATA)

with gr.Blocks(title="CyberSec Assistant") as demo:
    gr.Markdown(
        """
# 🔐 CyberSec Web Vulnerabilities Assistant
**نموذج مساعد الأمن السيبراني** - يشرح الثغرات الأمنية ويعرض أوامر الاختبار
        """
    )

    with gr.Row():
        with gr.Column(scale=3):
            chatbot = gr.Chatbot(label="المحادثة", height=500)
            with gr.Row():
                txt = gr.Textbox(
                    placeholder="اكتب اسم الثغرة مثل: XSS أو LFI أو SSRF ...",
                    label="سؤالك",
                    scale=4,
                )
                send_btn = gr.Button("إرسال", variant="primary", scale=1)

        with gr.Column(scale=1):
            gr.Markdown(f"### 📋 الثغرات المتاحة\n{CATEGORIES_MD}")
            gr.Markdown(
                """
### 💡 أمثلة
- `XSS` ← يعرض الوصف فقط
- `XSS commands` ← يعرض الأوامر
- `LFI Testing` ← وصف LFI
- `subdomain enumeration commands`
                """
            )

    # ربط الأحداث
    send_btn.click(respond, [txt, chatbot], [chatbot, txt])
    txt.submit(respond, [txt, chatbot], [chatbot, txt])

    gr.Markdown("---\n*تصنيف: هجومي | Web Vulnerabilities | v1.2*")

demo.launch()