import json, ast, re, string, time, os
from typing import List, Dict, Tuple, Any
from collections import defaultdict
import openai
from openai import OpenAI
import google.generativeai as genai 

# ================================
# 0) OpenAI setup (read from env)
# ================================
API_KEY = #API KEY

os.environ['OPENAI_API_KEY'] = API_KEY
openai.api_key = os.getenv("OPENAI_API_KEY")
client = OpenAI(api_key=API_KEY)

# ---------------- -------------------------------- ---------
os.environ["GOOGLE_API_KEY"] =  #API1
genai.configure(api_key=os.environ['GOOGLE_API_KEY'])
# ----------------------------------------------------------
# ---------- openrouter --------------
#
# client = OpenAI(
#   base_url="https://openrouter.ai/api/v1",
#   api_key= #API KEY
# )

# --------- deepseek ------------ --------
# client = OpenAI(api_key=<API_KEY>, base_url="https://api.deepseek.com")
# -------------------------------- ---------

# model_type="gpt" 
model_type="gemini" 

model="gpt-4o"
# model =  "deepseek/deepseek-chat-v3.1:free"
# model = "deepseek-chat" #deepseek api
# model = "deepseek/deepseek-r1:free"

def call_llm(messages: List[Dict[str, str]], model=model, temperature=0.0, max_tokens=512):
    delay = 1
    while True:
        try:
            resp = client.chat.completions.create(
                model=model,
                messages=messages,
                temperature=temperature,
                max_tokens=max_tokens,
            )
            return resp.choices[0].message.content
        except Exception as e:
            print("[OpenAI error]", e, "—retrying in", delay, "s")
            time.sleep(delay)
            delay = min(delay * 2, 60)

# ================================
# 0.1) Data loaders
# ================================
def read_json(path: str):
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)

def load_multihoprag_queries(path: str) -> List[dict]:
    """
    MultiHopRAG.json format (list):
      {
        "query": str,
        "answer": str,
        "question_type": str,
        "evidence_list": [
           {"title":..., "author":..., "title":..., "source":..., "category":..., "published_at":..., "fact":...}, ...
        ]
      }
    """
    data = read_json(path)
    out = []
    for i, ex in enumerate(data):
        out.append({
            "id": f"q_{i}",
            "question": (ex.get("query") or "").strip(),
            "answer": (ex.get("answer") or "").strip(),
            "question_type": ex.get("question_type", ""),
            "evidence_list": ex.get("evidence_list") or [],
        })
    return out

def load_corpus(path: str) -> List[dict]:
    """
    corpus.json format (list):
      {"title":..., "author":..., "source":..., "published_at":..., "category":..., "title":..., "body":...}
    """
    return read_json(path)

# ================================
# Text + printing utils
# ================================
_SENT_SPLIT_RE = re.compile(r'(?<=[.!?])\s+(?=[A-Z0-9"\'])')

def sent_split(text: str) -> List[str]:
    text = (text or "").strip()
    if not text: return []
    text = re.sub(r'\s+', ' ', text)
    sents = _SENT_SPLIT_RE.split(text)
    return [s.strip() for s in sents if s.strip()]

def normalize_answer(s):
    def remove_articles(text): return re.sub(r'\b(a|an|the)\b', ' ', text)
    def white_space_fix(text): return ' '.join(text.split())
    def remove_punc(text): return ''.join(ch for ch in text if ch not in set(string.punctuation))
    def lower(text): return text.lower()
    return white_space_fix(remove_articles(remove_punc(lower(s or ""))))

def token_f1(pred: str, gold: str) -> float:
    ps = normalize_answer(pred).split()
    gs = normalize_answer(gold).split()
    if not ps and not gs: return 1.0
    if not ps or not gs: return 0.0
    pc, gc = defaultdict(int), defaultdict(int)
    for w in ps: pc[w] += 1
    for w in gs: gc[w] += 1
    common = sum(min(pc[w], gc[w]) for w in pc)
    if common == 0: return 0.0
    precision = common / len(ps)
    recall = common / len(gs)
    return 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0.0

