import os
import gradio as gr
import torch
from sentence_transformers import CrossEncoder

# ──────────────────────────────────────────────────────────────────────────────
# Model
# ──────────────────────────────────────────────────────────────────────────────
MODEL_ID = os.environ.get("MODEL_ID", "ALJIACHI/Mizan-Rerank-v2")

model = CrossEncoder(
    MODEL_ID,
    max_length=8192,
    trust_remote_code=True,
    device="cuda" if torch.cuda.is_available() else "cpu",
)

# ──────────────────────────────────────────────────────────────────────────────
# Reranking logic
# ──────────────────────────────────────────────────────────────────────────────
def rerank(query: str, documents: str) -> str:
    """Score and rerank documents against a query."""
    if not query.strip():
        return "⚠️ يرجى إدخال استعلام (Please enter a query)"

    docs = [d.strip() for d in documents.strip().split("\n") if d.strip()]
    if not docs:
        return "⚠️ يرجى إدخال مستند واحد على الأقل (Please enter at least one document)"

    ranks = model.rank(query, docs)

    lines = []
    for i, entry in enumerate(ranks, 1):
        idx = entry["corpus_id"]
        score = entry["score"]
        doc_text = docs[idx]
        preview = doc_text[:200] + ("..." if len(doc_text) > 200 else "")
        score = score if score == score else 0.0  # NaN safety
        bar_len = int(score * 20)
        bar = "█" * bar_len + "░" * (20 - bar_len)
        lines.append(
            f"### #{i}  —  Score: {score:.4f}\n"
            f"`{bar}`\n\n"
            f"{preview}\n"
        )

    header = f"## 📊 Results — {len(docs)} documents reranked\n---\n"
    return header + "\n---\n".join(lines)


# ──────────────────────────────────────────────────────────────────────────────
# Examples
# ──────────────────────────────────────────────────────────────────────────────
EXAMPLES = [
    [
        "ما هو تفسير الآية وجعلنا من الماء كل شيء حي",
        "تعني الآية أن الماء هو عنصر أساسي في حياة جميع الكائنات الحية، وهو ضروري لاستمرار الحياة.\n"
        "تم اكتشاف كواكب خارج المجموعة الشمسية تحتوي على مياه متجمدة.\n"
        "تحدث القرآن الكريم عن البرق والرعد في عدة مواضع مختلفة.",
    ],
    [
        "ما هي فوائد فيتامين د؟",
        "يساعد فيتامين د في تعزيز صحة العظام وتقوية الجهاز المناعي، كما يلعب دوراً مهماً في امتصاص الكالسيوم.\n"
        "يستخدم فيتامين د في بعض الصناعات الغذائية كمادة حافظة.\n"
        "أطلقت وزارة الزراعة حملة وطنية لزيادة الوعي بأهمية الزراعة العضوية.",
    ],
    [
        "ما حكم الصلاة في الإسلام؟",
        "الصلاة هي الركن الثاني من أركان الإسلام وهي واجبة على كل مسلم بالغ عاقل خمس مرات في اليوم والليلة.\n"
        "يُستحب للمسلم أن يصلي صلاة النوافل لزيادة الأجر والثواب.\n"
        "تأسست الجامعة الإسلامية في المدينة المنورة عام 1961 ميلادي.\n"
        "يجب على المسلم الطهارة قبل أداء الصلاة وهي شرط من شروط صحتها.",
    ],
]


# ──────────────────────────────────────────────────────────────────────────────
# Gradio UI
# ──────────────────────────────────────────────────────────────────────────────
CSS = """
.gradio-container { max-width: 900px !important; }
.header-text { text-align: center; }
.rtl-text textarea { direction: rtl; text-align: right; font-size: 16px; }
"""

with gr.Blocks(css=CSS, theme=gr.themes.Soft(), title="Mizan-Rerank-v2 Demo") as demo:
    gr.HTML(
        """
        <div class="header-text">
            <h1>🔍 Mizan-Rerank-v2</h1>
            <p style="font-size:16px; color:#555;">
                Arabic Long-Context Reranking Model · 305M Parameters · Up to 8192 Tokens
            </p>
            <p style="font-size:14px;">
                <a href="https://huggingface.co/ALJIACHI/Mizan-Rerank-v2" target="_blank">Model Card</a> ·
            </p>
        </div>
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            query_input = gr.Textbox(
                label="🔎 Query / الاستعلام",
                placeholder="أدخل الاستعلام هنا...",
                lines=2,
                elem_classes=["rtl-text"],
            )
            docs_input = gr.Textbox(
                label="📄 Documents (one per line) / المستندات (سطر لكل مستند)",
                placeholder="أدخل كل مستند في سطر منفصل...",
                lines=8,
                elem_classes=["rtl-text"],
            )
            rerank_btn = gr.Button("⚡ Rerank / إعادة الترتيب", variant="primary", size="lg")

        with gr.Column(scale=1):
            output = gr.Markdown(label="Results / النتائج")

    gr.Examples(
        examples=EXAMPLES,
        inputs=[query_input, docs_input],
        label="📌 Try an Example / جرب مثالاً",
    )

    rerank_btn.click(fn=rerank, inputs=[query_input, docs_input], outputs=output)
    query_input.submit(fn=rerank, inputs=[query_input, docs_input], outputs=output)

demo.launch()