import subprocess
import time
import requests
import json
import gradio as gr
from huggingface_hub import hf_hub_download

# =========================
# 1. DOWNLOAD MODEL
# =========================
MODEL_PATH = hf_hub_download(
    repo_id="unsloth/Qwen3-0.6B-GGUF",
    filename="Qwen3-0.6B-UD-IQ2_XXS.gguf"
)
print("Model downloaded:", MODEL_PATH)

# =========================
# 2. START llama.cpp SERVER
# =========================
server = subprocess.Popen([
    "./llama/llama-server",
    "-m", MODEL_PATH,
    "-c", "3096",
    "-t", "4",
    "--host", "0.0.0.0",
    "--port", "8080"
       # disable prompt cache — prevents slowdown on repeated runs
])

print("Starting llama-server, waiting 8 seconds...")
time.sleep(8)

# =========================
# 3. SYSTEM PROMPT — compact (~150 tokens)
# =========================
SYSTEM_PROMPT = """أنت مساعد قانوني جزائري خبير. مهمتك هي تحديد التكييف القانوني (الجريمة) والنص القانوني (المادة) بناءً على الوقائع (الفعل) التي يقدمها المستخدم.

"""

# =========================
# 4. GRADIO CHAT FUNCTION — streaming
# =========================
def chat(user_fact, history):
    payload = {
        "messages": [
            {"role": "system", "content": SYSTEM_PROMPT},
             {"role": "user", "content": "الفعل: قيام الجاني بضرب الشاكي بواسطة عصى خشبية على مستوى الرأس، مسبباً له عجزاً عن العمل لمدة 20 يوماً."},
    {"role": "assistant", "content": "الجريمة: جنحة الضرب والجرح العمدي باستعمال سلاح.\nالنص القانوني: المادة 266 من قانون العقوبات."},
             {"role": "user", "content": "الفعل: قيام الجاني بضرب الشاكي بواسطة عصى خشبية على مستوى الرأس، مسبباً له عجزاً عن العمل لمدة 20 يوماً."},
    {"role": "assistant", "content": "الجريمة: جنحة الضرب والجرح العمدي باستعمال سلاح.\nالنص القانوني: المادة 266 من قانون العقوبات."},
    
    # 2
    {"role": "user", "content": "الفعل: اعتداء الابن على والده بالضرب المبرح إثر مناوشات شفوية داخل مسكنهما العائلي."},
    {"role": "assistant", "content": "الجريمة: جنحة ضرب الأصول.\nالنص القانوني: المادة 267 من قانون العقوبات."},
    
    # 3
    {"role": "user", "content": "الفعل: اعتداء زوج على زوجته بالضرب المبرح مستعملاً حزاماً جلدياً مما سبب لها عجزاً قدره 10 أيام."},
    {"role": "assistant", "content": "الجريمة: جنحة الضرب والجرح العمدي بين الأزواج.\nالنص القانوني: المادة 266 مكرر 1 من قانون العقوبات."},
    
    # 4
    {"role": "user", "content": "الفعل: تهديد شخص لآخر بالقتل شفهياً بم مفتاح براغي في رقبته أمام مرأى الناس بالمحطة."},
    {"role": "assistant", "content": "الجريمة: جنحة التهديد بالقتل.\nالنص القانوني: المادة 284 من قانون العقوبات."},
    
    # 5
    {"role": "user", "content": "الفعل: قيام أشخاص باختطاف طفل قاصر واحتجازه في مستودع لطلب فدية مالية من أهله."},
    {"role": "assistant", "content": "الجريمة: جناية الاختطاف بدافع طلب فدية.\nالنص القانوني: المادة 293 bis من قانون العقوبات."},
    
    # 6
    {"role": "user", "content": "الفعل: اتهام سليم لزميله بالسرقة أمام الشرطة مع علمه التام ببراءته بهدف تشويه سمعته."},
    {"role": "assistant", "content": "الجريمة: جنحة الوشاية الكاذبة.\nالنص القانوني: المادة 300 من قانون العقوبات."},
    
    # 7
    {"role": "user", "content": "الفعل: تسلل شخص ليلاً عبر تسلق جدران مستودع وسرقة كابلات نحاسية والفرار بها."},
    {"role": "assistant", "content": "الجريمة: جناية السرقة المقترنة بظرفي الليل والتسلق.\nالنص القانوني: المادة 354 من قانون العقوبات."},
    
    # 8
    {"role": "user", "content": "الفعل: قيام الوكيل العقاري بالاستيلاء على مبالغ تسبيقات الكراء التي سلمت له لتوصيلها للمالك."},
    {"role": "assistant", "content": "الجريمة: جنحة خيانة الأمانة.\nالنص القانوني: المادة 376 من قانون العقوبات."},
    
    # 9
    {"role": "user", "content": "الفعل: إنشاء صفحة تواصل وهمية تدعي بيع سيارات مستوردة، وطلب تسبيقات مالية من الضحايا ثم الاختفاء."},
    {"role": "assistant", "content": "الجريمة: جنحة النصب.\nالنص القانوني: المادة 372 من قانون العقوبات."},
    
    # 10
    {"role": "user", "content": "الفعل: قيام شقيقين بالاستيلاء على العقارات والأراضي الفلاحية المتروكة من والدهما المتوفى وحرمان شقيقاتهم من الإرث."},
    {"role": "assistant", "content": "الجريمة: جنحة الاستيلاء على أموال تركة قبل القسمة.\nالنص القانوني: المادة 363 من قانون العقوبات."},
    
    # 11
    {"role": "user", "content": "الفعل: قيام شخص برشق واجهة سيارة جاره بالحجارة عمداً مما أدى لتكسير زجاجها الأمامي."},
    {"role": "assistant", "content": "الجريمة: جنحة التحطيم العمدي لملك الغير.\nالنص القانوني: المادة 407 من قانون العقوبات."},
    
    # 12
    {"role": "user", "content": "الفعل: تسلل شخص خلسة إلى أرض فلاحية مملوكة للغير وحرثها وزراعتها والادعاء بملكيتها بدون سند قانوني."},
    {"role": "assistant", "content": "الجريمة: جنحة التعدي على الملكية العقارية.\nالنص القانوني: المادة 386 من قانون العقوبات."},
    
    # 13
    {"role": "user", "content": "الفعل: طلب موظف ببلدية مبلغ مالي من مواطن مقابل استخراج وثيقة إدارية بسرعة."},
    {"role": "assistant", "content": "الجريمة: جنحة الرشوة (طلب مزية غير مستحقة).\nالنص القانوني: المادة 25 من القانون 06-01 المتعلق بالوقاية من الفساد ومكافحته."},
    
    # 14
    {"role": "user", "content": "الفعل: امتناع محكوم عليه عن تنفيذ حكم نهائي بالهدم وتهديد القوة العمومية الحاضرة للتنفيذ."},
    {"role": "assistant", "content": "الجريمة: جنحة التقليل من شأن الأحكام القضائية وجنحة إهانة هيئة نظامية.\nالنص القانوني: المواد 144 و147 من قانون العقوبات."},
    
    # 15
    {"role": "user", "content": "الفعل: قيام شخص بالتعدي بالسب والشتم على شرطي أثناء تنظيمه لحركة المرور في الطريق العام."},
    {"role": "assistant", "content": "الجريمة: جنحة إهانة موظف أثناء تأدية مهامه.\nالنص القانوني: المادة 144 من قانون العقوبات."},
    
    # 16
    {"role": "user", "content": "الفعل: ضبط شخص في الحدود البرية ينقل كمية معتبرة من الوقود على متن مركبة مهيأة بمخازن سرية بدون فواتير أو رخصة."},
    {"role": "assistant", "content": "الجريمة: جنحة التهريب باستعمال وسيلة نقل مهيأة.\nالنص القانوني: المادتين 10 و12 من الأمر 05-06 المتعلق بمكافحة التهريب."},
    
    # 17
    {"role": "user", "content": "الفعل: العثور بحوزة شخص على شريط من المؤثرات العقلية بريجابالين (300 ملغ) بغرض استهلاكه الشخصي دون وصفة طبية."},
    {"role": "assistant", "content": "الجريمة: جنحة حيازة مؤثرات عقلية بطريقة غير مشروعة لغرض الاستهلاك الشخصي.\nالنص القانوني: المادة 12 من القانون 04-18 المتعلق بالوقاية من المخدرات والمؤثرات العقلية."},
    
    # 18
    {"role": "user", "content": "الفعل: ضبط شخص يبيع أقراصاً مهلوسة للشباب أمام مؤسسة تعليمية."},
    {"role": "assistant", "content": "الجريمة: جنحة حيازة مؤثرات عقلية لغرض العرض والبيع على الغير مع توافر ظرف تشديد (مؤسسة تعليمية).\nالنص القانوني: المادتين 13 و17 من القانون 04-18 المتعلق بالوقاية من المخدرات والمؤثرات العقلية."},
    
    # 19
    {"role": "user", "content": "الفعل: قيام تاجر بعرض سلع للبيع في محله دون تعليق لوحة الأسعار وإعلام المستهلك بها."},
    {"role": "assistant", "content": "الجريمة: جنحة عدم الإعلام بالأسعار.\nالنص القانوني: المادة 31 من القانون 04-02 المتعلق بالقواعد المطبقة على الممارسات التجارية."},
    
    # 20
    {"role": "user", "content": "الفعل: عثور مصالح الرقابة على أطعمة فاسدة ومتسخة داخل مطبخ مطعم يقدم وجبات للجمهور."},
    {"role": "assistant", "content": "الجريمة: جنحة عدم احترام إلزامية النظافة والنظافة الصحية للمواد الغذائية.\nالنص القانوني: المادة 72 من القانون 09-03 المتعلق بحماية المستهلك وقمع الغش."},
   
   
  
            {"role": "user",   "content": f" الفعل{user_fact}"}
        ],
        "stream": True,  # ← ADD THIS — you were missing the stream flag in payload
    }
    try:
        r = requests.post(
            "http://127.0.0.1:8080/v1/chat/completions",
            json=payload,
            stream=True,
            timeout=400
        )
        r.raise_for_status()
        result = ""
        yielded = False  # ← track whether we've yielded anything

        for raw in r.iter_lines():
            if not raw:
                continue
            line = raw if isinstance(raw, str) else raw.decode("utf-8")
            if not line.startswith("data: "):
                continue
            data = line[6:]
            if data.strip() == "[DONE]":
                break
            try:
                chunk = json.loads(data)
                delta = chunk["choices"][0]["delta"].get("content") or ""
                if delta:
                    result += delta
                    yielded = True
                    yield result
            except (KeyError, json.JSONDecodeError):
                continue

        if not yielded:  # ← always yield something
            yield "لم يتم استلام أي رد من النموذج."

    except Exception as e:
        yield f"حدث خطأ أثناء الاتصال بالخادم الداخلي: {str(e)}"

# =========================
# 5. GRADIO UI
# =========================
demo = gr.ChatInterface(
    fn=chat,
    title="مساعد تكييف الجرائم - قانون العقوبات الجزائري",
    description="أدخل تفاصيل الواقعة أو الفعل، وسيقوم النظام بتحديد الوصف القانوني للجريمة والمادة المعاقب عليها.",
)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)