"""Generate and score Qwen3.5 checkpoints on official LiveCodeBench v6.

This runner deliberately reuses LiveCodeBench's dataset objects, code
extraction, and executable-code evaluator, while rendering prompts with the
checkpoint's own Qwen3.5 tokenizer.  The upstream runner still hard-codes a
Qwen1.5 tokenizer for its Qwen prompt style and therefore cannot safely render
Qwen3.5's explicit non-thinking template.

The official repository currently requires ``datasets==3.6.0`` because its
dataset is implemented as a loading script.  Keep that dependency in an
isolated environment; do not downgrade the RL training environment.
"""

from __future__ import annotations

import argparse
import gc
import hashlib
import json
import os
import subprocess
import sys
from pathlib import Path
from typing import Any


SCHEMA = "livecodebench_qwen35_v1"


def _sha256_text(text: str) -> str:
    return hashlib.sha256(text.encode("utf-8")).hexdigest()


def _write_json(path: Path, value: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + ".tmp")
    with temporary.open("w") as handle:
        json.dump(value, handle, indent=2, sort_keys=True)
        handle.write("\n")
        handle.flush()
        os.fsync(handle.fileno())
    os.replace(temporary, path)


def _append_jsonl(handle, value: Any) -> None:
    handle.write(json.dumps(value, ensure_ascii=True) + "\n")
    handle.flush()
    os.fsync(handle.fileno())


def _read_jsonl(path: Path) -> list[dict[str, Any]]:
    if not path.exists():
        return []
    rows = []
    with path.open() as handle:
        for line_no, line in enumerate(handle, 1):
            try:
                rows.append(json.loads(line))
            except json.JSONDecodeError as exc:
                raise ValueError(f"invalid JSON at {path}:{line_no}") from exc
    return rows


def _load_lcb(lcb_root: Path, release_version: str):
    sys.path.insert(0, str(lcb_root))
    from lcb_runner.benchmarks import load_code_generation_dataset

    benchmark = load_code_generation_dataset(release_version)
    return sorted(benchmark, key=lambda row: row.question_id)


def _lcb_commit(lcb_root: Path) -> str:
    return subprocess.check_output(
        ["git", "-c", f"safe.directory={lcb_root}", "rev-parse", "HEAD"],
        cwd=lcb_root,
        text=True,
    ).strip()


def _format_prompts(benchmark, tokenizer) -> list[str]:
    from lcb_runner.prompts.code_generation import (
        PromptConstants,
        get_generic_question_template_answer,
    )

    prompts = []
    for problem in benchmark:
        messages = [
            {"role": "system", "content": PromptConstants.SYSTEM_MESSAGE_GENERIC},
            {"role": "user", "content": get_generic_question_template_answer(problem)},
        ]
        prompts.append(
            tokenizer.apply_chat_template(
                messages,
                tokenize=False,
                add_generation_prompt=True,
                enable_thinking=False,
            )
        )
    return prompts


def _config(args, lcb_commit: str, benchmark_size: int) -> dict[str, Any]:
    return {
        "schema": SCHEMA,
        "model": str(Path(args.model).resolve()),
        "model_label": args.model_label,
        "release_version": args.release_version,
        "limit": args.limit,
        "benchmark_size": benchmark_size,
        "lcb_commit": lcb_commit,
        "thinking": False,
        "n": args.n,
        "temperature": args.temperature,
        "top_p": args.top_p,
        "max_tokens": args.max_tokens,
        "max_model_len": args.max_model_len,
        "seed": args.seed,
        "stop": args.stop,
    }