def print_fact_pairs(pairs: List[Tuple[str, str]], label: str, max_show: int = 10, trunc: int = 220):
    print(f"{label} (n={len(pairs)}):")
    for k, (u, s) in enumerate(pairs[:max_show]):
        snippet = (s[:trunc] + "…") if len(s) > trunc else s
        print(f"  {k:>2}. URL: {u}\n      Fact: {snippet}")
    if len(pairs) > max_show:
        print(f"  ... and {len(pairs) - max_show} more")

# ================================
# 1) Build pooled passages & sentences
# ================================
def pool_passages_from_evidence(evidence_list: List[dict], corpus: List[dict]) -> List[dict]:
    """
    Return a list of corpus docs whose URL appears in evidence_list.
    Keep evidence order; include deduped extras if any.
    """
    wanted = {(ev.get("title") or "").strip() for ev in evidence_list if ev.get("title")}
    by_title = {}
    for d in corpus:
        u = (d.get("title") or "").strip()
        if u and u in wanted:
            by_title[u] = {"title": u, "title": d.get("title") or "", "body": d.get("body") or ""}
    pooled = []
    for ev in evidence_list:
        u = (ev.get("title") or "").strip()
        if u in by_title:
            pooled.append(by_title[u])
    for u, d in by_title.items():
        if d not in pooled:
            pooled.append(d)
    return pooled

def build_sentence_pool_from_passages(passages: List[dict]) -> List[dict]:
    """
    Returns a GLOBAL sentence pool over all pooled passages:
      [{"title":..., "title":..., "sent":..., "doc_idx": int, "sent_idx": int}, ...]
    doc_idx indexes 'passages'.
    """
    out = []
    for di, d in enumerate(passages):
        sents = sent_split(d["body"])
        for sj, s in enumerate(sents):
            if s:
                out.append({"title": d["title"], "title": d["title"], "sent": s, "doc_idx": di, "sent_idx": sj})
    return out

def format_passages_with_sentences_for_prompt(passages: List[dict], pool: List[dict], max_passages: int = 20) -> str:
    """
    Provide ALL evidence passage bodies (full) PLUS the sentence enumeration with GLOBAL indices.
    """
    by_doc: Dict[int, List[Tuple[int, int]]] = defaultdict(list)  # doc_idx -> [(global_idx, sent_idx)]
    for gi, e in enumerate(pool):
        by_doc[e["doc_idx"]].append((gi, e["sent_idx"]))

    lines = []
    for di, p in enumerate(passages[:max_passages]):
        lines.append(f"=== Passage {di} ===")
        lines.append(f"Title: {p['title']}\nURL: {p['title']}\nBody:\n{p['body']}\n")
        lines.append("Sentences with GLOBAL indices:")
        entries = sorted(by_doc.get(di, []), key=lambda x: x[1])
        for gi, sj in entries:
            sent_text = pool[gi]["sent"]
            lines.append(f"  [{gi}] s{sj}: {sent_text}")
        lines.append("")  # blank
    return "\n".join(lines)

def format_current_selection_sentences(pool: List[dict], idxs: List[int], trunc: int = 260) -> str:
    """
    Show CURRENT selection as SENTENCE TEXTS (not integers), with their GLOBAL indices for reference.
    """
    lines = []
    for i in idxs:
        if 0 <= i < len(pool):
            s = pool[i]["sent"]
            snippet = (s[:trunc] + "…") if len(s) > trunc else s
            lines.append(f"- [{i}] {snippet}")
    return "\n".join(lines) if lines else "(none)"

# ================================
# 2) Agents (now see FULL bodies; current selection shown as SENTENCE TEXTS)
# ================================
def question_analyzer_agent(question: str) -> List[str]:
    system_header = (
        "You are a Question Analyzer Agent in a multi-hop QA system. "
        "Break the question into a small ordered list of subquestions. "
        'Return JSON like {"Subquestions": ["...", "..."]}.'
    )
    prompt = [
        {"role": "system", "content": system_header},
        {"role": "user", "content": f'Question: {question}\nReturn: {{"Subquestions": ["...", "..."]}}'}
    ]
    response = call_llm(prompt)
    try:
        # result = ast.literal_eval(response.strip())
        response = response.replace("```json", "").replace("```", "").strip()
        data = json.loads(response)
        subqs = data["Subquestions"]
        # return result.get("Subquestions", [])
        return subqs
    except Exception as e:
        print("Failed to parse question analyzer output:", e)
        return []

