"""Shared pieces for the eval scripts: model loading, generation, a numpy
retriever, and results I/O. Paths and the generator are configured through
environment variables (see example.env; copy it to .env and edit).
"""
import os, json, time, gc, re

os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")
os.environ.setdefault("HF_HUB_DISABLE_XET", "1")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")


def _load_dotenv(path=".env"):
    """Read KEY=value lines into the environment (existing values win)."""
    try:
        for line in open(path):
            line = line.strip()
            if not line or line.startswith("#") or "=" not in line:
                continue
            k, _, v = line.partition("=")
            if v.strip():
                os.environ.setdefault(k.strip(), v.strip().strip("'\""))
    except FileNotFoundError:
        pass

_load_dotenv()

import numpy as np
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

# MODEL_TAG keeps comparison-model output files separate from earlier runs.
GEN_MODEL = os.environ.get("GEN_MODEL", "Qwen/Qwen3-4B-Instruct-2507")
MODEL_TAG = os.environ.get("MODEL_TAG", "")
EMB_MODEL = "BAAI/bge-base-en-v1.5"
BGE_QUERY_PREFIX = "Represent this sentence for searching relevant passages: "

# Same containment preamble as Check-In 3. The revised entailment preamble
# is intentionally NOT used here; see decision log D-16.
SYSTEM_PREAMBLE = (
    "You are a regulatory assistant. Answer the question using ONLY the "
    "context passages below. Cite the supporting document for every claim. "
    "If the context does not contain the answer, reply: \"The provided "
    "sources do not answer this question\" and briefly say what source "
    "likely would."
)
NORAG_PREAMBLE = (
    "You are a regulatory assistant. Answer the question from your own "
    "knowledge. If you do not know, say so."
)

_PIPE = None

def get_pipe(model_name=GEN_MODEL):
    global _PIPE
    if _PIPE is None:
        tok = AutoTokenizer.from_pretrained(model_name)
        mdl = AutoModelForCausalLM.from_pretrained(
            model_name, device_map="auto", dtype=torch.bfloat16)
        _PIPE = pipeline("text-generation", model=mdl, tokenizer=tok)
    return _PIPE

def generate(user_block, preamble, tokens=400):
    pipe = get_pipe()
    messages = [{"role": "user", "content": preamble + "\n\n" + user_block}]
    out = pipe(messages, max_new_tokens=tokens, do_sample=False,
               return_full_text=False)
    return out[0]["generated_text"].strip()


class Retriever:
    """Embed-and-search over passage dicts; caches the embedding matrix."""
    def __init__(self, passages, cache_path):
        from sentence_transformers import SentenceTransformer
        self.passages = passages
        self.model = SentenceTransformer(
            EMB_MODEL, device="cuda" if torch.cuda.is_available() else "cpu")
        if os.path.exists(cache_path):
            self.mat = np.load(cache_path)
            assert self.mat.shape[0] == len(passages), (
                f"cache {cache_path} has {self.mat.shape[0]} rows but "
                f"{len(passages)} passages were loaded; delete the cache")
        else:
            self.mat = self.model.encode(
                [p["text"] for p in passages], normalize_embeddings=True,
                batch_size=64, show_progress_bar=True, convert_to_numpy=True)
            os.makedirs(os.path.dirname(cache_path), exist_ok=True)
            np.save(cache_path, self.mat)

    def search(self, query, k=5):
        q = self.model.encode([BGE_QUERY_PREFIX + query],
                              normalize_embeddings=True,
                              convert_to_numpy=True)[0]
        scores = self.mat @ q
        idx = np.argsort(-scores)[:k]
        return [(self.passages[i], float(scores[i])) for i in idx]


def context_block(hits):
    return "\n\n".join(
        f"[{j+1}] ({h['id']}): \"{h['text']}\"" for j, (h, _) in enumerate(hits))


def norm_ws(s):
    return re.sub(r"\s+", " ", s).strip().lower()


def save_outputs(out_dir, name, mode, metrics, samples, extra=None):
    os.makedirs(out_dir, exist_ok=True)
    results = {"benchmark": name, "mode": mode, "model": GEN_MODEL,
               "embedding": EMB_MODEL, "n_items": len(samples),
               "metrics": metrics, "finished": time.strftime("%Y-%m-%d %H:%M")}
    if extra:
        results.update(extra)
    stem = f"{name}_{mode}" + (f"_{MODEL_TAG}" if MODEL_TAG else "")
    rp = os.path.join(out_dir, f"{stem}_results.json")
    sp = os.path.join(out_dir, f"{stem}_samples.json")
    with open(rp, "w") as fh:
        json.dump(results, fh, indent=2)
    with open(sp, "w") as fh:
        json.dump(samples, fh, indent=2)
    print(f"saved {rp}\nsaved {sp}")
    print(json.dumps(metrics, indent=2))


def cleanup():
    global _PIPE
    _PIPE = None
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()