#!/usr/bin/env python3
"""Thin runtime adapter for the locked Florence-Bridge FrameSkip continuation."""

from __future__ import annotations

import argparse
import json
import os
from dataclasses import asdict
from pathlib import Path

import numpy as np
from omegaconf import OmegaConf


def load_cache(path: str) -> tuple[dict[tuple[str, str], np.ndarray], dict]:
    payload = json.loads(Path(path).read_text(encoding="utf-8"))
    if payload.get("method") != "FrameSkip-adapted (Gaussian-TPI fallback)":
        raise RuntimeError("unexpected FrameSkip cache method")
    cache = {}
    for record in payload["records"]:
        key = (record["dataset_name"], str(record["trajectory_id"]))
        keep = np.asarray(record["keep_indices"], dtype=np.int64)
        if len(keep) == 0 or np.any(np.diff(keep) <= 0):
            raise RuntimeError(f"invalid retained indices for {key}")
        cache[key] = keep
    return cache, payload


def install_adapter(cache_path: str, start_global_draw_id: int) -> None:
    from starVLA.dataloader.gr00t_lerobot.datasets import LeRobotMixtureDataset

    cache, cache_payload = load_cache(cache_path)
    original_getitem = LeRobotMixtureDataset.__getitem__
    global_microbatch_size = 8
    period = 6

    def frameskip_getitem(self, index):
        trace = self.sample_trace(index)
        offset = trace.global_draw_id - start_global_draw_id
        if offset < 0:
            raise RuntimeError("FrameSkip draw precedes the certified continuation start")
        microbatch = offset // global_microbatch_size
        is_pruned = microbatch % period < 5
        logical_index = int(trace.base_index)
        actual_index = logical_index
        if is_pruned:
            key = (trace.dataset_name, str(trace.trajectory_id))
            if key not in cache:
                raise RuntimeError(f"FrameSkip cache is missing required trajectory {key}")
            retained = cache[key]
            position = int(np.searchsorted(retained, logical_index))
            position = int(np.clip(position, 0, len(retained) - 1))
            actual_index = int(retained[position])

        dataset = self.datasets[trace.dataset_index]
        original_get_step_data = dataset.get_step_data

        def redirected_get_step_data(trajectory_id, base_index):
            if str(trajectory_id) != str(trace.trajectory_id) or int(base_index) != logical_index:
                raise RuntimeError("unexpected nested dataset read during FrameSkip remapping")
            return original_get_step_data(trajectory_id, actual_index)

        dataset.get_step_data = redirected_get_step_data
        try:
            sample = original_getitem(self, index)
        finally:
            dataset.get_step_data = original_get_step_data
        sample["sampling_trace"]["frameskip"] = {
            "method": cache_payload["method"],
            "cache_record_count": len(cache_payload["records"]),
            "is_pruned_minibatch": is_pruned,
            "logical_base_index": logical_index,
            "physical_base_index": actual_index,
            "retention_ratio": 0.2 if is_pruned else 1.0,
        }
        return sample

    LeRobotMixtureDataset.__getitem__ = frameskip_getitem


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--config_yaml", required=True)
    args, clipargs = parser.parse_known_args()
    cache_path = os.environ.get("FRAMESKIP_CACHE")
    if not cache_path:
        raise RuntimeError("FRAMESKIP_CACHE is required")
    start_draw = int(os.environ.get("FRAMESKIP_START_GLOBAL_DRAW_ID", "1280000"))
    install_adapter(cache_path, start_draw)

    from starVLA.model.framework.share_tools import apply_config_compat
    from starVLA.training import train_starvla
    from starVLA.training.trainer_utils.trainer_tools import normalize_dotlist_args

    cfg = OmegaConf.load(args.config_yaml)
    cfg = OmegaConf.merge(cfg, OmegaConf.from_dotlist(normalize_dotlist_args(clipargs)))
    cfg = apply_config_compat(cfg, strict=True)
    cfg.config_yaml = args.config_yaml
    train_starvla.main(cfg)


if __name__ == "__main__":
    main()