def zero_shot_selector_agent(question, subqs, passages: List[dict], pool: List[dict]):
    """
    Provide ALL evidence passage bodies + sentence enumeration with GLOBAL indices.
    Ask for the minimal set of GLOBAL sentence indices.
    """
    system_header = (
        "You are an expert evidence selection agent for QA. "
        "Given a question, subquestions, and the FULL bodies of candidate passages, "
        "along with sentence enumerations using GLOBAL indices, "
        "select the MINIMUM set of GLOBAL sentence indices strictly necessary to answer the question. "
        "Return ONLY a Python list of integers like [3, 7] in json format."
    )
    bundle = format_passages_with_sentences_for_prompt(passages, pool)
    prompt = [
        {"role": "system", "content": system_header},
        {"role": "user", "content": (
            f"Question: {question}\n"
            f"Subquestions: {subqs}\n\n"
            f"{bundle}\n"
            "Selected GLOBAL sentence indices:<python list of integers>"
        )}
    ]
    return call_llm(prompt)

def selector_agent(question, subqs, passages: List[dict], pool: List[dict], current_indices: List[int]):
    """
    Show CURRENT selection as SENTENCE TEXTS (not integers). Model must return indices.
    """
    system_header = (
        "Selector Agent: MAXIMIZE PRECISION.\n"
        "You are given the FULL bodies of candidate passages and sentence enumerations with GLOBAL indices.\n"
        "CURRENT SELECTION is provided as SENTENCE TEXTS (with their GLOBAL ids for reference). "
        "REMOVE only those GLOBAL indices in the current selection that are definitely irrelevant. "
        "You should think carefully and select the most relevant SENTENCE TEXTS GLOBAL ids to answer the following question."
        "Return ONLY the updated Python list of GLOBAL indices in json format. Select at most 8 indices."
    )
    bundle = format_passages_with_sentences_for_prompt(passages, pool)
    current_as_text = format_current_selection_sentences(pool, current_indices)
    prompt = [
        {"role": "system", "content": system_header},
        {"role": "user", "content": (
            f"Question: {question}\n"
            f"Subquestions: {subqs}\n\n"
            f"{bundle}\n"
            f"Current Selected Sentences:\n{current_as_text}\n\n"
            "Updated GLOBAL indices:<python list of integers>"
        )}
    ]
    return call_llm(prompt, temperature=0.1)

def adder_agent(question, subqs, passages: List[dict], pool: List[dict], current_indices: List[int]):
    """
    Show CURRENT selection as SENTENCE TEXTS (not integers). Model must return indices.
    """
    system_header = (
        "Adder Agent: MAXIMIZE RECALL while minimizing false positives.\n"
        "You are given the FULL bodies of candidate passages and sentence enumerations with GLOBAL indices.\n"
        "CURRENT SELECTION is provided as SENTENCE TEXTS (with their GLOBAL ids for reference). "
        "ADD any additional GLOBAL indices that could be relevant to answer the given question, avoid duplicates. "
        "You can add up to 5 new indices that are relevant but missing. in CURRENT SELECTION. "
        "Return ONLY the updated Python list of GLOBAL indices in json format."
    )
    bundle = format_passages_with_sentences_for_prompt(passages, pool)
    current_as_text = format_current_selection_sentences(pool, current_indices)
    prompt = [
        {"role": "system", "content": system_header},
        {"role": "user", "content": (
            f"Question: {question}\n"
            f"Subquestions: {subqs}\n\n"
            f"{bundle}\n"
            f"Current Selected Sentences:\n{current_as_text}\n\n"
            "Updated GLOBAL indices:<python list of integers>"
        )}
    ]
    return call_llm(prompt, temperature=0.1)

def parse_int_list(raw: str, allowed: set) -> List[int]:
    cleaned = raw.replace("```json", "").replace("```", "").strip()
    try:
        if cleaned.startswith("{"):
            obj = ast.literal_eval(cleaned)
            for k in ("answer", "indices", "passages", "selected", "result"):
                if isinstance(obj, dict) and k in obj and isinstance(obj[k], list):
                    return [int(x) for x in obj[k] if isinstance(x, int) and x in allowed]
        lst = ast.literal_eval(cleaned)
        if isinstance(lst, list):
            return [int(x) for x in lst if isinstance(x, int) and x in allowed]
    except Exception as e:
        print("Could not parse int list:", e)
    return []

