from __future__ import annotations

import argparse
import json
import os
import sys
from pathlib import Path
from types import ModuleType
from typing import Any

EVAL_DIR = Path(__file__).resolve().parent
if str(EVAL_DIR) not in sys.path:
    sys.path.insert(0, str(EVAL_DIR))

from common.io import write_jsonl  # noqa: E402
from common.loader import SUPPORTED, load_bench  # noqa: E402
from common.paths import ensure_sys_path, resolve_out_dir  # noqa: E402

REPO = ensure_sys_path()  # put repo root on sys.path so examples / deepinfer can be imported

from examples._config import DEFAULT_MODEL_NAME  # noqa: E402

BENCH_DEFAULTS = {
    "aime": {"n_sample": 32, "temperature": 1.0, "top_p": 1.0},
    "hle": {"n_sample": 1, "temperature": 1.0, "top_p": 1.0},
    "gpqa": {"n_sample": 8, "temperature": 1.0, "top_p": 1.0},
}


_REASONING_EFFORT_LEVELS = ("low", "high", "max")


def _reasoning_effort(value: str) -> int | str:
    """Accept an integer 1-100, or a named level low/high/max."""
    if value in _REASONING_EFFORT_LEVELS:
        return value
    try:
        effort = int(value)
    except ValueError:
        raise argparse.ArgumentTypeError(f"expected an integer 1-100 or one of {_REASONING_EFFORT_LEVELS}, got {value!r}") from None
    if not 1 <= effort <= 100:
        raise argparse.ArgumentTypeError(f"expected 1-100, got {effort}")
    return effort


parser = argparse.ArgumentParser(description="eval benchmark")
# benchmark setting
parser.add_argument("--bench", default="aime", choices=list(SUPPORTED))
parser.add_argument("--model-name", default=DEFAULT_MODEL_NAME, help="HF model id or HF ckpt directory")
parser.add_argument("--data-file", required=True, help="problem jsonl")
parser.add_argument("--limit", type=int, default=None, help="run only the first N problems")
parser.add_argument("--out-dir", default=None)
parser.add_argument("--tag", default=None, help="output dir suffix, to compare settings of the same model")
# sample param
parser.add_argument("--n-sample", type=int, default=None, help="samples per problem (default: aime 32 / hle 1 / gpqa 8)")
parser.add_argument("--temperature", type=float, default=None, help="default 1.0 for every benchmark")
parser.add_argument("--top-p", type=float, default=None, help="default 1.0")
parser.add_argument(
    "--reasoning-effort", type=_reasoning_effort, default=100, help="integer 1-100, or low/high/max; higher means more thinking"
)
parser.add_argument("--max-len", type=int, default=1024 * 144, help="max prompt + output length per request")
# engine
parser.add_argument("--ep", type=int, default=8, help="expert parallel size = total GPUs in this job")
parser.add_argument("--cp", type=int, default=1)
parser.add_argument("--dp", type=int, default=None, help="defaults to ep // cp")
parser.add_argument("--max-num-seqs", type=int, default=384)
parser.add_argument("--max-num-batched-tokens", type=int, default=32768)
parser.add_argument("--gpu-memory-utilization", type=float, default=0.9)
parser.add_argument(
    "--env", action="append", default=[], metavar="KEY=VALUE", help="extra env vars, repeatable; applied before the engine is built"
)
parser.add_argument("--decode-blacklist", default="0,2,129262..129264", help="tokens blocked during generation; empty string disables it")
parser.add_argument(
    "--batch-invariant", action="store_true", help="DEEPINFER_BATCH_INVARIANT=1: results independent of batch composition (slower)"
)

args = parser.parse_args()

# fill in per-benchmark defaults (only for args not passed explicitly)
for _key, _default in BENCH_DEFAULTS[args.bench].items():
    if getattr(args, _key) is None:
        setattr(args, _key, _default)

REASONING_EFFORT = args.reasoning_effort

OUT_DIR = resolve_out_dir(
    args.out_dir,
    bench=args.bench,
    tag=args.tag or f"{args.model_name.replace('/', '_')}-n{args.n_sample}-j{args.reasoning_effort}",
)
META_PATH = os.path.join(OUT_DIR, "meta.json")