def generate(args, benchmark, lcb_commit: str) -> None:
    from transformers import AutoTokenizer
    from vllm import LLM, SamplingParams

    tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
    prompts = _format_prompts(benchmark, tokenizer)
    prompt_lengths = [
        len(tokenizer(prompt, add_special_tokens=False)["input_ids"]) for prompt in prompts
    ]
    budget = args.max_model_len - args.max_tokens
    overlong = [
        (problem.question_id, length)
        for problem, length in zip(benchmark, prompt_lengths)
        if length > budget
    ]
    if overlong:
        raise ValueError(
            f"{len(overlong)} LiveCodeBench prompts exceed the {budget}-token prompt "
            f"budget; refusing to truncate. First rows: {overlong[:5]}"
        )

    output = Path(args.output)
    manifest_path = Path(args.manifest)
    config = _config(args, lcb_commit, len(benchmark))
    config["prompt_tokens"] = {
        "minimum": min(prompt_lengths),
        "maximum": max(prompt_lengths),
        "mean": sum(prompt_lengths) / len(prompt_lengths),
    }
    if manifest_path.exists():
        existing_manifest = json.loads(manifest_path.read_text())
        comparable = {key: existing_manifest.get(key) for key in config if key != "prompt_tokens"}
        expected = {key: value for key, value in config.items() if key != "prompt_tokens"}
        if comparable != expected:
            raise ValueError("existing LiveCodeBench manifest does not match this run")
    else:
        _write_json(manifest_path, config)

    existing = _read_jsonl(output) if args.resume else []
    if output.exists() and not args.resume:
        output.unlink()
    by_id = {row["question_id"]: row for row in existing}
    unknown = set(by_id) - {row.question_id for row in benchmark}
    if unknown:
        raise ValueError(f"output contains unknown question IDs: {sorted(unknown)[:5]}")

    llm = LLM(
        model=args.model,
        tokenizer=args.model,
        trust_remote_code=True,
        language_model_only=True,
        max_model_len=args.max_model_len,
        gpu_memory_utilization=args.gpu_memory_utilization,
        seed=args.seed,
    )
    sampling = SamplingParams(
        n=args.n,
        max_tokens=args.max_tokens,
        temperature=args.temperature,
        top_p=args.top_p,
        stop=[args.stop] if args.stop else None,
        seed=args.seed,
    )

    output.parent.mkdir(parents=True, exist_ok=True)
    with output.open("a") as handle:
        for start in range(0, len(benchmark), args.batch_size):
            problems = benchmark[start : start + args.batch_size]
            batch_prompts = prompts[start : start + args.batch_size]
            missing = [index for index, problem in enumerate(problems) if problem.question_id not in by_id]
            if not missing:
                continue
            generated = llm.generate([batch_prompts[index] for index in missing], sampling)
            for index, request_output in zip(missing, generated):
                problem = problems[index]
                candidates = []
                for candidate in request_output.outputs:
                    candidates.append(
                        {
                            "text": candidate.text,
                            "token_count": len(candidate.token_ids),
                            "finish_reason": candidate.finish_reason,
                            "probably_truncated": (
                                candidate.finish_reason == "length"
                                or len(candidate.token_ids) >= args.max_tokens - args.truncation_buffer_tokens
                            ),
                        }
                    )
                row = {
                    "schema": SCHEMA,
                    "question_id": problem.question_id,
                    "prompt_sha256": _sha256_text(batch_prompts[index]),
                    "prompt_tokens": prompt_lengths[start + index],
                    "outputs": candidates,
                }
                _append_jsonl(handle, row)
                by_id[problem.question_id] = row
            print(f"[lcb] generated {len(by_id)}/{len(benchmark)} problems", flush=True)

    del llm
    gc.collect()