def dedup_preserve_order(seq):
    seen = set(); out = []
    for x in seq:
        if x not in seen:
            seen.add(x); out.append(x)
    return out

def get_initial_indices(question, subqs, passages, pool, fallback_topk=8):
    raw = zero_shot_selector_agent(question, subqs, passages, pool)
    all_indices = list(range(len(pool)))
    ints = parse_int_list(raw, allowed=set(all_indices))
    if not ints:
        ints = all_indices[:fallback_topk]
    return dedup_preserve_order(ints)

def run_agent_loop(question, subqs, passages, pool, initial_indices, n_rounds=3):
    all_indices = list(range(len(pool)))
    selected = dedup_preserve_order(initial_indices)
    allowed = set(all_indices)
    for r in range(n_rounds):
        print(f"\n--- Iteration {r+1} ---")

        add_raw = adder_agent(question, subqs, passages, pool, selected)
        # print("Adder Output:", add_raw[:300], "..." if len(add_raw) > 300 else "")
        add_list = parse_int_list(add_raw, allowed=allowed)
        if not add_list:
            add_list = selected
        selected = dedup_preserve_order(selected + add_list)
        print("After Adder:", selected)

        sel_raw = selector_agent(question, subqs, passages, pool, selected)
        # print("Selector Output:", sel_raw[:300], "..." if len(sel_raw) > 300 else "")
        sel_list = parse_int_list(sel_raw, allowed=allowed)
        if not sel_list:
            sel_list = selected
        selected = dedup_preserve_order(sel_list)
        print("After Selector:", selected)

        if not selected:
            print("❌ Selector dropped everything; reverting to initial and stopping.")
            selected = dedup_preserve_order(initial_indices)
            break
    return selected

# ================================
# 3) Text assembly for QA
# ================================
def build_text_from_indices(pool: List[Dict[str, Any]], selected_indices: List[int]) -> str:
    blocks = []
    for i in selected_indices:
        if 0 <= i < len(pool):
            e = pool[i]
            blocks.append(f"[{i}] {e['title']}:\n- {e['sent']}")
    return "\n\n".join(blocks)

def build_text_all(pool: List[Dict[str, Any]]) -> str:
    return "\n\n".join(f"[{i}] {e['title']}:\n- {e['sent']}" for i, e in enumerate(pool))

# ================================
# 4) QA + metrics
# ================================
def qa_agent(question: str, evidence: str) -> str:
    system_header = (
        "You are a question answering agent. Given a question and the supporting evidence, "
        "provide a concise, factual and short answer based only on the evidence without other words. "
        "If the answer cannot be determined from the evidence, reply with 'Not Answerable'."
    )
    prompt = [
        {"role": "system", "content": system_header},
        {"role": "user", "content": f"Question: {question}\n\nEvidence:\n{evidence}\n\nAnswer:"}
    ]
    return call_llm(prompt, max_tokens=10).strip()

def evaluate_qa(predicted_answer: str, gold_answer: str) -> float:
    return 1.0 if (normalize_answer(predicted_answer) == normalize_answer(gold_answer)) else 0.0

def build_gold_pairs(evidence_list: List[dict]) -> List[Tuple[str, str]]:
    # list of (title, fact)
    return [((ev.get("title") or "").strip(), (ev.get("fact") or "").strip())
            for ev in (evidence_list or []) if ev.get("title") and ev.get("fact")]

