import gradio as gr
import json
import pickle
import numpy as np
import faiss
import string
import re
import urllib.parse
from sentence_transformers import SentenceTransformer
from transformers import pipeline
from rank_bm25 import BM25Okapi

print("Loading data...")
with open("passages.json", "r") as f:
    passages = json.load(f)
with open("passage_meta.json", "r") as f:
    passage_meta = json.load(f)
with open("bm25.pkl", "rb") as f:
    bm25 = pickle.load(f)
faiss_index = faiss.read_index("faiss.index")
print("Data loaded!")

print("Loading models...")
embedder = SentenceTransformer(
    "sentence-transformers/msmarco-distilbert-base-v3", device="cpu"
)
qa_model = pipeline(
    task="question-answering",
    model="deepset/roberta-base-squad2",
    device=-1
)
print("All models loaded!")

def tokenize(text):
    text = text.lower()
    text = text.translate(str.maketrans("", "", string.punctuation))
    return [t for t in text.split() if len(t) > 2]

def get_indiankanoon_link(filename, meta):
    try:
        parts = filename.replace(".pdf", "").split("__")
        docid = parts[-1]
        if docid.isdigit():
            return f"https://indiankanoon.org/doc/{docid}/"
        pet = meta.get("pet", "")
        res = meta.get("res", "")
        if pet and pet != "Unknown" and res and res != "Unknown":
            query = urllib.parse.quote(f"{pet} vs {res}")
        else:
            query = urllib.parse.quote(filename[:40])
        return f"https://indiankanoon.org/search/?formInput={query}"
    except:
        return "https://indiankanoon.org"

def hybrid_retrieve(query, top_k=5):
    bm25_scores = bm25.get_scores(tokenize(query))
    bm25_max    = bm25_scores.max()
    if bm25_max > 0:
        bm25_scores = bm25_scores / bm25_max
    query_vec = embedder.encode([query]).astype("float32")
    faiss.normalize_L2(query_vec)
    dense_scores_raw, dense_indices = faiss_index.search(
        query_vec, len(passages)
    )
    dense_scores = np.zeros(len(passages))
    for rank, idx in enumerate(dense_indices[0]):
        if idx != -1:
            dense_scores[idx] = dense_scores_raw[0][rank]
    dense_max = dense_scores.max()
    if dense_max > 0:
        dense_scores = dense_scores / dense_max
    combined    = (0.4 * bm25_scores) + (0.6 * dense_scores)
    top_indices = combined.argsort()[::-1][:top_k]
    return [{
        "passage" : passages[idx],
        "score"   : float(combined[idx]),
        "metadata": passage_meta[idx]["metadata"],
        "filename": passage_meta[idx]["filename"],
    } for idx in top_indices]

def extract_answer(question, passages_list):
    all_answers = []
    for p in passages_list:
        try:
            results = qa_model(
                question=question, context=p["passage"],
                max_answer_len=100, top_k=5
            )
            if isinstance(results, dict):
                results = [results]
            for r in results:
                if r["score"] > 0.01:
                    all_answers.append({
                        "answer"  : r["answer"],
                        "score"   : r["score"],
                        "passage" : p["passage"],
                        "metadata": p["metadata"],
                        "filename": p["filename"],
                    })
        except:
            continue
    if not all_answers:
        return {
            "answer"  : "I could not find an answer in the available judgments.",
            "score"   : 0.0,
            "passage" : "",
            "metadata": {},
            "filename": ""
        }
    return max(all_answers, key=lambda x: x["score"])

def format_response(answer_dict):
    ans      = answer_dict["answer"]
    conf     = round(answer_dict["score"] * 100, 1)
    meta     = answer_dict["metadata"]
    filename = answer_dict["filename"]
    passage  = answer_dict["passage"]
    link     = get_indiankanoon_link(filename, meta)

    lines = []
    lines.append(ans)
    lines.append("")

    if meta and any(v not in ["Unknown","nan",""] for v in meta.values()):
        lines.append("---")
        lines.append("**Source case**")
        for key, label in [
            ("case_no","Case"),("pet","Petitioner"),
            ("res","Respondent"),("judgment_date","Date"),
            ("judgment_by","Judge")
        ]:
            v = meta.get(key,"")
            if v and v not in ["Unknown","nan"]:
                lines.append(f"**{label}:** {v}")
    else:
        lines.append(f"**Source:** {filename}")

    lines.append("")
    lines.append(f"[View full judgment on IndianKanoon]({link})")

    if passage:
        lines.append("")
        lines.append("---")
        lines.append("**Relevant excerpt**")
        short = passage[:800]+"..." if len(passage)>800 else passage
        lines.append(f"*{short}*")

    lines.append("")
    lines.append(f"*Confidence: {conf}%*")
    return "\n".join(lines)

