"""Multi-GPU closed-loop ALFWorld evaluation from immutable best checkpoints."""
from __future__ import annotations

import argparse
from collections import Counter, defaultdict
import heapq
import json
import os
from pathlib import Path
import random
import shutil
import sys
import uuid

PROJECT = Path("/share/shenmingxiang/closing-square")
WS = PROJECT / "verl_workspace"
sys.path.insert(0, str(WS / "third_party" / "Agent-R1"))
sys.path.insert(0, str(PROJECT))

import pandas as pd
import torch
import torch.distributed as dist
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

from recipes.alfworld.env.alfworld_wrapper import AlfworldTextworldEnv
from src.prompts import student_prompt, teacher_prompt
from src.prepare_structured_hints import build_hint_fields, render_hint, _load_traj

DEFAULT_MODEL = PROJECT / "models" / "Qwen2.5-3B-Instruct"
DATA_ROOT = WS / "data" / "alfworld"


def set_seed(seed: int) -> None:
    random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


@torch.no_grad()
def constrained_action(model, tokenizer, prompt: str, actions: list[str], device, max_prompt_length: int) -> str:
    old_side = tokenizer.truncation_side
    tokenizer.truncation_side = "left"
    encoded = tokenizer(
        prompt, return_tensors="pt", add_special_tokens=True,
        truncation=True, max_length=max_prompt_length,
    )
    tokenizer.truncation_side = old_side
    input_ids = encoded["input_ids"].to(device)
    attention_mask = encoded["attention_mask"].to(device)
    prompt_length = input_ids.shape[1]
    token_to_action: dict[tuple[int, ...], str] = {}
    sequences: list[tuple[int, ...]] = []
    for action in actions:
        tokens = tuple(tokenizer(action, add_special_tokens=False)["input_ids"])
        token_to_action.setdefault(tokens, action)
        sequences.append(tokens + (tokenizer.eos_token_id,))
    max_new_tokens = max(len(sequence) for sequence in sequences)

    def allowed_tokens(_batch_id: int, generated: torch.Tensor) -> list[int]:
        prefix = tuple(generated[prompt_length:].tolist())
        allowed = {
            sequence[len(prefix)]
            for sequence in sequences
            if len(sequence) > len(prefix) and sequence[: len(prefix)] == prefix
        }
        return sorted(allowed) if allowed else [tokenizer.eos_token_id]

    generated = model.generate(
        input_ids=input_ids,
        attention_mask=attention_mask,
        do_sample=False,
        temperature=1.0,
        top_p=1.0,
        top_k=50,
        repetition_penalty=1.0,
        num_beams=1,
        max_new_tokens=max_new_tokens,
        eos_token_id=tokenizer.eos_token_id,
        pad_token_id=tokenizer.pad_token_id,
        prefix_allowed_tokens_fn=allowed_tokens,
        use_cache=True,
    )
    continuation = generated[0, prompt_length:].tolist()
    if tokenizer.eos_token_id in continuation:
        continuation = continuation[: continuation.index(tokenizer.eos_token_id)]
    key = tuple(continuation)
    if key not in token_to_action:
        raise RuntimeError(f"constrained decoding returned unknown token path: {continuation}")
    return token_to_action[key]


def run_episode(
    model,
    tokenizer,
    env,
    item,
    device,
    max_steps: int,
    max_prompt_length: int,
    prompt_mode: str,
):
    obs, info = env.reset_with_info(item["game"])
    hint = item.get("H", "")
    if prompt_mode == "teacher" and not hint:
        expert_policy = [str(value) for value in info.get("policy_commands", [])]
        hint = render_hint(
            build_hint_fields(item["task_family"], _load_traj(item["game"]), expert_policy)
        )
    history: list[str] = []
    done = False
    steps = 0
    repeated = 0
    last_action = None
    trajectory = []
    for _ in range(max_steps):
        admissible = [str(value) for value in info.get("admissible_commands", [])]
        if done or not admissible:
            break
        if prompt_mode == "teacher":
            prompt = teacher_prompt(item["goal"], str(obs), admissible, hint, history)
        else:
            prompt = student_prompt(item["goal"], str(obs), admissible, history)
        action = constrained_action(model, tokenizer, prompt, admissible, device, max_prompt_length)
        repeated += int(action == last_action)
        next_obs, _, done, info = env.step(action)
        trajectory.append(
            {
                "step": steps,
                "observation": str(obs),
                "admissible": admissible,
                "action": action,
                "response": str(next_obs),
                "success": bool(info.get("success", False)) if info else False,
            }
        )
        history.append(f"Action taken: {action}\nEnvironment response: {next_obs}")
        obs = next_obs
        last_action = action
        steps += 1
    return {
        "success": bool(info.get("success", False)) if info else False,
        "steps": steps,
        "repeated_actions": repeated,
        "H": hint or None,
        "trajectory": trajectory,
    }


