"""Build a preference dataset by sampling two answers per prompt and ranking them.

Purpose: the data half of the DPO lab. Samples two candidate answers from the model
    you are about to tune, then decides which is better either by asking you or by
    asking a local judge model through an OpenAI-compatible endpoint, and writes the
    result in the preference shape TRL's DPOTrainer expects: a prompt, a chosen
    completion and a rejected one. Ties and disagreements are dropped rather than
    guessed at, because a pair whose direction you are unsure of is training signal
    pointing at nothing.
Platform: all. Sampling runs either locally through transformers (any track) or
    against a served model over HTTP; judging always goes over HTTP.
Minimum memory: 16 GB when sampling locally from a 1B to 4B model; very little when
    both the sampler and the judge are served from another machine.
Assumes: Python 3.10 or newer. Local sampling needs torch, transformers and, for an
    adapter, peft. Judging needs a server that accepts OpenAI-compatible chat
    completions with a JSON schema response format, such as llama-server from Part 6
    or the gateway from Part 9.

Usage: python3 make-preference-pairs.py --model Qwen/Qwen3-1.7B --adapter runs/sft-my-format \
           --rank manual --out-dir pairs
       python3 make-preference-pairs.py --base-url http://127.0.0.1:8080/v1 --served-model local \
           --rank judge --judge-url http://127.0.0.1:8081/v1 --judge-model qwen3-8b --out-dir pairs
"""

from __future__ import annotations

import argparse
import hashlib
import json
import random
import sys
import time
import urllib.error
import urllib.request
from pathlib import Path
from typing import Any, Optional

# The default prompts are about running local models, which is what this course's
# readers write about; replace them with your own work and the adapter becomes
# genuinely yours. One line per prompt, or a JSONL file with a "prompt" field.
DEFAULT_PROMPTS = [
    "Explain what a KV cache is and why it grows with context length.",
    "A colleague asks whether to quantise a model to Q4_K_M or Q8_0. What do you tell them?",
    "Write a short note explaining why a 30B mixture-of-experts model can be faster than a 14B dense one.",
    "Summarise the difference between prefill and decode for someone who has never served a model.",
    "How would you decide whether to fine-tune a model or add retrieval?",
    "Describe what happens when a model does not fit in GPU memory on a discrete card.",
    "Explain a chat template to someone who has only ever used a hosted chat interface.",
    "What should be recorded about a training run so the result still means something in a month?",
    "Give practical advice on choosing a context length for a local server.",
    "Why is a benchmark score without hardware, version and date close to useless?",
    "Explain LoRA to an engineer who understands matrix multiplication but not fine-tuning.",
    "A model answers correctly but far too verbosely. What are the options, cheapest first?",
    "How do you tell whether a quantised model has been damaged by the quantisation?",
    "Explain why sampling temperature changes reproducibility, and what to do about it.",
    "What does it mean for a model to be open weight but not open source?",
    "Describe how to expose a local model to another machine on your network safely.",
    "Explain the difference between a tool call and a normal completion.",
    "What is the first thing to check when a served model produces gibberish?",
    "How should someone size a machine for running an 8B model comfortably?",
    "Explain why training loss falling is not the same as the model getting better.",
    "What is the point of a held-out split, in one paragraph?",
    "Describe when a smaller model is the right answer even though a larger one scores higher.",
    "Explain speculative decoding without using the word 'draft' more than twice.",
    "A fine-tune made the model worse on everything else. What happened and what next?",
]

JUDGE_SYSTEM = """You compare two answers to the same question and say which better matches
the house style described below. Reply with JSON only.

Judge only against the house style and the correctness of the content. Ignore which answer is
longer, which sounds more confident, and which is shown first. If neither is clearly better,
say tie."""

JUDGE_SCHEMA = {
    "type": "object",
    "properties": {
        "winner": {"type": "string", "enum": ["first", "second", "tie"]},
        "why": {"type": "string"},
    },
    "required": ["winner", "why"],
    "additionalProperties": False,
}

