import os
import json
import pickle
import time
import gradio as gr
from pathlib import Path

# ── LLM ──────────────────────────────────────────────────────────────────────
from llama_cpp import Llama

# ── Retrieval ─────────────────────────────────────────────────────────────────
import faiss
import numpy as np
from sentence_transformers import SentenceTransformer, CrossEncoder
from rank_bm25 import BM25Okapi

# ─────────────────────────────────────────────────────────────────────────────
# Config
# ─────────────────────────────────────────────────────────────────────────────
HF_MODEL_REPO   = "pushkarsharma/LegalSahayk_q4_k_m"
HF_MODEL_FILE   = "LegalSahyak_q4_k_m.gguf"
EMBED_MODEL     = "BAAI/bge-small-en-v1.5"
RERANKER_MODEL  = "BAAI/bge-reranker-large"
N_CTX           = 4096
N_GPU_LAYERS    = int(os.getenv("N_GPU_LAYERS", "0"))   # set to 20+ on GPU Space
TOP_K_RETRIEVE  = 10
TOP_K_RERANK    = 3
DEVICE          = os.getenv("DEVICE", "cpu")

# ─────────────────────────────────────────────────────────────────────────────
# Load everything once at startup
# ─────────────────────────────────────────────────────────────────────────────
print("⚙️  Loading embedding model …")
embedder = SentenceTransformer(EMBED_MODEL, device=DEVICE)

print("⚙️  Loading reranker …")
reranker = CrossEncoder(RERANKER_MODEL, device=DEVICE)

print("⚙️  Loading FAISS indices …")
def load_index(prefix):
    index    = faiss.read_index(f"{prefix}.faiss")
    with open(f"{prefix}_bm25.pkl", "rb") as f:
        bm25 = pickle.load(f)
    with open(f"{prefix}_meta.json") as f:
        meta = json.load(f)
    return index, bm25, meta

statute_index,  statute_bm25,  statute_meta  = load_index("db_statutes")
contract_index, contract_bm25, contract_meta = load_index("db_contract")

print("⚙️  Loading GGUF model from HuggingFace …")
llm = Llama.from_pretrained(
    repo_id      = HF_MODEL_REPO,
    filename     = HF_MODEL_FILE,
    n_ctx        = N_CTX,
    n_gpu_layers = N_GPU_LAYERS,
    verbose      = False,
)
print("✅  All components loaded.")

# ─────────────────────────────────────────────────────────────────────────────
# Retrieval helpers
# ─────────────────────────────────────────────────────────────────────────────
def hybrid_search(query: str, faiss_index, bm25, meta, top_k=TOP_K_RETRIEVE):
    """Dense + sparse retrieval, deduplicated."""
    q_emb = embedder.encode([query], normalize_embeddings=True).astype("float32")

    # Dense
    _, dense_ids = faiss_index.search(q_emb, top_k)
    dense_ids = dense_ids[0].tolist()

    # Sparse
    tokens    = query.lower().split()
    bm25_scores = bm25.get_scores(tokens)
    sparse_ids  = np.argsort(bm25_scores)[::-1][:top_k].tolist()

    # Merge unique
    seen, docs = set(), []
    for idx in dense_ids + sparse_ids:
        if idx not in seen and 0 <= idx < len(meta):
            seen.add(idx)
            docs.append(meta[idx])
    return docs[:top_k]

def rerank(query: str, docs: list, top_k=TOP_K_RERANK):
    """Cross-encoder reranking."""
    if not docs:
        return []
    pairs  = [(query, d.get("text", "")) for d in docs]
    scores = reranker.predict(pairs)
    ranked = sorted(zip(scores, docs), key=lambda x: x[0], reverse=True)
    return [d for _, d in ranked[:top_k]]

def search_statutes(query: str):
    docs = hybrid_search(query, statute_index, statute_bm25, statute_meta)
    return rerank(query, docs)

def search_contracts(query: str):
    docs = hybrid_search(query, contract_index, contract_bm25, contract_meta)
    return rerank(query, docs)

# ─────────────────────────────────────────────────────────────────────────────
# Agent / answer generation
# ─────────────────────────────────────────────────────────────────────────────
SYSTEM_PROMPT = (
    "You are LegalSahyak, an expert AI legal assistant specialising in Indian law. "
    "Answer only based on the retrieved context provided. "
    "If the context does not contain enough information, say so honestly. "
    "Always cite the statute name and section when available. "
    "Do NOT invent legal provisions."
)

def build_context(statute_docs, contract_docs):
    parts = []
    if statute_docs:
        parts.append("## Relevant Statutes")
        for i, d in enumerate(statute_docs, 1):
            title = d.get("act_title", d.get("title", "Unknown Act"))
            sec   = d.get("section_title", d.get("section", ""))
            text  = d.get("text", "")
            parts.append(f"[S{i}] {title} — {sec}\n{text}")
    if contract_docs:
        parts.append("## Relevant Contract Clauses")
        for i, d in enumerate(contract_docs, 1):
            title = d.get("title", f"Clause {i}")
            text  = d.get("text", "")
            parts.append(f"[C{i}] {title}\n{text}")
    return "\n\n".join(parts) if parts else "No relevant context found."