def atomic_json(path: Path, value, token: str) -> None:
    temp = path.parent / f".{path.name}.tmp.{token}"
    with temp.open("w", encoding="utf-8") as handle:
        json.dump(value, handle, ensure_ascii=False, indent=2)
        handle.flush()
        os.fsync(handle.fileno())
    os.replace(temp, path)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--name", required=True)
    parser.add_argument("--adapter", default="none")
    parser.add_argument("--model", type=Path, default=DEFAULT_MODEL)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--max_steps", type=int, default=50)
    parser.add_argument("--max_prompt_length", type=int, default=1536)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--limit_per_split", type=int, default=0)
    parser.add_argument("--limit_per_family", type=int, default=0)
    parser.add_argument(
        "--dataset_mode", choices=("standard", "pilot", "sft_val"), default="standard"
    )
    parser.add_argument("--prompt_mode", choices=("student", "teacher"), default="student")
    parser.add_argument(
        "--hint_data",
        type=Path,
        default=WS / "data" / "privileged_plan_sft.jsonl",
    )
    args = parser.parse_args()

    world_size = int(os.environ.get("WORLD_SIZE", "1"))
    distributed = world_size > 1
    if distributed:
        dist.init_process_group("nccl")
        rank = dist.get_rank()
        local_rank = int(os.environ["LOCAL_RANK"])
        torch.cuda.set_device(local_rank)
    else:
        rank = 0
        local_rank = 0
    is_main = rank == 0
    device = torch.device("cuda", local_rank)
    set_seed(args.seed)

    items = []
    if args.dataset_mode == "sft_val":
        by_game = {}
        with args.hint_data.open("r", encoding="utf-8") as source:
            for line in source:
                row = json.loads(line)
                if row["sft_split"] != "val":
                    continue
                by_game.setdefault(
                    row["game_id"],
                    {
                        "split": "sft_val",
                        "game": row["game_id"],
                        "goal": row["goal"],
                        "task_family": row["task_family"],
                        "H": row["H"],
                    },
                )
        by_family = defaultdict(list)
        for item in by_game.values():
            by_family[item["task_family"]].append(item)
        rng = random.Random(args.seed)
        for family in sorted(by_family):
            family_items = sorted(by_family[family], key=lambda value: value["game"])
            rng.shuffle(family_items)
            if args.limit_per_family:
                family_items = family_items[: args.limit_per_family]
            items.extend(family_items)
    else:
        source_root = DATA_ROOT / "pilot" / "pick_and_place" if args.dataset_mode == "pilot" else DATA_ROOT
        for split in ("valid_seen", "valid_unseen"):
            frame = pd.read_parquet(source_root / f"{split}.parquet")
            if args.limit_per_split:
                frame = frame.head(args.limit_per_split)
            for _, row in frame.iterrows():
                extra = dict(row["extra_info"])
                items.append(
                    {
                        "split": split,
                        "game": str(extra["game_relative_path"]),
                        "goal": str(extra.get("goal_text", "")),
                        "task_family": str(extra["task_family"]),
                    }
                )
    indexed_items = [(index, item) for index, item in enumerate(items) if index % world_size == rank]
    if is_main:
        print(
            f"eval name={args.name} episodes={len(items)} world_size={world_size} "
            f"adapter={args.adapter}",
            flush=True,
        )

    tokenizer = AutoTokenizer.from_pretrained(args.model)
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "right"
    base = AutoModelForCausalLM.from_pretrained(
        args.model, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True
    ).to(device)
    if args.adapter.lower() != "none":
        adapter = Path(args.adapter).resolve()
        model = PeftModel.from_pretrained(base, adapter)
    else:
        adapter = None
        model = base
    model.eval().requires_grad_(False)
    model.config.use_cache = True

    token_holder = [f"{os.getpid()}.{uuid.uuid4().hex}" if is_main else ""]
    if distributed:
        dist.broadcast_object_list(token_holder, src=0)
    run_token = token_holder[0]
    part_dir = args.output.parent / f".{args.output.name}.parts.{run_token}"
    if is_main:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        part_dir.mkdir(parents=False, exist_ok=False)
    if distributed:
        dist.barrier()

    env = AlfworldTextworldEnv(data_root=str(DATA_ROOT), max_episode_steps=args.max_steps)
    temp_part = part_dir / f".rank_{rank:02d}.jsonl.tmp"
    final_part = part_dir / f"rank_{rank:02d}.jsonl"
    with temp_part.open("w", encoding="utf-8") as handle:
        for local_index, (global_index, item) in enumerate(indexed_items, start=1):
            try:
                result = run_episode(
                    model, tokenizer, env, item, device, args.max_steps,
                    args.max_prompt_length, args.prompt_mode,
                )
                error = None
            except Exception as exception:
                result = {
                    "success": False,
                    "steps": 0,
                    "repeated_actions": 0,
                    "H": None,
                    "trajectory": [],
                }
                error = f"{type(exception).__name__}: {exception}"
            finally:
                env._close_env()
            record = {"record_index": global_index, **item, **result, "error": error}
            handle.write(json.dumps(record, ensure_ascii=False) + "\n")
            if local_index % 5 == 0:
                print(f"eval rank={rank} completed={local_index}/{len(indexed_items)}", flush=True)
        handle.flush()
        os.fsync(handle.fileno())
    os.replace(temp_part, final_part)
    if distributed:
        dist.barrier()

    if is_main:
        def lines(path: Path):
            with path.open("r", encoding="utf-8") as source:
                for line in source:
                    value = json.loads(line)
                    yield value["record_index"], line, value

        streams = [lines(path) for path in sorted(part_dir.glob("rank_*.jsonl"))]
        temp_output = args.output.parent / f".{args.output.name}.tmp.{run_token}"
        merged_records = []
        with temp_output.open("w", encoding="utf-8") as output_handle:
            for _, line, value in heapq.merge(*streams, key=lambda triple: triple[0]):
                output_handle.write(line)
                merged_records.append(value)
            output_handle.flush()
            os.fsync(output_handle.fileno())
        os.replace(temp_output, args.output)

        groups = defaultdict(lambda: Counter(total=0, success=0, errors=0, steps=0, repeated=0))
        for record in merged_records:
            for key in (record["split"], f"{record['split']}::{record['task_family']}", "pooled"):
                groups[key]["total"] += 1
                groups[key]["success"] += int(record["success"])
                groups[key]["errors"] += int(record["error"] is not None)
                groups[key]["steps"] += int(record["steps"])
                groups[key]["repeated"] += int(record["repeated_actions"])
        summary_groups = {}
        for key, counts in sorted(groups.items()):
            total = max(counts["total"], 1)
            summary_groups[key] = {
                **dict(counts),
                "success_rate": counts["success"] / total,
                "mean_steps": counts["steps"] / total,
                "mean_repeated_actions": counts["repeated"] / total,
            }
        summary = {
            "name": args.name,
            "adapter": str(adapter) if adapter else None,
            "base_model": str(args.model.resolve()),
            "seed": args.seed,
            "dataset_mode": args.dataset_mode,
            "prompt_mode": args.prompt_mode,
            "decoding": "greedy token-trie constrained to admissible commands",
            "groups": summary_groups,
        }
        summary_path = args.output.with_suffix(args.output.suffix + ".summary.json")
        atomic_json(summary_path, summary, run_token)
        print(json.dumps(summary, ensure_ascii=False, indent=2), flush=True)
    if distributed:
        dist.barrier()
    if is_main:
        shutil.rmtree(part_dir)
    if distributed:
        dist.destroy_process_group()


if __name__ == "__main__":
    main()