#!/usr/bin/env python3
"""Convert the official RAGTruth release into our TokenEvalRow format.

RAGTruth ships character-span annotations on response text. This pipeline needs
one label per *model* token, so the conversion re-tokenises each response with
the scoring model's own tokenizer and maps spans onto tokens by character
overlap. Their word-level `generated_tokens` are deliberately ignored: the only
durable part of a span annotation is its character offsets.

Two things make this replay faithful rather than approximate:

* **Model matching.** Only responses generated by `--generator-model` are kept,
  so the model whose internals we later read is the model that wrote the text.
  Scoring another model's output would measure surprise at someone else's
  continuation, not self-detection. LLM-Check (NeurIPS 2024) filters RAGTruth
  the same way in `utils_ragtruth.py`.
* **Prompt matching.** RAGTruth's README states its generation wrapper exactly:
  "For Llama and Mistral models, we used `<s>[INST] {prompt} [/INST]`". We
  rebuild that string, so the hidden state at each prediction row is the state
  the model was actually in when it emitted the next token.

Label semantics follow this repo's convention: 0 = grounded, 1 = hallucinated,
2 = uncertain. RAGTruth has no uncertain class, so 2 is never assigned -- every
token outside a gold span is grounded.

Usage:
    python scripts/convert_ragtruth.py \
        --ragtruth-dir ~/Desktop/datasets/hallucinations/ragtruth_official \
        --generator-model llama-2-7b-chat \
        --tokenizer NousResearch/Llama-2-7b-chat-hf \
        --output outputs/ragtruth_llama2_7b_chat.jsonl
"""

from __future__ import annotations

import argparse
import collections
import hashlib
import json
from pathlib import Path
import sys

REPO_ROOT = Path(__file__).resolve().parents[1]
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

# RAGTruth README: "For Llama and Mistral models, we used `<s>[INST] {prompt} [/INST]`".
# The BOS token is added by the tokenizer, so it is not written literally here.
LLAMA_INSTRUCT_WRAPPER = "[INST] {prompt} [/INST]"


def parse_args(argv=None):
    p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    p.add_argument("--ragtruth-dir", type=Path, required=True)
    p.add_argument("--output", type=Path, required=True)
    p.add_argument("--report", type=Path, default=None, help="Defaults to <output>.report.json")
    p.add_argument(
        "--generator-model",
        default="llama-2-7b-chat",
        help="Keep only responses from this generator; it must match the scoring model.",
    )
    p.add_argument(
        "--tokenizer",
        default="NousResearch/Llama-2-7b-chat-hf",
        help=(
            "Tokenizer of the scoring model. The official meta-llama repo is "
            "gated; NousResearch mirrors it byte-for-byte."
        ),
    )
    p.add_argument(
        "--scoring-model-id",
        default=None,
        help=(
            "HuggingFace id of the model that will score these rows. Written to "
            "metadata.generator_model so the replay guard can verify that the "
            "scoring model IS the generating model. Defaults to --tokenizer. "
            "RAGTruth's own label is preserved in extra.ragtruth_generator."
        ),
    )
    p.add_argument(
        "--splits",
        nargs="*",
        default=["train", "test"],
        help="RAGTruth splits to convert.",
    )
    p.add_argument(
        "--cross-check",
        type=Path,
        default=None,
        help="Optional teammate-normalised JSONL to verify our spans against.",
    )
    return p.parse_args(argv)


def token_labels_from_spans(offsets, spans, answer_len):
    """One label per token; a token is H if it overlaps any gold span.

    Overlap rule ``max(tok_start, span_start) < min(tok_end, span_end)``, the
    same half-open convention RAGTruth's own annotation tooling uses. Zero-width
    tokens (some tokenizers emit them for specials) can never overlap and stay
    grounded.
    """
    labels = [0] * len(offsets)
    covered = []
    for span in spans:
        s, e = int(span["start"]), int(span["end"])
        if s >= e or s < 0 or e > answer_len:
            covered.append(0)
            continue
        hit = 0
        for index, (ts, te) in enumerate(offsets):
            if ts >= te:
                continue
            if max(ts, s) < min(te, e):
                labels[index] = 1
                hit += 1
        covered.append(hit)
    return labels, covered


def build_row(record, source, tokenizer, generator_model, scoring_model_id):
    answer = record["response"]
    encoded = tokenizer(answer, add_special_tokens=False, return_offsets_mapping=True)
    token_ids = list(encoded["input_ids"])
    offsets = [tuple(o) for o in encoded["offset_mapping"]]
    if not token_ids:
        return None, "empty_answer"

    spans = [
        {"start": int(s["start"]), "end": int(s["end"]), "text": s.get("text")}
        for s in (record.get("labels") or [])
    ]
    # Drop RAGTruth's `meta`/`label_type` from the span itself (our schema has no
    # slot for it) but keep the taxonomy in metadata so the error types survive.
    span_types = [s.get("label_type") for s in (record.get("labels") or [])]

    labels, covered = token_labels_from_spans(offsets, spans, len(answer))
    t_star = next((i for i, v in enumerate(labels) if v == 1), None)

    prompt = LLAMA_INSTRUCT_WRAPPER.format(prompt=source["prompt"])
    row = {
        "id": f"ragtruth_{record['id']}",
        "prompt": prompt,
        "answer": answer,
        "labels": labels,
        "hallucinated_spans": spans,
        "metadata": {
            "source": "ragtruth",
            "dataset": "ParticleMedia/RAGTruth",
            "split": record["split"],
            # The id the replay guard checks. RAGTruth generated these with
            # Meta's llama-2-7b-chat; we score with the same weights, so this
            # records the repo we actually load rather than RAGTruth's label.
            "generator_model": scoring_model_id,
            "generated_token_ids": token_ids,
            "generated_tokens": tokenizer.convert_ids_to_tokens(token_ids),
            "t_star_index": t_star,
            "extra": {
                "ragtruth_generator": generator_model,
                "ragtruth_response_id": record["id"],
                "ragtruth_source_id": record["source_id"],
                "task_type": source["task_type"],
                "upstream_source": source.get("source"),
                "temperature": record.get("temperature"),
                "quality": record.get("quality"),
                "hallucination_types": span_types,
                "raw_prompt": source["prompt"],
                "prompt_wrapper": "llama_inst_v1",
                "split_group": f"ragtruth_source_{record['source_id']}",
            },
        },
    }
    return row, {"covered": covered, "n_spans": len(spans), "n_tokens": len(token_ids)}