def answer_question(query: str, history: list):
    if not query.strip():
        return history, "", []

    # Retrieve from both corpora
    stat_docs = search_statutes(query)
    cont_docs = search_contracts(query)
    context   = build_context(stat_docs, cont_docs)

    messages = [
        {"role": "system",    "content": SYSTEM_PROMPT},
        {"role": "user",      "content": f"Context:\n{context}\n\nQuestion: {query}"},
    ]

    resp = llm.create_chat_completion(
        messages    = messages,
        max_tokens  = 512,
        temperature = 0.1,
    )
    answer = resp["choices"][0]["message"]["content"].strip()

    # Build source cards for display
    sources = []
    for d in stat_docs:
        sources.append({
            "type"   : "statute",
            "title"  : d.get("act_title", d.get("title", "Statute")),
            "section": d.get("section_title", d.get("section", "")),
            "snippet": d.get("text", "")[:200] + "…",
        })
    for d in cont_docs:
        sources.append({
            "type"   : "contract",
            "title"  : d.get("title", "Contract Clause"),
            "section": "",
            "snippet": d.get("text", "")[:200] + "…",
        })

    history = history + [[query, answer]]
    return history, "", sources

# ─────────────────────────────────────────────────────────────────────────────
# Gradio UI
# ─────────────────────────────────────────────────────────────────────────────
EXAMPLE_QUESTIONS = [
    "What is the penalty for GST evasion under Section 122?",
    "Explain the notice period clause in a standard employment contract.",
    "What are the rights of an employee under the Factories Act 1948?",
    "What does Section 138 of the Negotiable Instruments Act say about cheque bounce?",
    "What is the process for registering a startup under Indian law?",
]

CSS = """
@import url('https://fonts.googleapis.com/css2?family=Cormorant+Garamond:wght@400;600;700&family=DM+Sans:wght@300;400;500&display=swap');

:root {
    --bg:        #0d0f14;
    --surface:   #161a22;
    --border:    #252c3a;
    --gold:      #c9a84c;
    --gold-dim:  #8a6f2e;
    --text:      #e8e4dc;
    --muted:     #7a8090;
    --user-bg:   #1a2035;
    --bot-bg:    #131720;
    --tag-statute:  #1a3a2a;
    --tag-contract: #1a2a3a;
}

* { box-sizing: border-box; }

body, .gradio-container {
    background: var(--bg) !important;
    font-family: 'DM Sans', sans-serif !important;
    color: var(--text) !important;
}

/* Header */
.ls-header {
    text-align: center;
    padding: 2.5rem 1rem 1rem;
    border-bottom: 1px solid var(--border);
    margin-bottom: 1.5rem;
}
.ls-header h1 {
    font-family: 'Cormorant Garamond', serif;
    font-size: 2.8rem;
    font-weight: 700;
    color: var(--gold);
    letter-spacing: 0.02em;
    margin: 0;
}
.ls-header p {
    color: var(--muted);
    font-size: 0.9rem;
    margin: 0.4rem 0 0;
    font-weight: 300;
}
.ls-badge {
    display: inline-block;
    background: var(--gold-dim);
    color: var(--gold);
    font-size: 0.7rem;
    font-weight: 500;
    letter-spacing: 0.1em;
    text-transform: uppercase;
    padding: 2px 10px;
    border-radius: 20px;
    margin-top: 0.6rem;
}

/* Chatbot */
.ls-chat { border-radius: 12px !important; background: var(--surface) !important; border: 1px solid var(--border) !important; }
.ls-chat .message { font-family: 'DM Sans', sans-serif !important; font-size: 0.92rem !important; line-height: 1.7 !important; }
.ls-chat .user  { background: var(--user-bg) !important; border: 1px solid var(--border) !important; }
.ls-chat .bot   { background: var(--bot-bg) !important;  border: 1px solid var(--border) !important; }

/* Input row */
.ls-input textarea {
    background: var(--surface) !important;
    border: 1px solid var(--border) !important;
    color: var(--text) !important;
    font-family: 'DM Sans', sans-serif !important;
    border-radius: 10px !important;
    font-size: 0.92rem !important;
}
.ls-input textarea:focus { border-color: var(--gold-dim) !important; outline: none !important; }

.ls-send button {
    background: var(--gold) !important;
    color: #0d0f14 !important;
    font-weight: 600 !important;
    border-radius: 10px !important;
    border: none !important;
    font-family: 'DM Sans', sans-serif !important;
    letter-spacing: 0.03em;
}
.ls-send button:hover { background: #e0b955 !important; }

/* Source cards */
.ls-sources h3 {
    font-family: 'Cormorant Garamond', serif;
    font-size: 1.1rem;
    color: var(--gold);
    font-weight: 600;
    margin-bottom: 0.8rem;
    letter-spacing: 0.03em;
}
.source-card {
    background: var(--surface);
    border: 1px solid var(--border);
    border-radius: 10px;
    padding: 0.9rem 1rem;
    margin-bottom: 0.6rem;
    transition: border-color 0.2s;
}
.source-card:hover { border-color: var(--gold-dim); }
.source-card .tag {
    font-size: 0.65rem;
    font-weight: 600;
    letter-spacing: 0.12em;
    text-transform: uppercase;
    padding: 2px 8px;
    border-radius: 4px;
    display: inline-block;
    margin-bottom: 0.4rem;
}
.tag-statute  { background: var(--tag-statute);  color: #6bcf8e; }
.tag-contract { background: var(--tag-contract); color: #6baacf; }
.source-card .src-title   { font-weight: 500; font-size: 0.85rem; color: var(--text); }
.source-card .src-section { font-size: 0.78rem; color: var(--muted); margin-bottom: 0.3rem; }
.source-card .src-snippet { font-size: 0.78rem; color: var(--muted); line-height: 1.5; }

/* Examples */
.ls-examples { margin-top: 0.8rem; }
.ls-examples button {
    background: var(--surface) !important;
    border: 1px solid var(--border) !important;
    color: var(--muted) !important;
    font-size: 0.8rem !important;
    border-radius: 8px !important;
    transition: all 0.2s !important;
}
.ls-examples button:hover { border-color: var(--gold-dim) !important; color: var(--text) !important; }

/* Disclaimer */
.ls-disclaimer {
    text-align: center;
    font-size: 0.72rem;
    color: var(--muted);
    margin-top: 1.2rem;
    padding-top: 1rem;
    border-top: 1px solid var(--border);
    line-height: 1.6;
}
"""

