"""Prompt trajectories and fixed FFFF targets for Predictor-v4 Stage 2."""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any

import torch
from safetensors import safe_open
from torch.utils.data import Dataset


class PredictorV4TrajectoryDataset(Dataset):
    """One deterministic seven-chunk trajectory per training prompt."""

    def __init__(
        self,
        cases_path: str | Path,
        *,
        data_root: str | Path,
        max_cases: int | None = None,
    ) -> None:
        self.cases_path = Path(cases_path).resolve()
        self.data_root = Path(data_root).resolve()
        with self.cases_path.open("r", encoding="utf-8") as handle:
            cases = [json.loads(line) for line in handle if line.strip()]
        if max_cases is not None:
            cases = cases[: int(max_cases)]
        if not cases:
            raise ValueError(f"No trajectory cases in {self.cases_path}")
        for position, case in enumerate(cases):
            if int(case["case_id"]) != position:
                raise ValueError("Stage-2 case_id values must be dense and ordered")
            if int(case.get("seed", -1)) != 0:
                raise ValueError(f"case {position} does not use latent seed 0")
            if not str(case.get("prompt", "")).strip():
                raise ValueError(f"case {position} has an empty prompt")
        self.cases = cases

    def __len__(self) -> int:
        return len(self.cases)

    def __getitem__(self, index: int) -> dict[str, Any]:
        case = self.cases[index]
        return {
            "case_id": int(case["case_id"]),
            "prompt": str(case["prompt"]),
            "seed": int(case["seed"]),
        }

    def offline_step_path(self, case_id: int, chunk_id: int) -> Path:
        path = (
            self.data_root
            / "steps"
            / f"case_{int(case_id):06d}"
            / f"chunk_{int(chunk_id):02d}.safetensors"
        )
        if not path.is_file():
            raise FileNotFoundError(path)
        return path


def trajectory_collate(items: list[dict[str, Any]]) -> dict[str, Any]:
    if len(items) != 1:
        raise ValueError(
            "Stage-2 rollout currently requires one trajectory per rank"
        )
    return items[0]


def load_offline_ffff_target(
    path: str | Path,
    *,
    step_id: int,
    device: torch.device,
) -> dict[str, torch.Tensor]:
    """Load one immutable Full-trajectory hidden/flow target."""

    path = Path(path)
    names = {
        "hidden": f"step_{int(step_id)}_final_hidden",
        "flow": f"step_{int(step_id)}_flow",
        "timestep": f"step_{int(step_id)}_timestep",
    }
    with safe_open(str(path), framework="pt", device="cpu") as handle:
        missing = set(names.values()).difference(handle.keys())
        if missing:
            raise KeyError(f"{path} lacks fixed FFFF tensors {sorted(missing)}")
        result = {key: handle.get_tensor(name) for key, name in names.items()}
    if result["hidden"].dtype != torch.bfloat16:
        raise TypeError(f"{path}: target hidden must be BF16")
    if result["flow"].dtype != torch.bfloat16:
        raise TypeError(f"{path}: target flow must be BF16")
    if result["timestep"].dtype != torch.int64:
        raise TypeError(f"{path}: legacy offline timestep must be INT64")
    return {
        key: value.to(device=device, non_blocking=True)
        for key, value in result.items()
    }


__all__ = [
    "PredictorV4TrajectoryDataset",
    "load_offline_ffff_target",
    "trajectory_collate",
]