import gradio as gr
from llama_cpp import Llama
from huggingface_hub import hf_hub_download

# ----------------------------------------------------------------------
# 1. إعداد النموذج (سحب من مستودع عام لتجاوز حدود التخزين)
# ----------------------------------------------------------------------

# نستخدم نموذج Qwen2.5-3B القوي والصغير (متوفر مسبقاً ولا يحتاج رفع)
REPO_ID = "Qwen/Qwen2.5-3B-Instruct-GGUF"
FILENAME = "qwen2.5-3b-instruct-q4_k_m.gguf"

def load_model():
    print(f"Downloading model {FILENAME} from {REPO_ID}...")
    try:
        # تحميل النموذج إلى الذاكرة المؤقتة للسيرفر
        model_path = hf_hub_download(
            repo_id=REPO_ID, 
            filename=FILENAME
        )
        print(f"Model downloaded to: {model_path}")
        
        # تشغيل النموذج (CPU)
        llm = Llama(
            model_path=model_path,
            n_ctx=2048,
            n_threads=2, 
            verbose=True
        )
        return llm
    except Exception as e:
        print(f"Error loading model: {e}")
        return None

# تحميل النموذج عند الإقلاع
ZEROCYBER_MODEL = load_model()

# ----------------------------------------------------------------------
# 2. منطق التحليل
# ----------------------------------------------------------------------

def generate_response(prompt_text, file_obj):
    if ZEROCYBER_MODEL is None:
        return "❌ Error: Model failed to load."

    # قراءة الملف إذا وجد
    context = ""
    if file_obj:
        try:
            with open(file_obj.name, 'r', encoding='utf-8', errors='ignore') as f:
                content = f.read()[:2000]
            context = f"\n\nFile Content to Analyze:\n{content}\n"
        except Exception as e:
            return f"Error reading file: {e}"

    # تجهيز البرومبت
    # نماذج Qwen/Mistral تفضل هذا التنسيق
    full_prompt = f"<|im_start|>system\nYou are a Cybersecurity Analyst.<|im_end|>\n<|im_start|>user\n{prompt_text}\n{context}<|im_end|>\n<|im_start|>assistant\n"

    output = ZEROCYBER_MODEL(
        full_prompt,
        max_tokens=512,
        stop=["<|im_end|>"],
        echo=False
    )
    
    return output['choices'][0]['text']

# ----------------------------------------------------------------------
# 3. واجهة المستخدم
# ----------------------------------------------------------------------

if __name__ == "__main__":
    interface = gr.Interface(
        fn=generate_response,
        inputs=[
            gr.Textbox(label="1. استفسارك الأمني:", placeholder="كيف أقوم بتأمين قاعدة البيانات؟"),
            gr.File(label="2. تحليل ملف (Log/Code)")
        ],
        outputs=gr.Textbox(label="التقرير الأمني"),
        title="🛡️ ZeroCyber Cloud Platform",
        description="تطبيق سحابي للتحليل الأمني يعمل بنموذج Qwen2.5-3B (GGUF).",
        allow_flagging="never"
    )
    
    interface.launch()