def main(argv=None):
    args = parse_args(argv)
    from transformers import AutoTokenizer

    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer)
    scoring_model_id = args.scoring_model_id or args.tokenizer
    if not tokenizer.is_fast:
        raise SystemExit("a fast tokenizer is required for character offsets")

    rt = args.ragtruth_dir
    sources = {}
    with (rt / "source_info.jsonl").open(encoding="utf-8") as f:
        for line in f:
            s = json.loads(line)
            sources[s["source_id"]] = s

    rows, stats = [], []
    skipped = collections.Counter()
    task_types = collections.Counter()
    splits = collections.Counter()
    with (rt / "response.jsonl").open(encoding="utf-8") as f:
        for line in f:
            rec = json.loads(line)
            if rec["model"] != args.generator_model:
                skipped["other_generator"] += 1
                continue
            if rec["split"] not in args.splits:
                skipped["other_split"] += 1
                continue
            source = sources.get(rec["source_id"])
            if source is None:
                skipped["missing_source"] += 1
                continue
            row, stat = build_row(
                rec, source, tokenizer, args.generator_model, scoring_model_id
            )
            if row is None:
                skipped[stat] += 1
                continue
            rows.append(row)
            stats.append(stat)
            task_types[source["task_type"]] += 1
            splits[rec["split"]] += 1

    if not rows:
        raise SystemExit(f"no rows matched generator {args.generator_model!r}")

    args.output.parent.mkdir(parents=True, exist_ok=True)
    with args.output.open("w", encoding="utf-8") as f:
        for row in rows:
            f.write(json.dumps(row, ensure_ascii=False) + "\n")

    total_tokens = sum(s["n_tokens"] for s in stats)
    h_tokens = sum(sum(1 for v in r["labels"] if v == 1) for r in rows)
    all_spans = sum(s["n_spans"] for s in stats)
    lost_spans = sum(1 for s in stats for c in s["covered"] if c == 0)
    rows_with_h = sum(1 for r in rows if any(v == 1 for v in r["labels"]))

    report = {
        "output": str(args.output),
        "output_sha256": hashlib.sha256(args.output.read_bytes()).hexdigest(),
        "generator_model": args.generator_model,
        "tokenizer": args.tokenizer,
        "scoring_model_id": scoring_model_id,
        "prompt_wrapper": LLAMA_INSTRUCT_WRAPPER,
        "rows": len(rows),
        "splits": dict(splits),
        "task_types": dict(task_types),
        "answer_tokens": total_tokens,
        "hallucinated_tokens": h_tokens,
        "hallucinated_token_rate": h_tokens / max(total_tokens, 1),
        "rows_with_hallucination": rows_with_h,
        "gold_spans": all_spans,
        "spans_covering_zero_tokens": lost_spans,
        "span_survival_rate": 1.0 - lost_spans / max(all_spans, 1),
        "skipped": dict(skipped),
    }

    if args.cross_check and args.cross_check.exists():
        report["cross_check"] = cross_check(args.cross_check, rows, args.generator_model)

    report_path = args.report or args.output.with_suffix(args.output.suffix + ".report.json")
    report_path.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8")

    print(f"wrote {len(rows):,} rows -> {args.output}")
    print(f"  answer tokens {total_tokens:,} | H {h_tokens:,} ({report['hallucinated_token_rate']:.2%})")
    print(f"  rows with H {rows_with_h:,} | task types {dict(task_types)}")
    print(f"  gold spans {all_spans:,} | covering zero tokens {lost_spans} "
          f"(survival {report['span_survival_rate']:.4%})")
    if "cross_check" in report:
        print(f"  cross-check: {report['cross_check']}")
    print(f"wrote report -> {report_path}")
    return 0


def cross_check(path: Path, rows, generator_model):
    """Compare our character spans against a teammate's normalisation.

    Two independent normalisations agreeing on spans is real evidence the
    conversion is right; it is not a substitute for the coverage report above.
    """
    theirs = {}
    with path.open(encoding="utf-8") as f:
        for line in f:
            r = json.loads(line)
            if r.get("model") != generator_model:
                continue
            key = str(r.get("id"))
            theirs[key] = {(int(s["start"]), int(s["end"])) for s in (r.get("hallucinated_spans") or [])}
    matched = agreed = 0
    for row in rows:
        key = str(row["metadata"]["extra"]["ragtruth_response_id"])
        if key not in theirs:
            continue
        matched += 1
        ours = {(s["start"], s["end"]) for s in row["hallucinated_spans"]}
        if ours == theirs[key]:
            agreed += 1
    return {
        "rows_matched_by_id": matched,
        "rows_with_identical_spans": agreed,
        "agreement_rate": agreed / max(matched, 1),
        "rows_only_in_ours": len(rows) - matched,
    }


if __name__ == "__main__":
    raise SystemExit(main())