import os, json, re, time
import numpy as np
import gradio as gr
from sentence_transformers import SentenceTransformer, CrossEncoder
from rank_bm25 import BM25Okapi
from pinecone import Pinecone
from huggingface_hub import InferenceClient

# CONFIG
PINECONE_API_KEY = os.getenv("PINECONE_API_KEY")
HF_TOKEN         = os.getenv("HF_TOKEN")
PINECONE_INDEX   = "rag-nlp-project"
LLM_MODEL        = "meta-llama/Meta-Llama-3-8B-Instruct"

# ── LOAD RESOURCES ──
print("Loading resources...")

with open("chunks_recursive.json") as f:
    ALL_CHUNKS = json.load(f)

tokenized = [c["text"].lower().split() for c in ALL_CHUNKS]
bm25 = BM25Okapi(tokenized)

embedder = SentenceTransformer("all-MiniLM-L6-v2")
reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")

pc = Pinecone(api_key=PINECONE_API_KEY)
pine_index = pc.Index(PINECONE_INDEX)

llm = InferenceClient(token=HF_TOKEN)
print("All ready!")

# ── RETRIEVAL ──
def semantic_search(query, namespace="recursive", top_k=20):
    qvec = embedder.encode(query).tolist()
    res = pine_index.query(vector=qvec, top_k=top_k,
                           include_metadata=True, namespace=namespace)
    return [{"id": m["id"], "text": m["metadata"]["text"],
             "title": m["metadata"]["title"], "score": m["score"]}
            for m in res["matches"]]

def bm25_search(query, top_k=20):
    tokens = query.lower().split()
    scores = bm25.get_scores(tokens)
    top_idx = np.argsort(scores)[::-1][:top_k]
    return [{"id": ALL_CHUNKS[i]["id"], "text": ALL_CHUNKS[i]["text"],
             "title": ALL_CHUNKS[i]["title"], "score": float(scores[i])}
            for i in top_idx if scores[i] > 0]

def rrf_fuse(lists_of_results, k=60):
    scores, data = {}, {}
    for results in lists_of_results:
        for rank, item in enumerate(results):
            did = item["id"]
            scores[did] = scores.get(did, 0) + 1.0 / (k + rank + 1)
            data[did] = {"text": item["text"], "title": item["title"]}
    ranked = sorted(scores, key=lambda x: scores[x], reverse=True)
    return [{"id": d, "rrf_score": scores[d], **data[d]} for d in ranked]

def cross_encoder_rerank(query, candidates, top_k=5):
    if not candidates:
        return []
    pool = candidates[:30]
    pairs = [(query, c["text"]) for c in pool]
    ce_scores = reranker.predict(pairs)
    for i, s in enumerate(ce_scores):
        pool[i]["ce_score"] = float(s)
    pool.sort(key=lambda x: x["ce_score"], reverse=True)
    return pool[:top_k]

# ── LLM ──
def call_llm(prompt, max_tokens=512, temperature=0.3):
    for model in ["mistralai/Mistral-7B-Instruct-v0.2", "meta-llama/Meta-Llama-3-8B-Instruct"]:
        try:
            resp = llm.chat_completion(
                model=model,
                messages=[{"role": "user", "content": prompt}],
                max_tokens=max_tokens, temperature=temperature
            )
            return resp.choices[0].message.content.strip()
        except:
            continue
    return "[LLM Error: All models failed]"

def generate_answer(query, contexts):
    ctx = "\n\n".join([f"{i+1}. [{c['title']}] {c['text']}" for i, c in enumerate(contexts)])
    prompt = f"""Based on the following information:

{ctx}

Please provide a detailed answer to the question: {query}.
Your answer should integrate the diverse perspectives or data points provided by the retrieved passages.
If the passages are irrelevant to the question, say that you couldn't find a good response in the database."""
    return call_llm(prompt)

# ── EVALUATION ──
def eval_faithfulness(answer, contexts):
    context_str = "\n".join([c["text"] for c in contexts])[:3000]
    claims_raw = call_llm(
        f"Extract all factual claims as a numbered list.\n\nAnswer: {answer}\n\nClaims:",
        max_tokens=400, temperature=0.1
    )
    claims = [re.sub(r"^[\d]+[\.\)]\s*", "", l.strip())
              for l in claims_raw.split("\n")
              if len(re.sub(r"^[\d]+[\.\)]\s*", "", l.strip())) > 15]
    if not claims:
        return 1.0, "No claims extracted."

    supported = 0
    details = []
    for claim in claims[:8]:
        verdict = call_llm(
            f"Is this claim supported by the context? Reply ONLY 'SUPPORTED' or 'NOT SUPPORTED'.\n\n"
            f"Context: {context_str}\n\nClaim: {claim}\n\nVerdict:",
            max_tokens=10, temperature=0.1
        ).upper()
        ok = "SUPPORTED" in verdict and "NOT" not in verdict
        if ok:
            supported += 1
        details.append(f"{'[Y]' if ok else '[N]'} {claim}")
    score = supported / len(claims[:8])
    return score, "\n".join(details)