def score(args, benchmark, lcb_commit: str) -> None:
    from lcb_runner.evaluation import codegen_metrics
    from lcb_runner.lm_styles import LMStyle
    from lcb_runner.utils.extraction_utils import extract_code

    rows = _read_jsonl(Path(args.output))
    by_id = {row["question_id"]: row for row in rows}
    missing = [row.question_id for row in benchmark if row.question_id not in by_id]
    if missing:
        raise ValueError(f"generation output is incomplete: missing {len(missing)} problems")
    generations = []
    for problem in benchmark:
        outputs = by_id[problem.question_id]["outputs"]
        if len(outputs) != args.n:
            raise ValueError(f"{problem.question_id} has {len(outputs)} outputs, expected {args.n}")
        generations.append(
            [extract_code(candidate["text"], LMStyle.CodeQwenInstruct) for candidate in outputs]
        )
    samples = [problem.get_evaluation_sample() for problem in benchmark]
    metrics, results, metadata = codegen_metrics(
        samples,
        generations,
        k_list=[1, 5, 10],
        num_process_evaluate=args.num_process_evaluate,
        timeout=args.timeout,
    )
    cap_hits = sum(
        candidate["probably_truncated"]
        for row in rows
        for candidate in row["outputs"]
    )
    total = sum(len(row["outputs"]) for row in rows)
    summary = {
        **_config(args, lcb_commit, len(benchmark)),
        "metrics": metrics,
        "truncated_generations": cap_hits,
        "total_generations": total,
        "truncation_rate": cap_hits / total,
        "per_problem_results": {str(key): value for key, value in results.items()},
        "evaluator_metadata": metadata,
    }
    _write_json(Path(args.summary), summary)
    print(json.dumps({"metrics": metrics, "truncation_rate": cap_hits / total}, indent=2))


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--lcb-root", required=True)
    parser.add_argument("--model", required=True, help="A base model or already-merged checkpoint")
    parser.add_argument("--model-label", required=True)
    parser.add_argument("--output", required=True)
    parser.add_argument("--manifest", required=True)
    parser.add_argument("--summary", required=True)
    parser.add_argument("--mode", choices=("generate", "score", "both"), default="both")
    parser.add_argument("--release-version", default="release_v6")
    parser.add_argument(
        "--limit",
        type=int,
        default=0,
        help="Evaluate only the first N sorted problems (0 means the full release).",
    )
    parser.add_argument("--n", type=int, default=10)
    parser.add_argument("--temperature", type=float, default=0.2)
    parser.add_argument("--top-p", type=float, default=0.95)
    parser.add_argument("--max-tokens", type=int, default=32768)
    parser.add_argument("--max-model-len", type=int, default=36864)
    parser.add_argument("--truncation-buffer-tokens", type=int, default=24)
    parser.add_argument("--stop", default="###")
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--batch-size", type=int, default=8)
    parser.add_argument("--gpu-memory-utilization", type=float, default=0.9)
    parser.add_argument("--num-process-evaluate", type=int, default=12)
    parser.add_argument("--timeout", type=int, default=6)
    parser.add_argument("--resume", action="store_true")
    args = parser.parse_args()
    if args.max_model_len <= args.max_tokens:
        parser.error("--max-model-len must exceed --max-tokens")
    if args.n < 1 or args.batch_size < 1 or args.limit < 0:
        parser.error("--n and --batch-size must be positive; --limit must be non-negative")
    return args


def main() -> None:
    args = parse_args()
    lcb_root = Path(args.lcb_root).resolve()
    # Upstream prompt modules load few-shot fixtures relative to the process
    # working directory. Resolve all of our paths first, then enter the pinned
    # checkout so those official assets are found regardless of the caller's
    # cwd.
    args.model = str(Path(args.model).resolve())
    args.output = str(Path(args.output).resolve())
    args.manifest = str(Path(args.manifest).resolve())
    args.summary = str(Path(args.summary).resolve())
    os.chdir(lcb_root)
    benchmark = _load_lcb(lcb_root, args.release_version)
    if args.limit:
        benchmark = benchmark[: args.limit]
    commit = _lcb_commit(lcb_root)
    if args.mode in ("generate", "both"):
        generate(args, benchmark, commit)
    if args.mode in ("score", "both"):
        score(args, benchmark, commit)


if __name__ == "__main__":
    main()