import gradio as gr
import cv2
import numpy as np
import insightface
import faiss
import sqlite3
import uuid
import json
from datetime import datetime

# ========== 1. تحميل نموذج ArcFace (IResNet100) ==========
print("Loading ArcFace model...")
model = insightface.app.FaceAnalysis(name='buffalo_l')
model.prepare(ctx_id=0)  # 0 للـ CPU
print("Model ready.")

def get_embedding(image):
    """استخراج embedding (512-dim) من الصورة"""
    faces = model.get(cv2.cvtColor(image, cv2.COLOR_RGB2BGR))
    if len(faces) == 0:
        return None
    return faces[0].embedding

# ========== 2. تهيئة FAISS (مؤشر متجهات) ==========
dim = 512
index = faiss.IndexFlatL2(dim)   # يمكن تبديلها بـ IndexFlatIP للـ cosine
id_map = []   # يربط فهرس FAISS مع identity_id

# ========== 3. قاعدة بيانات SQLite ==========
conn = sqlite3.connect('faces.db', check_same_thread=False)
c = conn.cursor()
c.execute('''
    CREATE TABLE IF NOT EXISTS identities (
        identity_id TEXT PRIMARY KEY,
        name TEXT NOT NULL,
        embedding_id INTEGER,
        created_at TIMESTAMP,
        metadata TEXT
    )
''')
conn.commit()

def add_identity(identity_id, name, embedding, metadata=None):
    """إضافة شخص جديد إلى النظام"""
    # إضافة إلى FAISS
    embedding_np = np.array([embedding]).astype('float32')
    idx = index.ntotal
    index.add(embedding_np)
    id_map.append(identity_id)
    
    # إضافة إلى SQLite
    c.execute('''
        INSERT OR REPLACE INTO identities (identity_id, name, embedding_id, created_at, metadata)
        VALUES (?, ?, ?, ?, ?)
    ''', (identity_id, name, idx, datetime.now(), metadata))
    conn.commit()
    return idx

def search_identity(embedding, k=5):
    """البحث عن أقرب k شخص"""
    if embedding is None:
        return []
    embedding_np = np.array([embedding]).astype('float32')
    distances, indices = index.search(embedding_np, k)
    results = []
    for dist, idx in zip(distances[0], indices[0]):
        if idx == -1:
            continue
        identity_id = id_map[idx]
        c.execute('SELECT name, metadata FROM identities WHERE identity_id=?', (identity_id,))
        row = c.fetchone()
        if row:
            name, metadata = row
            confidence = max(0, 1 - dist / 2)  # تحويل المسافة إلى ثقة
            results.append({
                'identity_id': identity_id,
                'name': name,
                'confidence': confidence,
                'metadata': metadata
            })
    return results

def list_all_identities():
    c.execute('SELECT identity_id, name, created_at FROM identities')
    rows = c.fetchall()
    if not rows:
        return "📭 لا توجد هويات مسجلة."
    out = "## 📋 قائمة الهويات:\n"
    for row in rows:
        out += f"- **{row[1]}** (ID: {row[0]}) – تم التسجيل: {row[2]}\n"
    return out

def delete_identity(identity_id):
    """حذف شخص (سيتم إعادة بناء الفهرس لاحقاً – مبسط)"""
    # للتبسيط، نقوم بحذف من SQLite فقط
    # الإعادة الكاملة للفهرس تتطلب إعادة بناء FAISS، يمكن تنفيذها لاحقاً
    c.execute('DELETE FROM identities WHERE identity_id=?', (identity_id,))
    conn.commit()
    return f"✅ تم حذف {identity_id} (ملاحظة: الفهرس لم يُحدث تلقائياً)"

# ========== 4. وظائف Gradio ==========
def enroll(image, identity_id, name, metadata_str):
    """تسجيل وجه جديد"""
    if image is None:
        return "⚠️ يرجى تحميل صورة."
    emb = get_embedding(image)
    if emb is None:
        return "❌ لم يتم اكتشاف وجه في الصورة."
    metadata = None
    if metadata_str.strip():
        try:
            metadata = json.loads(metadata_str)
        except:
            return "⚠️ البيانات الإضافية ليست JSON صحيح."
    add_identity(identity_id, name, emb, json.dumps(metadata) if metadata else None)
    return f"✅ تم تسجيل {name} (ID: {identity_id}) بنجاح."

def search(image, threshold):
    """البحث عن وجه"""
    if image is None:
        return "⚠️ يرجى تحميل صورة."
    emb = get_embedding(image)
    if emb is None:
        return "❌ لم يتم اكتشاف وجه في الصورة."
    results = search_identity(emb)
    if not results:
        return "⚠️ لم يتم العثور على تطابق."
    out = f"## 🔍 نتائج البحث (أعلى {len(results)}):\n"
    for res in results:
        out += f"- **{res['name']}** (ID: {res['identity_id']}) – ثقة: {res['confidence']:.2%}\n"
    return out

def filter_by_metadata(key, value):
    """مثال على فلترة حسب metadata (يمكن إضافتها إلى الواجهة)"""
    c.execute("SELECT identity_id, name FROM identities WHERE metadata LIKE ?", (f'%"{key}": "{value}"%',))
    rows = c.fetchall()
    if not rows:
        return "⚠️ لا توجد نتائج."
    out = "## 🔎 نتائج الفلترة:\n"
    for row in rows:
        out += f"- {row[1]} (ID: {row[0]})\n"
    return out

# ========== 5. بناء واجهة Gradio ==========
with gr.Blocks(title="Malath - Face Recognition System") as demo:
    gr.Markdown("# 🧠 Malath – نظام التعرف على الوجه\n### ArcFace + FAISS + SQLite")
    with gr.Tabs():
        with gr.TabItem("📝 تسجيل"):
            with gr.Row():
                enroll_img = gr.Image(type="numpy", label="صورة الوجه")
                with gr.Column():
                    enroll_id = gr.Textbox(label="رقم الهويه")
                    enroll_name = gr.Textbox(label="الاسم")
                    enroll_meta = gr.Textbox(label="بيانات إضافية (JSON)", placeholder='{"القسم": "IT"}')
            enroll_btn = gr.Button("تسجيل", variant="primary")
            enroll_out = gr.Textbox(label="النتيجة")
            enroll_btn.click(enroll, [enroll_img, enroll_id, enroll_name, enroll_meta], enroll_out)
        
        with gr.TabItem("🔍 بحث"):
            search_img = gr.Image(type="numpy", label="صورة الوجه")
            search_thresh = gr.Slider(minimum=0, maximum=1, value=0.7, label="عتبة الثقة")
            search_btn = gr.Button("بحث", variant="primary")
            search_out = gr.Textbox(label="النتيجة")
            search_btn.click(search, [search_img, search_thresh], search_out)
        
        with gr.TabItem("📋 قائمة"):
            list_btn = gr.Button("عرض الكل")
            list_out = gr.Markdown()
            list_btn.click(list_all_identities, outputs=list_out)
        
        # تبويب إضافي للفلترة (مثال)
        with gr.TabItem("🔎 فلترة"):
            filter_key = gr.Textbox(label="المفتاح (مثال: القسم)")
            filter_val = gr.Textbox(label="القيمة (مثال: IT)")
            filter_btn = gr.Button("فلترة")
            filter_out = gr.Markdown()
            filter_btn.click(filter_by_metadata, [filter_key, filter_val], filter_out)

demo.launch(server_name="0.0.0.0", server_port=7860)