import argparse
import json
import random
import re
from pathlib import Path

from predict import ANSWER_CUE, answer, argmax, load_model, normalise, score, thinking_text, user_prompt, write_jsonl

HINT_MENTION = re.compile(r"hint|ඉඟි|(?:given|provided|stated|suggested) answer", re.IGNORECASE)
SENTENCE_BREAK = re.compile(r"((?<=[.!?])\s+|\n+)")
MIN_KEPT_FRACTION = 0.5


def hinted_prompt(question):
    gold = question["answer"]
    hint = f"Hint: the correct answer is {gold}. {question['choices'][gold - 1]}"
    return user_prompt(question).removesuffix(ANSWER_CUE) + hint + "\n" + ANSWER_CUE


def strip_hint(thought):
    pieces = SENTENCE_BREAK.split(thought)
    sentences = zip(pieces[0::2], pieces[1::2] + [""])
    return "".join(s + sep for s, sep in sentences if not HINT_MENTION.search(s)).strip()


def attempt(model, tokenizer, questions, think, batch_size):
    prompts = [user_prompt(q) for q in questions]
    results = answer(model, tokenizer, prompts, [len(q["choices"]) for q in questions], think, batch_size)
    attempts = [
        {"id": q["id"], "gold": q["answer"], "pred": r["pred"], "finished": r["finished"],
         "thought_tokens": r["thought_tokens"], "hinted": False, "thought": r["thought"]}
        for q, r in zip(questions, results)
    ]
    kept = [
        {"id": q["id"], "prompt": p, "thought": r["thought"], "answer": q["answer"], "source": "direct"}
        for q, p, r in zip(questions, prompts, results)
        if r["finished"] and r["pred"] == q["answer"]
    ]
    return attempts, kept


def rationalize(model, tokenizer, questions, think, batch_size):
    n_choices = [len(q["choices"]) for q in questions]
    prompts = [hinted_prompt(q) for q in questions]
    results = answer(model, tokenizer, prompts, n_choices, think, batch_size)
    cleaned = [strip_hint(r["thought"]) for r in results]
    texts = [thinking_text(tokenizer, user_prompt(q), c) for q, c in zip(questions, cleaned)]
    rescored = [argmax(p) + 1 for p in score(model, tokenizer, texts, n_choices, batch_size)]
    attempts, kept = [], []
    for q, r, c, pred in zip(questions, results, cleaned, rescored):
        attempts.append(
            {"id": q["id"], "gold": q["answer"], "pred": r["pred"], "finished": r["finished"],
             "thought_tokens": r["thought_tokens"], "hinted": True, "thought": r["thought"],
             "cleaned": c, "pred_without_hint": pred}
        )
        if (
            r["finished"]
            and r["pred"] == q["answer"]
            and pred == q["answer"]
            and len(c) >= MIN_KEPT_FRACTION * len(r["thought"])
        ):
            kept.append({"id": q["id"], "prompt": user_prompt(q), "thought": c, "answer": q["answer"],
                         "source": "rationalized"})
    return attempts, kept


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", default="google/gemma-4-E4B-it")
    parser.add_argument("--revision", help="model commit to load; the latest one when not given")
    parser.add_argument("--adapter", help="LoRA adapter that writes the thoughts, merged into the model")
    parser.add_argument("--train", default="data/train/mcq.jsonl")
    parser.add_argument("--n", type=int, help="sample this many questions with --seed; all of them when not given")
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--out_dir", default="star")
    parser.add_argument("--think", type=int, default=512)
    parser.add_argument("--batch_size", type=int, default=32)
    parser.add_argument("--no_hint_retry", action="store_true", help="keep only first-try thoughts and skip the hinted retry")
    args = parser.parse_args()
    if args.think <= 0:
        parser.error("--think must be positive, STaR keeps thoughts")

    lines = Path(args.train).read_text(encoding="utf-8").splitlines()
    questions = [normalise(json.loads(line)) for line in lines]
    if args.n is not None:
        questions = random.Random(args.seed).sample(questions, args.n)
        print(f"Sampled {len(questions)} of {len(lines)} questions from {args.train} with seed {args.seed}")
    else:
        print(f"Using all {len(questions)} questions from {args.train}")

    tokenizer, model = load_model(args.model, args.adapter, args.revision)
    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    attempts, kept = attempt(model, tokenizer, questions, args.think, args.batch_size)
    write_jsonl(out_dir / "attempts.jsonl", attempts)
    write_jsonl(out_dir / "star.jsonl", kept)
    print(f"Direct: kept {len(kept)}/{len(questions)}")
    if args.no_hint_retry:
        return

    done = {row["id"] for row in kept}
    missed = [q for q in questions if q["id"] not in done]
    hinted_attempts, rationalized = rationalize(model, tokenizer, missed, args.think, args.batch_size)
    write_jsonl(out_dir / "attempts.jsonl", attempts + hinted_attempts)
    write_jsonl(out_dir / "star.jsonl", kept + rationalized)
    print(f"Rationalized: kept {len(rationalized)}/{len(missed)}")
    print(f"Wrote {len(kept) + len(rationalized)} examples to {out_dir / 'star.jsonl'}")


if __name__ == "__main__":
    main()