def eval_relevancy(query, answer):
    qs_raw = call_llm(
        f"Generate exactly 3 questions that this answer directly addresses. "
        f"One per line, no numbering.\n\nAnswer: {answer}\n\nQuestions:",
        max_tokens=200, temperature=0.3
    )
    questions = [re.sub(r"^[\d]+[\.\)]\s*", "", l.strip())
                 for l in qs_raw.split("\n")
                 if len(re.sub(r"^[\d]+[\.\)]\s*", "", l.strip())) > 10][:3]
    if not questions:
        return 0.0, "Could not generate questions."

    embs = embedder.encode([query] + questions)
    q_emb = embs[0]
    sims, detail_lines = [], []
    for i, q in enumerate(questions):
        sim = float(np.dot(q_emb, embs[i+1]) /
                    (np.linalg.norm(q_emb) * np.linalg.norm(embs[i+1])))
        sims.append(sim)
        detail_lines.append(f"  Q{i+1}: {q}  (sim={sim:.3f})")
    return float(np.mean(sims)), "\n".join(detail_lines)

# ── MAIN PIPELINE ──
def run_query(query, run_eval):
    if not query.strip():
        return "Please enter a question.", "", "", ""

    t0 = time.time()
    sem = semantic_search(query)
    kw = bm25_search(query)
    fused = rrf_fuse([sem, kw])
    reranked = cross_encoder_rerank(query, fused)
    t_retrieve = time.time() - t0

    t1 = time.time()
    answer = generate_answer(query, reranked)
    t_generate = time.time() - t1

    ctx_display = ""
    for i, c in enumerate(reranked):
        ctx_display += f"**[{i+1}] {c['title']}**  (score: {c.get('ce_score', 0):.3f})\n"
        ctx_display += f"{c['text']}\n\n---\n\n"

    scores_display = ""
    t_eval = 0
    if run_eval:
        t2 = time.time()
        faith_score, faith_detail = eval_faithfulness(answer, reranked)
        rel_score, rel_detail = eval_relevancy(query, answer)
        t_eval = time.time() - t2
        scores_display = (
            f"### Faithfulness: {faith_score:.0%}\n{faith_detail}\n\n"
            f"### Relevancy: {rel_score:.0%}\n{rel_detail}"
        )
    else:
        scores_display = "*(Check the box to run evaluation)*"

    timing = (f"Retrieval: {t_retrieve:.2f}s | Generation: {t_generate:.2f}s | "
              f"Evaluation: {t_eval:.2f}s | Total: {t_retrieve + t_generate + t_eval:.2f}s")
    return answer, ctx_display, scores_display, timing

# ── GRADIO UI ──
with gr.Blocks(title="RAG Q&A — AI/ML Domain", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        "# RAG Question-Answering System\n"
        "*AI/ML Domain - Hybrid Search (BM25 + Semantic + RRF) - Cross-Encoder Reranking - LLM-as-a-Judge*"
    )
    with gr.Row():
        query_box = gr.Textbox(label="Your Question",
                               placeholder="e.g. What is backpropagation?", scale=4)
        eval_check = gr.Checkbox(label="Run Evaluation (slower)", value=True)
        btn = gr.Button("Ask", variant="primary", scale=1)

    with gr.Tabs():
        with gr.TabItem("Answer"):
            answer_out = gr.Markdown()
        with gr.TabItem("Retrieved Context"):
            context_out = gr.Markdown()
        with gr.TabItem("Evaluation Scores"):
            scores_out = gr.Markdown()

    timing_out = gr.Textbox(label="Timing", interactive=False)
    btn.click(fn=run_query, inputs=[query_box, eval_check],
              outputs=[answer_out, context_out, scores_out, timing_out])

    gr.Markdown("---\n*Embedding: all-MiniLM-L6-v2 | Reranker: ms-marco-MiniLM | "
                "LLM: Meta-Llama-3-8B-Instruct | Vector DB: Pinecone*")

if __name__ == "__main__":
    demo.launch()