def chat(message, history):
    if not message.strip():
        return history, ""
    history = history or []
    history.append((message, "Searching judgments..."))
    yield history, ""
    retrieved = hybrid_retrieve(message, top_k=5)
    answer    = extract_answer(message, retrieved)
    response  = format_response(answer)
    history[-1] = (message, response)
    yield history, ""

CSS = """
@import url(https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&display=swap);

* { font-family: Inter, sans-serif !important; box-sizing: border-box; }

body { background: #0f1117 !important; }

.gradio-container {
    max-width: 900px !important;
    margin: 0 auto !important;
    background: #0f1117 !important;
    padding: 0 16px !important;
}

.header-wrap {
    background: linear-gradient(135deg, #1a1f2e 0%, #161b27 100%);
    border: 1px solid #2d3348;
    border-radius: 16px;
    padding: 24px;
    margin: 16px 0;
    text-align: center;
}

.app-title {
    font-size: 28px !important;
    font-weight: 700 !important;
    background: linear-gradient(135deg, #c9a84c, #f0d080, #c9a84c);
    -webkit-background-clip: text !important;
    -webkit-text-fill-color: transparent !important;
    background-clip: text !important;
    margin: 0 0 6px !important;
    letter-spacing: -0.5px;
}

.app-subtitle {
    color: #8892a4 !important;
    font-size: 14px !important;
    margin: 0 0 16px !important;
}

.stats-row {
    display: flex;
    gap: 8px;
    justify-content: center;
    flex-wrap: wrap;
}

.stat-pill {
    background: #1e2436;
    border: 1px solid #2d3348;
    border-radius: 20px;
    padding: 5px 14px;
    font-size: 12px;
    color: #c9a84c;
    font-weight: 500;
}

#chatbot {
    background: #161b27 !important;
    border: 1px solid #2d3348 !important;
    border-radius: 16px !important;
}

#chatbot .message.user {
    background: #1e2d4a !important;
    border: 1px solid #2d4a7a !important;
    border-radius: 16px 16px 4px 16px !important;
    color: #e8eaf0 !important;
    padding: 12px 16px !important;
    max-width: 75% !important;
    margin-left: auto !important;
}

#chatbot .message.bot {
    background: #1a1f2e !important;
    border: 1px solid #2d3348 !important;
    border-radius: 16px 16px 16px 4px !important;
    color: #e8eaf0 !important;
    padding: 12px 16px !important;
    max-width: 88% !important;
}

#chatbot .message.bot a {
    color: #c9a84c !important;
    text-decoration: underline !important;
}

.input-wrap {
    background: #161b27;
    border: 1px solid #2d3348;
    border-radius: 14px;
    padding: 8px 8px 8px 16px;
    display: flex;
    align-items: center;
    gap: 8px;
    margin-top: 10px;
}

#msg-box textarea {
    background: transparent !important;
    border: none !important;
    color: #e8eaf0 !important;
    font-size: 15px !important;
    outline: none !important;
    box-shadow: none !important;
}

#msg-box textarea::placeholder { color: #4a5568 !important; }

#send-btn {
    background: linear-gradient(135deg, #c9a84c, #e8c96a) !important;
    color: #0f1117 !important;
    border: none !important;
    border-radius: 10px !important;
    font-weight: 600 !important;
    font-size: 14px !important;
    padding: 10px 20px !important;
    min-width: 80px !important;
}

#send-btn:hover {
    background: linear-gradient(135deg, #e8c96a, #f0d080) !important;
}

#clear-btn {
    background: #1a1f2e !important;
    border: 1px solid #2d3348 !important;
    border-radius: 10px !important;
    color: #8892a4 !important;
    font-size: 13px !important;
    padding: 8px 16px !important;
}

#clear-btn:hover {
    background: #1e2436 !important;
    color: #c9a84c !important;
}

.example-label {
    color: #8892a4;
    font-size: 12px;
    margin: 12px 0 6px;
    text-transform: uppercase;
    letter-spacing: 0.08em;
    font-weight: 500;
}

.example-btn button {
    background: #1a1f2e !important;
    border: 1px solid #2d3348 !important;
    border-radius: 20px !important;
    font-size: 12px !important;
    color: #c9a84c !important;
    padding: 6px 14px !important;
    font-weight: 500 !important;
    transition: all 0.2s !important;
}

.example-btn button:hover {
    background: #1e2d4a !important;
    border-color: #c9a84c !important;
}

.disclaimer {
    text-align: center;
    font-size: 11px;
    color: #4a5568;
    margin-top: 12px;
    padding-bottom: 16px;
}

footer { display: none !important; }
"""