# --------------------------------------------------------------------------- #
# generation
# --------------------------------------------------------------------------- #


def _render_prompts(encoder: ModuleType, records: list[dict]) -> list[str]:
    return [
        encoder.encode_messages(
            rec["messages"],
            thinking_mode="thinking",
            reasoning_effort=REASONING_EFFORT,
        )
        for rec in records
    ]


def _multi_modal_entry(record: dict) -> dict:
    """One multi_modal_data entry: image bytes plus dimensions."""
    return {
        "image_bytes": [Path(path).read_bytes() for path in record.get("image_paths") or []],
        "image_sizes": [tuple(size) for size in record.get("image_sizes") or []],
    }


def _setup_env() -> None:
    # MTP / DFlash speculative decoding, set before setup_common_env (which only setdefaults it).
    os.environ["DEEPINFER_USE_NEXTN_DECODE"] = "5"
    os.environ["DEEPINFER_ENABLE_MTP_VARLEN"] = "1"

    from examples._config import setup_common_env

    setup_common_env(model_name=args.model_name, sampler_force_topk=0, set_override_experts=True, enable_eplb=True)
    for item in args.env:
        key, _, value = item.partition("=")
        assert key and "=" in item, f"--env expects KEY=VALUE, got {item!r}"
        os.environ[key] = value
    if args.batch_invariant:
        os.environ["DEEPINFER_BATCH_INVARIANT"] = "1"
    if args.decode_blacklist:
        os.environ["DEEPINFER_DECODE_TOKEN_BLACKLIST"] = args.decode_blacklist


def _seed(global_index: int, sample_index: int) -> int:
    """Per-sample seed: identical seeds produce identical output, so every sample needs its own."""
    return (global_index * args.n_sample + sample_index) % (2**64)


def _rows_from_outputs(tokenizer: Any, plan: list, records: list[dict], outputs: list) -> list[dict]:
    rows = []
    for (problem, sample), generated in zip(plan, outputs, strict=True):
        assert generated, f"{records[problem]['task_id']}#{sample}: generated no tokens"
        text = tokenizer.decode([token.token_id for token in generated], skip_special_tokens=False, clean_up_tokenization_spaces=False)
        rows.append(
            {
                "task_id": records[problem]["task_id"],
                "sample_idx": sample,
                "seq_seed": _seed(problem, sample),
                "prompt_len": generated[0].prompt_len,
                "num_output_tokens": len(generated),
                "finish_reason": generated[-1].finish_reason,
                "raw_text": text,
            }
        )
    return rows