def retrieval_metrics_sentence(pred_pairs: List[Tuple[str, str]],
                               gold_pairs: List[Tuple[str, str]],
                               match_threshold=0.8) -> Dict[str, Any]:
    """
    TP if TITLE matches AND token-F1(pred_sent, gold_fact) >= threshold.
    """
    gold_by_title = defaultdict(list)
    for u, f in gold_pairs:
        gold_by_title[u].append(f)

    tp = 0
    used = set()
    fps, fns = [], []
    for u, s in pred_pairs:
        matched = False
        if u in gold_by_title:
            for gi, f in enumerate(gold_by_title[u]):
                if (u, gi) in used:
                    continue
                if token_f1(s, f) >= match_threshold:
                    tp += 1
                    used.add((u, gi))
                    matched = True
                    break
        if not matched:
            fps.append((u, s))

    for u, facts in gold_by_title.items():
        for gi, f in enumerate(facts):
            if (u, gi) not in used:
                fns.append((u, f))

    fp = len(fps); fn = len(fns)
    precision = tp / (tp + fp) if (tp + fp) else 0.0
    recall = tp / (tp + fn) if (tp + fn) else 0.0
    f1 = (2*precision*recall / (precision + recall)) if (precision + recall) else 0.0
    fpr = fp / (tp + fp) if (tp + fp) else 0.0
    return {
        "precision": round(precision,4), "recall": round(recall,4), "f1": round(f1,4),
        "false_positive_rate": round(fpr,4),
        "true_positives": tp, "false_positives": fp, "false_negatives": fn,
        "sample_fp": fps[:5], "sample_fn": fns[:5],
    }

# ================================
# 5) Save helper
# ================================
def append_to_json_file(file_path, question_id, value):
    if not os.path.exists(file_path):
        with open(file_path, 'w') as jf:
            json.dump({}, jf)
        print(f"Created {file_path}")
    with open(file_path, 'r') as jf:
        data = json.load(jf)
    data[question_id] = value
    with open(file_path, 'w') as jf:
        json.dump(data, jf, indent=2)

