"""Housing Statute QA (reglab/housing_qa), New York subset, with and
without retrieval. Metrics: yes/no accuracy and retrieval hit rate.

Usage: python run_housingqa.py --mode rag --k 5
"""
import argparse, os, re
from bench_common import (Retriever, generate, context_block, norm_ws,
                          save_outputs, SYSTEM_PREAMBLE, NORAG_PREAMBLE)

STATE = "NY"
SCRATCH = os.environ.get("SCRATCH", ".")  # results/ and emb/ land here


def chunk_text(text, size=1500, overlap=150):
    out, start = [], 0
    while start < len(text):
        out.append(text[start:start + size])
        start += size - overlap
    return out


def load_data():
    """Download the questions and statutes files and filter both to NY.
    (datasets>=3 can't load this old-style script dataset directly.)"""
    import json, zipfile
    import pandas as pd
    from huggingface_hub import hf_hub_download

    def is_target_state(v):
        return str(v).strip().lower() in {"ny", "new york"}

    qzip = hf_hub_download("reglab/housing_qa", "data/questions.json.zip",
                           repo_type="dataset")
    with zipfile.ZipFile(qzip) as z:
        name = [n for n in z.namelist() if n.endswith(".json")][0]
        all_qs = json.loads(z.read(name).decode("utf-8"))
    qs = [q for q in all_qs if is_target_state(q.get("state"))]
    if not qs:
        seen = sorted({str(q.get("state")) for q in all_qs})[:60]
        raise RuntimeError(f"no questions matched state {STATE}; "
                           f"distinct state values in data: {seen}")

    szip = hf_hub_download("reglab/housing_qa", "data/statutes.tsv.zip",
                           repo_type="dataset")
    kept, seen_states = [], set()
    for chunk in pd.read_csv(szip, sep="\t", chunksize=200_000,
                             dtype=str, on_bad_lines="skip"):
        seen_states.update(chunk["state"].dropna().unique().tolist())
        kept.append(chunk[chunk["state"].map(is_target_state)])
    statutes = pd.concat(kept, ignore_index=True)
    if not len(statutes):
        raise RuntimeError(f"no statutes matched state {STATE}; "
                           f"distinct values: {sorted(seen_states)[:60]}")
    print(f"{len(qs)} {STATE} questions, {len(statutes)} {STATE} statutes")

    passages = []
    for _, s in statutes.iterrows():
        text = s.get("text") or ""
        if not isinstance(text, str) or not text.strip():
            continue
        for c in chunk_text(text):
            passages.append({"text": c, "id": str(s["citation"]),
                             "statute_idx": str(s["idx"])})
    print(f"{len(passages)} statute chunks")
    return qs, passages


def parse_yesno(response):
    m = re.search(r"\b(yes|no)\b", response.lower())
    return m.group(1) if m else "unparsed"


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--mode", choices=["rag", "norag"], required=True)
    ap.add_argument("--k", type=int, default=5)
    ap.add_argument("--limit", type=int, default=0,
                    help="debug only; leave 0 for the full NY set")
    args = ap.parse_args()

    qs, passages = load_data()
    if args.limit:
        qs = qs[:args.limit]
        print(f"DEBUG LIMIT: {args.limit} questions")

    retriever = None
    if args.mode == "rag":
        retriever = Retriever(passages, f"{SCRATCH}/emb/housingqa_{STATE}.npy")

    samples, n_correct, n_hit = [], 0, 0
    for q in qs:
        instruction = (f"Answer with Yes or No first, then one sentence of "
                       f"support.\n\nQuestion (state: {STATE}): {q['question']}")
        gold_cites = [st["citation"] for st in q.get("statutes", [])]
        if args.mode == "rag":
            hits = retriever.search(q["question"], k=args.k)
            block = f"Context:\n{context_block(hits)}\n\n{instruction}"
            resp = generate(block, SYSTEM_PREAMBLE)
            retrieved = [{"id": h["id"], "score": s} for h, s in hits]
            hit = any(norm_ws(h["id"]) in {norm_ws(c) for c in gold_cites}
                      for h, _ in hits)
        else:
            resp = generate(instruction, NORAG_PREAMBLE)
            retrieved, hit = [], False
        pred = parse_yesno(resp)
        correct = pred == q["answer"].lower()
        n_correct += correct
        n_hit += hit
        samples.append({"idx": q["idx"], "question": q["question"],
                        "gold": q["answer"], "pred": pred, "correct": correct,
                        "gold_citations": gold_cites, "retrieved": retrieved,
                        "retrieval_hit": hit, "response": resp})

    metrics = {"accuracy": n_correct / len(samples)}
    if args.mode == "rag":
        metrics["retrieval_hit_rate"] = n_hit / len(samples)
    save_outputs(f"{SCRATCH}/results", "housingqa", args.mode, metrics,
                 samples, extra={"state": STATE, "k": args.k})


if __name__ == "__main__":
    main()