def format_sources_html(sources):
    if not sources:
        return "<div class='ls-sources'><p style='color:var(--muted);font-size:0.82rem'>No sources retrieved.</p></div>"
    cards = ""
    for s in sources:
        tag_cls  = "tag-statute" if s["type"] == "statute" else "tag-contract"
        tag_label = "Statute" if s["type"] == "statute" else "Contract"
        sec_html = f"<div class='src-section'>{s['section']}</div>" if s["section"] else ""
        cards += f"""
        <div class='source-card'>
            <span class='tag {tag_cls}'>{tag_label}</span>
            <div class='src-title'>{s['title']}</div>
            {sec_html}
            <div class='src-snippet'>{s['snippet']}</div>
        </div>"""
    return f"<div class='ls-sources'><h3>📚 Sources Retrieved</h3>{cards}</div>"

def on_submit(query, history, sources_state):
    history, _, sources = answer_question(query, history)
    sources_html = format_sources_html(sources)
    return history, "", sources_html

def load_example(example):
    return example

with gr.Blocks(css=CSS, title="LegalSahyak") as demo:
    gr.HTML("""
    <div class='ls-header'>
        <h1>⚖ LegalSahyak</h1>
        <p>AI-powered Indian Legal Assistant · Grounded in Statutes & Contracts</p>
        <span class='ls-badge'>Llama 3.1 · Fine-tuned · RAG · Offline-capable</span>
    </div>
    """)

    with gr.Row():
        with gr.Column(scale=3):
            chatbot = gr.Chatbot(
                label="",
                height=480,
                elem_classes=["ls-chat"],
                show_label=False,
                avatar_images=(None, "https://huggingface.co/front/assets/huggingface_logo-noborder.svg"),
            )
            with gr.Row():
                query_box = gr.Textbox(
                    placeholder="Ask a legal question… e.g. What is Section 138 NI Act?",
                    show_label=False,
                    lines=2,
                    elem_classes=["ls-input"],
                    scale=5,
                )
                send_btn = gr.Button("Ask →", elem_classes=["ls-send"], scale=1)

            gr.Examples(
                examples=EXAMPLE_QUESTIONS,
                inputs=query_box,
                label="Try an example",
                elem_classes=["ls-examples"],
            )

        with gr.Column(scale=2):
            sources_box = gr.HTML(
                value="<div class='ls-sources'><p style='color:var(--muted);font-size:0.82rem;margin-top:1rem'>Sources will appear here after your first question.</p></div>"
            )

    sources_state = gr.State([])
    history_state = gr.State([])

    send_btn.click(
        fn=on_submit,
        inputs=[query_box, chatbot, sources_state],
        outputs=[chatbot, query_box, sources_box],
    )
    query_box.submit(
        fn=on_submit,
        inputs=[query_box, chatbot, sources_state],
        outputs=[chatbot, query_box, sources_box],
    )

    gr.HTML("""
    <div class='ls-disclaimer'>
        ⚠️ LegalSahyak is an AI tool for informational purposes only.<br>
        It is <strong>not a substitute for licensed legal counsel</strong>. Always consult a qualified advocate for legal decisions.
    </div>
    """)

demo.launch()