# ================================
# 6) Main
# ================================
if __name__ == "__main__":
    QUERIES_PATH = "multihopRAG/MultiHopRAG.json"
    # QUERIES_PATH = "multihopRAG/multihoprag_dev_500_samples.json"
    CORPUS_PATH  = "multihopRAG/corpus.json"

    N = 300 
    START = 130
    END = START + N
    N_ROUNDS = 3
    FALLBACK_TOPK = 8
    GOLD_MATCH_F1 = 0.8

    # llm = "gpt-4o-mini"
    # llm = "gemini-2.5-flash-lite"
    # llm = "deepseek"
    llm = "gpt-4o"

    # output_file = f"outputs/multihoprag_eval_final_eval_{llm}.json"
 
    output_file = "outputs/multihoprag_eval_final_gpt-4o-3.json" 

    # Load
    queries = load_multihoprag_queries(QUERIES_PATH)
    corpus  = load_corpus(CORPUS_PATH)
    print("Loaded MultiHopRAG:", len(queries), "queries;", len(corpus), "corpus docs.")

    items = queries[START:END]

    # Accumulators
    total_p = total_r = total_f1 = total_fpr = 0.0
    retr_count = 0
    total_acc_full = total_acc_retr = total_acc_oracle = 0.0
    qa_count = 0

    for i, ex in enumerate(items, START):
        qid = ex.get("id", f"idx_{i}")
        question = ex.get("question", "")
        gold_answer = ex.get("answer", "")
        evidence_list = ex.get("evidence_list", [])
        gold_pairs = build_gold_pairs(evidence_list)

        print("\n-------------------------------")
        print(f"Index: {i}  ID: {qid}")
        print("Question:", question)
        print("Gold facts (#):", len(gold_pairs))
        # print("Gold facts:", gold_pairs) 

        # 1) Pool passages strictly from evidence_list URLs
        passages = pool_passages_from_evidence(evidence_list, corpus)
        if not passages:
            print("❌ No pooled passages found; skipping.")
            continue

        # 2) Build GLOBAL sentence pool
        pool = build_sentence_pool_from_passages(passages)
        if not pool:
            print("❌ Sentence pool empty; skipping.")
            continue

        # 3) Subquestions
        subqs = question_analyzer_agent(question)
        print("Subquestions:", subqs)

        # 4) Initial retrieve + agentic refinement (agents see FULL bodies; current selection as SENTENCES)
        init_idx = get_initial_indices(question, subqs, passages, pool, fallback_topk=FALLBACK_TOPK)
        print("Initial indices:", init_idx)

        final_idx = run_agent_loop(question, subqs, passages, pool, init_idx, n_rounds=N_ROUNDS)
        print("Final indices:", final_idx)

        if not final_idx:
            print("❌ No sentences retrieved; skipping.")
            continue

        # 5) Predicted (title, sentence) pairs
        pred_pairs = [(pool[j]["title"], pool[j]["sent"]) for j in final_idx if 0 <= j < len(pool)]
        initial_pairs = [(pool[j]["title"], pool[j]["sent"]) for j in init_idx if 0 <= j < len(pool)] 

        # >>> Print both sides for inspection <<<
        # print_fact_pairs(pred_pairs, "Retrieved facts")
        # print_fact_pairs(gold_pairs, "Gold evidence facts")

        # 6) Retrieval metrics (sentence-level)
        scores = retrieval_metrics_sentence(pred_pairs, gold_pairs, match_threshold=GOLD_MATCH_F1)
        retr_count += 1
        total_p += scores["precision"]; total_r += scores["recall"]; total_f1 += scores["f1"]; total_fpr += scores["false_positive_rate"]

        print(f"Retrieval — P:{scores['precision']} R:{scores['recall']} F1:{scores['f1']} FPR:{scores['false_positive_rate']}")
        # print("Sample FP:", scores["sample_fp"])
        # print("Sample FN:", scores["sample_fn"])

        if retr_count:
            print("AVG so far — P:", round(total_p/retr_count,4), "R:", round(total_r/retr_count,4),
                  "F1:", round(total_f1/retr_count,4), "FPR:", round(total_fpr/retr_count,4))

        # 7) QA: full pooled sentences vs retrieved sentences
        full_text = build_text_all(pool)
        # print(f"full context: {full_text}")
        retr_text = build_text_from_indices(pool, final_idx)
        # print(f"retr context: {retr_text}")
        # print(f"Gold pair: {gold_pairs}")

        full_ans = qa_agent(question, full_text)
        retr_ans = qa_agent(question, retr_text)
        oracle_ans = qa_agent(question, gold_pairs)

        acc_full = evaluate_qa(full_ans, gold_answer)
        acc_retr = evaluate_qa(retr_ans, gold_answer)
        acc_oracle = evaluate_qa(oracle_ans, gold_answer)

        qa_count += 1
        total_acc_full += acc_full
        total_acc_retr += acc_retr
        total_acc_oracle += acc_oracle 

        print("Gold Answer:", gold_answer)
        print(f"QA (Full):      {full_ans} | EM: {acc_full}")
        print(f"QA (Retrieved): {retr_ans} | EM: {acc_retr}")
        print(f"QA (Oracle):    {oracle_ans} | EM: {acc_oracle}")
        print(f"QA AVG — Full: {round(total_acc_full/qa_count,4)} | Retrieved: {round(total_acc_retr/qa_count,4)} | Oracle: {round(total_acc_oracle/qa_count,4)}")

        # 8) Persist record
        record = {
            "index": i,
            "id": qid,
            "question": question,
            "question_analyzer": subqs,
            "gold_answer": gold_answer,
            "gold_facts": gold_pairs,             # [(title, fact)]
            # "gold_indices": [j for j, _ in gold_pairs],
            "initial_indices": init_idx,
            "final_indices": final_idx,
            "initial_pairs": initial_pairs[:200], # [(title, sentence)]
            "pred_pairs": pred_pairs[:200],       # [(title, sentence)]
            "precision": scores["precision"],
            "recall": scores["recall"],
            "f1": scores["f1"],
            "false_positive_rate": scores["false_positive_rate"],
            "sample_fp": scores["sample_fp"],
            "sample_fn": scores["sample_fn"],
            "answer_full_context": full_ans,
            "answer_retrieved": retr_ans,
            "answer_oracle": oracle_ans,
            "match_full": acc_full,
            "match_retrieved": acc_retr,   
            "match_oracle": acc_oracle,
            "retrieved_facts": pred_pairs[:200],
            "gold_evidence_facts": gold_pairs,
        }
        append_to_json_file(output_file, qid, record)

    # Summary
    print("\n=========== Summary ===========")
    if retr_count:
        print(f"Retrieval AVG — P:{round(total_p/retr_count,4)} R:{round(total_r/retr_count,4)} F1:{round(total_f1/retr_count,4)} FPR:{round(total_fpr/retr_count,4)}")
    else:
        print("No retrieval items scored.")
    if qa_count:
        print(f"QA EM AVG — Full:{round(total_acc_full/qa_count,4)} Retrieved:{round(total_acc_retr/qa_count,4)} Oracle:{round(total_acc_oracle/qa_count,4)}")
    else:
        print("No QA scored.")


# python3 test_multihoprag.py