def run_rank(dp_size: int, dp_rank: int) -> None:
    _setup_env()

    from common.encoder import load_encoder

    from examples._config import create_llm

    encoder = load_encoder(args.model_name)

    bench = load_bench(args.bench)
    problems = bench.load(args.data_file)
    records = bench.build_records(problems, limit=args.limit)

    # Shard by sample so every rank gets an equal share and no rank inherits a problem's long tail.
    plan = [(problem, sample) for problem in range(len(records)) for sample in range(args.n_sample)][dp_rank::dp_size]

    llm = create_llm(
        args.model_name,
        dp_rank,
        ep_size=args.ep,
        cp_size=args.cp,
        data_parallel_size=dp_size,
        max_num_batched_tokens=args.max_num_batched_tokens,
        # The engine's runtime_max_model_len is max_model_len + 1, and every request's
        # prompt + max_tokens is already clamped to <= max_len, so a tiny margin is enough.
        max_model_len=args.max_len + 64,
        max_num_seqs=args.max_num_seqs,
        gpu_memory_utilization=args.gpu_memory_utilization,
        lookahead_tokens=min(1, int(os.environ.get("DEEPINFER_USE_NEXTN_DECODE", "0"))),
        decode_sequence=False,
    )

    try:
        # Import after create_llm: deepinfer imports torch, which must not happen before the engine
        # has set up the accelerator devices.
        from deepinfer.sampling_params import SamplingParams

        tokenizer = llm.get_tokenizer()

        if not plan:
            # dp_size > number of requests leaves some ranks with an empty shard: write an empty
            # file and stop, do not build a batch.
            os.makedirs(OUT_DIR, exist_ok=True)
            write_jsonl(os.path.join(OUT_DIR, f"raw_rank{dp_rank:03d}.jsonl"), [])
            print(f"[rank {dp_rank}] empty shard (dp_size={dp_size} > request count), skipping", flush=True)
            return

        # Render and tokenize each needed prompt once: a rank may use the same prompt for several
        # samples of the same problem.
        needed = sorted({problem for problem, _ in plan})
        texts = dict(zip(needed, _render_prompts(encoder, [records[problem] for problem in needed]), strict=True))
        has_images = any(records[problem].get("image_paths") for problem in needed)
        multi_modal: dict[int, dict] = {}
        if has_images:
            # Replace the <image> placeholder with the image token the engine understands. Image
            # questions then take the string path and the engine expands the token.
            from examples._prompts import get_engine_image_token, replace_image_tokens

            image_token_str, _ = get_engine_image_token(llm)
            texts = dict(zip(needed, replace_image_tokens([texts[problem] for problem in needed], image_token_str), strict=True))
            # The engine consumes multi_modal_data in place, so every request gets its own copy.
            multi_modal = {problem: _multi_modal_entry(records[problem]) for problem in needed}
        token_ids = {problem: tokenizer.encode(text, add_special_tokens=False) for problem, text in texts.items()}

        # max_tokens is clamped to the remaining budget: max_len - prompt_len, at least 1.
        sampling_params = [
            SamplingParams(
                temperature=args.temperature,
                top_p=args.top_p,
                max_tokens=max(args.max_len - len(token_ids[problem]), 1),
                ignore_eos=False,
                seq_seed=_seed(problem, sample),
            )
            for problem, sample in plan
        ]

        lens = [len(token_ids[problem]) for problem, _ in plan]
        print(
            f"[rank {dp_rank}] problems={len(needed)} requests={len(plan)} "
            f"prompt_tokens[min/avg/max]={min(lens)}/{sum(lens) // len(lens)}/{max(lens)}"
            + (f" multimodal={len(plan)}" if has_images else ""),
            flush=True,
        )

        if has_images:
            outputs = llm.generate(
                prompts=[texts[problem] for problem, _ in plan],
                multi_modal_data=[dict(multi_modal[problem]) for problem, _ in plan],
                sampling_params=sampling_params,
                use_tqdm=False,
            )
        else:
            outputs = llm.generate(
                prompts=None,
                prompt_token_ids=[token_ids[problem] for problem, _ in plan],
                sampling_params=sampling_params,
                use_tqdm=False,
            )

        rows = _rows_from_outputs(tokenizer, plan, records, outputs)

        os.makedirs(OUT_DIR, exist_ok=True)
        path = os.path.join(OUT_DIR, f"raw_rank{dp_rank:03d}.jsonl")
        write_jsonl(path, rows)
        print(f"[rank {dp_rank}] wrote {path} rows={len(rows)}", flush=True)
    finally:
        llm.close()


def main() -> int:
    dp_size = max(args.ep // args.cp, 1) if args.dp is None else args.dp
    os.makedirs(OUT_DIR, exist_ok=True)

    meta = {
        "bench": args.bench,
        "model_name": args.model_name,
        "ep": args.ep,
        "cp": args.cp,
        "dp": dp_size,
        "data_file": args.data_file,
        "limit": args.limit,
        "n_sample": args.n_sample,
        "temperature": args.temperature,
        "top_p": args.top_p,
        "reasoning_effort": REASONING_EFFORT,
        "max_len": args.max_len,
        "mtp": 5,
        "mtp_varlen": True,
        "batch_invariant": args.batch_invariant,
        "decode_blacklist": args.decode_blacklist,
        "max_num_seqs": args.max_num_seqs,
        "max_num_batched_tokens": args.max_num_batched_tokens,
        "env": args.env,
    }

    for old in Path(OUT_DIR).glob("raw_rank*.jsonl"):
        old.unlink()
    with open(META_PATH, "w") as f:
        json.dump(meta, f, indent=2, ensure_ascii=False)

    print(f"REPO={REPO}\nOUT_DIR={OUT_DIR}\nmeta={json.dumps(meta, ensure_ascii=False)}", flush=True)

    from examples._launch import launch_dp

    launch_dp(run_rank, dp_size)
    return 0


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