DEFAULT_STYLE = (
    "Answers are direct and short. They open with the answer itself, not with a restatement of "
    "the question. They use plain words, British spelling, and no bulleted lists unless the "
    "content is genuinely a list. They say plainly when something is uncertain, and they never "
    "close with an offer of further help."
)


# --------------------------------------------------------------------------------------
# HTTP
# --------------------------------------------------------------------------------------

def post_json(url: str, payload: dict, api_key: Optional[str], timeout: int) -> dict:
    body = json.dumps(payload).encode("utf-8")
    headers = {"Content-Type": "application/json"}
    if api_key:
        headers["Authorization"] = f"Bearer {api_key}"
    request = urllib.request.Request(url, data=body, headers=headers, method="POST")
    try:
        with urllib.request.urlopen(request, timeout=timeout) as response:
            return json.loads(response.read().decode("utf-8"))
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode("utf-8", "replace")[:400]
        raise RuntimeError(f"{url} returned HTTP {exc.code}: {detail}") from exc
    except urllib.error.URLError as exc:
        raise RuntimeError(f"cannot reach {url}: {exc.reason}") from exc


# --------------------------------------------------------------------------------------
# Sampling
# --------------------------------------------------------------------------------------

class ServedSampler:
    """Two answers per prompt from an OpenAI-compatible endpoint."""

    def __init__(self, base_url: str, model: str, api_key: Optional[str], timeout: int,
                 temperature: float, max_tokens: int):
        self.endpoint = base_url.rstrip("/") + "/chat/completions"
        self.model = model
        self.api_key = api_key
        self.timeout = timeout
        self.temperature = temperature
        self.max_tokens = max_tokens

    def sample(self, prompt: str, seed: int) -> str:
        payload = {
            "model": self.model,
            "messages": [{"role": "user", "content": prompt}],
            "temperature": self.temperature,
            "max_tokens": self.max_tokens,
            "seed": seed,
        }
        body = post_json(self.endpoint, payload, self.api_key, self.timeout)
        return (body["choices"][0]["message"]["content"] or "").strip()


class LocalSampler:
    """Two answers per prompt from a model loaded in this process."""

    def __init__(self, model_id: str, adapter: Optional[str], temperature: float, max_tokens: int):
        import torch  # noqa: PLC0415 - only needed for local sampling
        from transformers import AutoModelForCausalLM, AutoTokenizer  # noqa: PLC0415

        self.torch = torch
        self.temperature = temperature
        self.max_tokens = max_tokens
        device = 'cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu'
        dtype = torch.bfloat16 if device == 'cuda' and torch.cuda.is_bf16_supported() else torch.float32
        if adapter:
            from peft import AutoPeftModelForCausalLM  # noqa: PLC0415
            self.model = AutoPeftModelForCausalLM.from_pretrained(adapter, dtype=dtype).to(device)
            self.tokenizer = AutoTokenizer.from_pretrained(adapter)
        else:
            self.model = AutoModelForCausalLM.from_pretrained(model_id, dtype=dtype).to(device)
            self.tokenizer = AutoTokenizer.from_pretrained(model_id)
        self.model.eval()
        print(f"sampling locally from {adapter or model_id} on {self.model.device}")

    def sample(self, prompt: str, seed: int) -> str:
        self.torch.manual_seed(seed)
        messages = [{"role": "user", "content": prompt}]
        inputs = self.tokenizer.apply_chat_template(
            messages, tokenize=True, add_generation_prompt=True, return_dict=True, return_tensors="pt"
        ).to(self.model.device)
        with self.torch.no_grad():
            out = self.model.generate(**inputs, max_new_tokens=self.max_tokens,
                                      do_sample=True, temperature=self.temperature, top_p=0.95)
        generated = out[0][inputs["input_ids"].shape[-1]:]
        return self.tokenizer.decode(generated, skip_special_tokens=True).strip()