EXAMPLES = [
    "What is the punishment for murder?",
    "What are the grounds for bail?",
    "What is habeas corpus?",
    "What is the burden of proof?",
    "What is anticipatory bail?",
    "What is contempt of court?",
    "What is res judicata?",
    "What is the right to legal aid?",
]

INITIAL = [(
    None,
    "Namaste! I am **LexBot** — your Indian Supreme Court legal research assistant.\n\n"
    "I am trained on **1000 Indian Supreme Court judgments** and can answer questions "
    "about Indian law, legal principles, and court procedures.\n\n"
    "Each answer includes a **direct link to the full judgment** on IndianKanoon.org.\n\n"
    "Ask me anything about Indian law!"
)]

with gr.Blocks(css=CSS, title="LexBot - Indian Legal Assistant") as demo:

    gr.HTML("""
    <div class="header-wrap">
        <div style="font-size:36px;margin-bottom:8px">⚖️</div>
        <div class="app-title">LexBot</div>
        <div class="app-subtitle">
            Indian Supreme Court Legal Research Assistant
        </div>
        <div class="stats-row">
            <span class="stat-pill">MRR@5: 86.61%</span>
            <span class="stat-pill">F1: 8.73%</span>
            <span class="stat-pill">Response: 96.67%</span>
            <span class="stat-pill">1000 judgments</span>
            <span class="stat-pill">RoBERTa + FAISS</span>
        </div>
    </div>
    """)

    chatbot = gr.Chatbot(
        value=INITIAL,
        elem_id="chatbot",
        height=500,
        show_label=False,
        bubble_full_width=False,
        show_copy_button=True,
    )

    with gr.Row(elem_classes="input-wrap"):
        msg_input = gr.Textbox(
            placeholder="Ask a legal question e.g. What is habeas corpus?",
            show_label=False,
            scale=9,
            container=False,
            lines=1,
            elem_id="msg-box",
        )
        send_btn = gr.Button(
            "Send",
            elem_id="send-btn",
            scale=1,
            min_width=80
        )

    with gr.Row():
        clear_btn = gr.Button(
            "Clear conversation",
            elem_id="clear-btn",
            scale=1
        )

    gr.HTML("<div class=\'example-label\'>Try these questions</div>")

    with gr.Row():
        for q in EXAMPLES[:4]:
            gr.Button(
                q, elem_classes="example-btn", size="sm"
            ).click(fn=lambda x=q: x, outputs=msg_input)

    with gr.Row():
        for q in EXAMPLES[4:]:
            gr.Button(
                q, elem_classes="example-btn", size="sm"
            ).click(fn=lambda x=q: x, outputs=msg_input)

    gr.HTML("""
    <div class="disclaimer">
        Answers extracted directly from Supreme Court judgment texts.
        Click IndianKanoon links to read full judgments.
        Always verify legal information with a qualified advocate.
    </div>
    """)

    msg_input.submit(
        fn=chat, inputs=[msg_input, chatbot], outputs=[chatbot, msg_input]
    )
    send_btn.click(
        fn=chat, inputs=[msg_input, chatbot], outputs=[chatbot, msg_input]
    )
    clear_btn.click(
        fn=lambda: (INITIAL, ""), outputs=[chatbot, msg_input]
    )

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