"""
Second pass over segments whose greedy translation failed validation.

Two recovery strategies, tried in order, keeping the first candidate that validates:
  1. beam search (beam_width=4) — what the Seed-X authors recommend
  2. sampled best-of-8 — a different part of the distribution when beam search repeats the error

Recovered translations are appended to the cache, overriding the greedy result.
"""
import json
import os
import sys
from pathlib import Path

sys.path.insert(0, ".")
from build_dataset import check
from gsm_common import SRC_PROMPT

OUT = Path("out_gsm")
CACHE = OUT / "translations.jsonl"
MODEL = "./models/Seed-X-PPO-7B"


def load():
    trans = {}
    with open(CACHE, encoding="utf-8") as fh:
        for line in fh:
            try:
                r = json.loads(line)
            except json.JSONDecodeError:
                continue
            trans[r["src"]] = r["tgt"]
    rows = [json.loads(l) for l in open(OUT / "selected_rows.jsonl", encoding="utf-8")]
    kind = {}
    for r in rows:
        kind.setdefault(r["question"], "question")
        kind.setdefault(r["thinking"], "thinking")
    return trans, kind


def main():
    trans, kind = load()
    failed = [
        s for s, t in trans.items()
        if check(s, t, strict=(kind.get(s) == "thinking")) is not None
    ]
    print(f"[*] {len(failed)}/{len(trans)} segments failed validation ({len(failed)/len(trans):.2%})")
    if not failed:
        return

    from vllm import LLM, SamplingParams
    from vllm.sampling_params import BeamSearchParams

    llm = LLM(model=MODEL, max_num_seqs=256, gpu_memory_utilization=0.92, max_model_len=1024)
    prompts = [SRC_PROMPT.format(text=s) for s in failed]
    recovered = {}

    # Beam search is disabled by default: vLLM 0.8.5 runs it as a Python-level loop and it took
    # >50 min on ~4.8k segments without finishing, versus ~2 min for the batched sampling path
    # below. Set RETRY_BEAM=1 to use it anyway.
    if os.environ.get("RETRY_BEAM") == "1":
        print("[*] pass 1: beam search")
        outs = llm.beam_search([{"prompt": p} for p in prompts], BeamSearchParams(beam_width=4, max_tokens=256))
        still = []
        for src, o in zip(failed, outs):
            strict = kind.get(src) == "thinking"
            for seq in o.sequences:
                cand = seq.text.strip()
                if check(src, cand, strict=strict) is None:
                    recovered[src] = cand
                    break
            else:
                still.append(src)
        print(f"    recovered {len(recovered)}, still failing {len(still)}")
    else:
        print("[*] pass 1: skipped (beam search disabled)")
        still = list(failed)

    if still:
        print("[*] pass 2: sampled best-of-8")
        params = SamplingParams(n=8, temperature=0.8, top_p=0.95, max_tokens=256, skip_special_tokens=True)
        outs = llm.generate([SRC_PROMPT.format(text=s) for s in still], params)
        final = []
        for src, o in zip(still, outs):
            strict = kind.get(src) == "thinking"
            for cand in o.outputs:
                text = cand.text.strip()
                if check(src, text, strict=strict) is None:
                    recovered[src] = text
                    break
            else:
                final.append(src)
        print(f"    recovered {len(recovered)} total, unrecoverable {len(final)}")

    with open(CACHE, "a", encoding="utf-8") as fh:
        for src, tgt in recovered.items():
            fh.write(json.dumps({"src": src, "tgt": tgt}, ensure_ascii=False) + "\n")
    print(f"[+] appended {len(recovered)} recovered translations to {CACHE}")


if __name__ == "__main__":
    main()