# --------------------------------------------------------------------------------------
# Ranking
# --------------------------------------------------------------------------------------

def rank_manually(prompt: str, first: str, second: str, index: int, total: int) -> Optional[int]:
    """Ask the reader. Returns 0 or 1 for the better answer, or None to skip the pair."""
    print("\n" + "=" * 78)
    print(f"[{index}/{total}] {prompt}")
    for label, answer in (("A", first), ("B", second)):
        print(f"\n--- {label} " + "-" * 70)
        print(answer if answer else "(empty)")
    print("\n" + "-" * 78)
    while True:
        choice = input("Better answer? [a/b/s to skip/q to stop] ").strip().lower()
        if choice in ("a", "b"):
            return 0 if choice == "a" else 1
        if choice == "s":
            return None
        if choice == "q":
            raise KeyboardInterrupt
        print("Type a, b, s or q.")


def rank_with_judge(prompt: str, first: str, second: str, args) -> Optional[int]:
    """Ask the judge twice with the order swapped; disagreement means the pair is dropped.

    A judge that changes its answer when the answers change places has told you it is
    scoring position, not quality, and that pair carries no signal. Part 10's judge
    does the same swap for the same reason.
    """
    endpoint = args.judge_url.rstrip("/") + "/chat/completions"

    def ask(a: str, b: str) -> Optional[str]:
        user = (f"House style:\n{args.style}\n\nQuestion:\n{prompt}\n\n"
                f"First answer:\n{a}\n\nSecond answer:\n{b}")
        payload = {
            "model": args.judge_model,
            "messages": [{"role": "system", "content": JUDGE_SYSTEM},
                         {"role": "user", "content": user}],
            "temperature": 0.0,
            "max_tokens": 300,
            "seed": args.seed,
            "response_format": {"type": "json_schema",
                                "json_schema": {"name": "verdict", "schema": JUDGE_SCHEMA, "strict": True}},
        }
        try:
            body = post_json(endpoint, payload, args.judge_api_key, args.timeout)
            return json.loads(body["choices"][0]["message"]["content"]).get("winner")
        except (RuntimeError, KeyError, ValueError, TypeError) as exc:
            print(f"    judge call failed: {exc}", file=sys.stderr)
            return None

    forward = ask(first, second)
    backward = ask(second, first)
    if forward is None or backward is None:
        return None
    if forward == "tie" or backward == "tie":
        return None
    if forward == "first" and backward == "second":
        return 0
    if forward == "second" and backward == "first":
        return 1
    return None  # the judge flipped with the order; the pair tells us nothing


# --------------------------------------------------------------------------------------

def load_prompts(path: Optional[str]) -> list[str]:
    if path is None:
        return list(DEFAULT_PROMPTS)
    text = Path(path).read_text(encoding="utf-8")
    prompts = []
    for line in text.splitlines():
        line = line.strip()
        if not line:
            continue
        if line.startswith("{"):
            row = json.loads(line)
            prompts.append(row["prompt"] if isinstance(row.get("prompt"), str) else row["prompt"][-1]["content"])
        else:
            prompts.append(line)
    return prompts


def write_jsonl(path: Path, rows: list[dict]) -> str:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as handle:
        for row in rows:
            handle.write(json.dumps(row, ensure_ascii=False) + "\n")
    return hashlib.sha256(path.read_bytes()).hexdigest()


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--prompts", default=None,
                        help="a file of prompts, one per line or as JSONL; the built-in set is used if omitted")
    parser.add_argument("--model", default="Qwen/Qwen3-1.7B", help="local model to sample from")
    parser.add_argument("--adapter", default=None, help="a LoRA adapter to sample through")
    parser.add_argument("--base-url", default=None,
                        help="sample from a served model instead of loading one locally")
    parser.add_argument("--served-model", default=None, help="model name the sampling server answers to")
    parser.add_argument("--api-key", default=None)
    parser.add_argument("--rank", default="manual", choices=["manual", "judge"])
    parser.add_argument("--judge-url", default="http://127.0.0.1:8080/v1")
    parser.add_argument("--judge-model", default=None, help="a different, preferably larger model")
    parser.add_argument("--judge-api-key", default=None)
    parser.add_argument("--style", default=DEFAULT_STYLE,
                        help="the house style the judge is asked to prefer; write your own")
    parser.add_argument("--temperature", type=float, default=0.9,
                        help="above zero, or the two samples will be the same answer twice")
    parser.add_argument("--max-tokens", type=int, default=400)
    parser.add_argument("--valid-fraction", type=float, default=0.2)
    parser.add_argument("--out-dir", default="pairs")
    parser.add_argument("--timeout", type=int, default=300)
    parser.add_argument("--seed", type=int, default=0)
    args = parser.parse_args()

    if args.rank == "judge" and not args.judge_model:
        raise SystemExit("--rank judge needs --judge-model, and it should not be the model under test")

    prompts = load_prompts(args.prompts)
    print(f"{len(prompts)} prompts; two samples each at temperature {args.temperature}")

    if args.base_url:
        if not args.served_model:
            raise SystemExit("--base-url needs --served-model")
        sampler: Any = ServedSampler(args.base_url, args.served_model, args.api_key,
                                     args.timeout, args.temperature, args.max_tokens)
    else:
        sampler = LocalSampler(args.model, args.adapter, args.temperature, args.max_tokens)

    raw: list[dict] = []
    kept: list[dict] = []
    skipped = 0
    started = time.time()

    try:
        for i, prompt in enumerate(prompts, start=1):
            first = sampler.sample(prompt, args.seed + 2 * i)
            second = sampler.sample(prompt, args.seed + 2 * i + 1)
            if first == second:
                print(f"  [{i}/{len(prompts)}] both samples identical; skipped")
                skipped += 1
                continue

            if args.rank == "manual":
                better = rank_manually(prompt, first, second, i, len(prompts))
            else:
                better = rank_with_judge(prompt, first, second, args)
                verdict = "skipped" if better is None else ("A" if better == 0 else "B")
                print(f"  [{i}/{len(prompts)}] judge: {verdict}")

            raw.append({"prompt": prompt, "a": first, "b": second,
                        "better": better, "ranked_by": args.rank})
            if better is None:
                skipped += 1
                continue
            chosen, rejected = (first, second) if better == 0 else (second, first)
            kept.append({
                "prompt": [{"role": "user", "content": prompt}],
                "chosen": [{"role": "assistant", "content": chosen}],
                "rejected": [{"role": "assistant", "content": rejected}],
            })
    except KeyboardInterrupt:
        print("\nstopped early; writing what has been ranked so far")

    if not kept:
        raise SystemExit("no pairs were kept. With --rank judge, a judge that ties or flips on "
                         "every pair usually means the two samples are too similar: raise "
                         "--temperature, or write prompts where style has room to differ.")

    rng = random.Random(args.seed)
    rng.shuffle(kept)
    split = max(1, int(len(kept) * args.valid_fraction))
    valid, train = kept[:split], kept[split:]

    out = Path(args.out_dir)
    train_hash = write_jsonl(out / "train.jsonl", train)
    valid_hash = write_jsonl(out / "valid.jsonl", valid)
    write_jsonl(out / "raw.jsonl", raw)

    elapsed = time.time() - started
    print(f"\nranked {len(raw)} pairs in {elapsed:.0f} s; kept {len(kept)}, skipped {skipped}")
    print(f"  train {len(train)}  sha256 {train_hash[:12]}...")
    print(f"  valid {len(valid)}  sha256 {valid_hash[:12]}...")
    print(f"written to {out}/  (train.jsonl, valid.jsonl, raw.jsonl)")
    print("Record the train hash in the run log; it is what ties a DPO result to this exact set.")


if __name__ == "__main__":
    main()