"""Transactional FLock adapter for a self-contained SigLIP2 action-chunk policy.

The exporter copies this file to ``flock_robotics_adapter.py`` next to
``siglip2_action_chunk_model.py``, ``vla_config.json``,
``model.safetensors``, and the bundled SigLIP2 tokenizer assets. Runtime
loading is deliberately local-only: the complete SigLIP2 backbone, policy
head, and tokenizer must already be present in the submission directory.

The validator can retry the same ``policy.act(obs)`` request. State is
therefore copy-on-write and is committed only after a valid action has been
produced. An exact same-step retry returns the cached action without running
the model again.
"""

from __future__ import annotations

import hashlib
import json
import math
import numbers
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable, Mapping, Sequence

import numpy as np
import torch

ACTION_DIM = 7
DEFAULT_HISTORY = 4
DEFAULT_ACTION_CHUNK = 8
SIGLIP2_TEXT_MAX_LENGTH = 64
REQUIRED_TOKENIZER_ASSETS = (
    "tokenizer.model",
    "tokenizer.json",
    "tokenizer_config.json",
)


@dataclass(frozen=True)
class _ParsedObservation:
    image: np.ndarray
    proprio: np.ndarray
    task: str
    instruction: str
    difficulty: str
    step: int
    horizon: int
    episode_key: tuple[str, str, str, int]
    digest: str


@dataclass(frozen=True)
class _RuntimeState:
    episode_key: tuple[str, str, str, int]
    last_step: int
    last_obs_digest: str
    visual_history: tuple[torch.Tensor, ...]
    rgb_history: tuple[torch.Tensor, ...]
    proprio_history: tuple[torch.Tensor, ...]
    phase_history: tuple[torch.Tensor, ...]
    prediction_chunks: tuple[tuple[int, np.ndarray], ...]
    gripper_state: float
    last_action: np.ndarray


def _resolve_device(requested: str) -> torch.device:
    value = str(requested or "cpu").strip().lower()
    if value == "cpu":
        return torch.device("cpu")
    if not value.startswith("cuda"):
        raise ValueError(
            f"unsupported policy device {requested!r}; use 'cpu' or 'cuda[:index]'"
        )
    if not torch.cuda.is_available():
        raise RuntimeError(
            f"CUDA device {requested!r} was requested but "
            "torch.cuda.is_available() is false"
        )
    try:
        device = torch.device(value)
    except (RuntimeError, ValueError) as exc:
        raise ValueError(f"invalid CUDA device {requested!r}") from exc
    index = device.index
    if index is not None and index >= torch.cuda.device_count():
        raise RuntimeError(
            f"CUDA device index {index} is unavailable; "
            f"device_count={torch.cuda.device_count()}"
        )
    return device


def _resolve_dtype(requested: str, device: torch.device) -> torch.dtype:
    # CPU fp16/bf16 coverage varies across LayerNorm, attention, and GRU
    # kernels. Keeping CPU inference in fp32 is the portable validator path.
    if device.type == "cpu":
        return torch.float32
    value = str(requested or "auto").strip().lower().replace("torch.", "")
    if value in {"float32", "fp32"}:
        return torch.float32
    if value in {"float16", "fp16", "half"}:
        return torch.float16
    if value in {"bfloat16", "bf16", "auto"}:
        return torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
    raise ValueError(
        f"unsupported policy dtype {requested!r}; use float32, float16, "
        "bfloat16, or auto"
    )


def _strict_nonnegative_int(value: Any, name: str) -> int:
    if isinstance(value, bool) or not isinstance(value, numbers.Integral):
        raise ValueError(
            f"obs[{name!r}] must be a non-negative integer, got {value!r}"
        )
    result = int(value)
    if result < 0:
        raise ValueError(
            f"obs[{name!r}] must be a non-negative integer, got {value!r}"
        )
    return result


def _canonical_image(value: Any) -> np.ndarray:
    image = np.asarray(value)
    if image.ndim == 2:
        image = np.repeat(image[..., None], 3, axis=-1)
    if image.ndim != 3 or image.shape[0] <= 0 or image.shape[1] <= 0:
        raise ValueError(
            f"obs['image'] must be HxW, HxWx1, HxWx3, or HxWx4; got "
            f"{image.shape}"
        )
    channels = int(image.shape[-1])
    if channels == 1:
        image = np.repeat(image, 3, axis=-1)
    elif channels == 4:
        image = image[..., :3]
    elif channels != 3:
        raise ValueError(
            f"obs['image'] must have 1, 3, or 4 channels; got {image.shape}"
        )
    if not np.issubdtype(image.dtype, np.number):
        raise ValueError(f"obs['image'] must be numeric, got dtype={image.dtype}")
    if not np.isfinite(image).all():
        raise ValueError("obs['image'] contains NaN or Inf")
    if image.dtype != np.uint8:
        image = image.astype(np.float32, copy=False)
        # Accept the common floating-point [0, 1] convention while preserving
        # the same canonical bytes as validator uint8 observations.
        if image.size and float(image.min()) >= 0.0 and float(image.max()) <= 1.0:
            image = image * 255.0
        image = np.rint(np.clip(image, 0.0, 255.0)).astype(np.uint8)
    return np.ascontiguousarray(image)


def _canonical_proprio(value: Any, expected_dim: int) -> np.ndarray:
    try:
        proprio = np.asarray(value, dtype=np.float32)
    except (TypeError, ValueError) as exc:
        raise ValueError(
            "obs['proprio'] must be a one-dimensional numeric array"
        ) from exc
    if proprio.ndim != 1:
        raise ValueError(
            f"obs['proprio'] must have shape ({expected_dim},), got "
            f"{proprio.shape}"
        )
    if not np.isfinite(proprio).all():
        raise ValueError("obs['proprio'] contains NaN or Inf")

    # Robosuite's Wipe task uses the zero-DoF WipingGripper. The official
    # validator consequently emits only the first 21 fields and omits the final
    # gripper_qpos(2) and gripper_qvel(2). Preserve the trained 25D contract by
    # representing those four not-applicable values as zero. No other malformed
    # length is padded, so reordered/corrupt observations remain visible.
    if expected_dim == 25 and proprio.shape == (21,):
        proprio = np.pad(proprio, (0, 4), mode="constant", constant_values=0.0)
    elif proprio.shape != (expected_dim,):
        raise ValueError(
            f"obs['proprio'] must have shape ({expected_dim},), got "
            f"{proprio.shape}"
        )
    return np.ascontiguousarray(proprio)


def _observation_digest(
    image: np.ndarray,
    proprio: np.ndarray,
    episode_key: tuple[str, str, str, int],
    step: int,
) -> str:
    digest = hashlib.sha256()
    for field in (*episode_key[:3], str(episode_key[3]), str(step)):
        encoded = field.encode("utf-8")
        digest.update(len(encoded).to_bytes(8, "little"))
        digest.update(encoded)
    for array in (image, proprio):
        descriptor = (
            f"{array.dtype.str}:{','.join(map(str, array.shape))}".encode("ascii")
        )
        digest.update(len(descriptor).to_bytes(8, "little"))
        digest.update(descriptor)
        digest.update(array.tobytes(order="C"))
    return digest.hexdigest()


def _read_observation(
    obs: Mapping[str, Any], proprio_dim: int
) -> _ParsedObservation:
    if not isinstance(obs, Mapping):
        raise ValueError(
            f"policy observation must be a mapping, got {type(obs).__name__}"
        )
    for required in ("image", "proprio", "step", "horizon"):
        if required not in obs:
            raise ValueError(
                f"policy observation is missing required field {required!r}"
            )
    image = _canonical_image(obs["image"])
    proprio = _canonical_proprio(obs["proprio"], proprio_dim)
    step = _strict_nonnegative_int(obs["step"], "step")
    horizon = _strict_nonnegative_int(obs["horizon"], "horizon")
    if horizon <= 0:
        raise ValueError(f"obs['horizon'] must be positive, got {horizon}")
    task = str(obs.get("task") or "")
    instruction = str(obs.get("instruction") or "")
    difficulty = str(obs.get("difficulty") or "")
    episode_key = (task, instruction, difficulty, horizon)
    return _ParsedObservation(
        image=image,
        proprio=proprio,
        task=task,
        instruction=instruction,
        difficulty=difficulty,
        step=step,
        horizon=horizon,
        episode_key=episode_key,
        digest=_observation_digest(image, proprio, episode_key, step),
    )


def _append_history(
    history: tuple[torch.Tensor, ...], value: torch.Tensor, limit: int
) -> tuple[torch.Tensor, ...]:
    return (*history, value)[-limit:]


def _stack_left_padded(
    history: Sequence[torch.Tensor], expected_length: int
) -> torch.Tensor:
    if not history:
        raise RuntimeError("cannot stack an empty temporal history")
    if len(history) > expected_length:
        raise RuntimeError(
            f"history contains {len(history)} entries; maximum is "
            f"{expected_length}"
        )
    padded = [history[0]] * (expected_length - len(history)) + list(history)
    return torch.stack(padded, dim=0).unsqueeze(0)


def _mean_head_chunks(
    raw: torch.Tensor, ensemble_heads: int, action_chunk: int
) -> np.ndarray:
    expected = (ensemble_heads, 1, action_chunk, ACTION_DIM)
    if not isinstance(raw, torch.Tensor) or tuple(raw.shape) != expected:
        shape = (
            tuple(raw.shape)
            if isinstance(raw, torch.Tensor)
            else type(raw).__name__
        )
        raise RuntimeError(f"policy head must return shape {expected}, got {shape}")
    raw_float = raw.detach().float()
    if not bool(torch.isfinite(raw_float).all().item()):
        raise RuntimeError("policy head returned NaN or Inf")
    chunk = torch.tanh(raw_float).mean(dim=0).squeeze(0)
    result = chunk.cpu().numpy().astype(np.float32, copy=True)
    if result.shape != (action_chunk, ACTION_DIM) or not np.isfinite(result).all():
        raise RuntimeError("head ensemble produced an invalid action chunk")
    result.setflags(write=False)
    return result


def _temporal_ensemble(
    chunks: Mapping[int, np.ndarray],
    step: int,
    action_chunk: int,
    decay: float,
) -> tuple[np.ndarray, tuple[tuple[int, np.ndarray], ...]]:
    retained: list[tuple[int, np.ndarray]] = []
    candidates: list[np.ndarray] = []
    weights: list[float] = []
    for start in sorted(chunks):
        predicted = np.asarray(chunks[start], dtype=np.float32)
        if (
            predicted.shape != (action_chunk, ACTION_DIM)
            or not np.isfinite(predicted).all()
        ):
            raise RuntimeError(f"stored action chunk at step {start} is invalid")
        offset = step - int(start)
        if 0 <= offset < action_chunk:
            retained.append((int(start), chunks[start]))
            candidates.append(predicted[offset])
            weights.append(math.exp(-decay * float(offset)))
    if not candidates:
        raise RuntimeError(
            "temporal ensemble has no action covering the current step"
        )
    candidate_array = np.stack(candidates, axis=0).astype(np.float64)
    weight_array = np.asarray(weights, dtype=np.float64)
    total_weight = float(weight_array.sum())
    if not math.isfinite(total_weight) or total_weight <= 0.0:
        raise RuntimeError("temporal ensemble weights are invalid")
    action = (candidate_array * weight_array[:, None]).sum(axis=0) / total_weight
    result = action.astype(np.float32)
    if result.shape != (ACTION_DIM,) or not np.isfinite(result).all():
        raise RuntimeError("temporal ensemble produced an invalid action")
    return result, tuple(retained)


class SigLIP2ActionChunkPolicy:
    """Stateful SigLIP2 action-chunk policy with transactional semantics."""

    def __init__(
        self,
        model: Any,
        config: Mapping[str, Any],
        device: torch.device,
        dtype: torch.dtype,
        *,
        tokenizer: Any,
        phase_vector_fn: Callable[[int, int], np.ndarray],
    ) -> None:
        self.model = model
        self.config = dict(config)
        self.device = torch.device(device)
        self.dtype = dtype
        self.tokenizer = tokenizer
        self.phase_vector_fn = phase_vector_fn

        self.history = int(self.config.get("history", DEFAULT_HISTORY))
        self.action_chunk = int(
            self.config.get("action_chunk", DEFAULT_ACTION_CHUNK)
        )
        self.ensemble_heads = int(self.config["ensemble_heads"])
        self.proprio_dim = int(self.config["proprio_dim"])
        self.phase_dim = int(self.config.get("phase_dim", 4))
        self.text_dim = int(self.config["text_dim"])
        self.text_max_length = int(
            self.config.get("text_max_length", SIGLIP2_TEXT_MAX_LENGTH)
        )
        self.text_lowercase = self.config.get("text_lowercase")
        self.temporal_decay = float(
            self.config.get("temporal_ensemble_decay", 0.55)
        )
        self.gripper_hysteresis = float(
            self.config.get("gripper_hysteresis", 0.12)
        )
        self.task_to_id = _validate_id_map(
            self.config.get("task_to_id"), "task_to_id"
        )
        self.difficulty_to_id = _validate_id_map(
            self.config.get("difficulty_to_id"), "difficulty_to_id"
        )

        for name, value in (
            ("history", self.history),
            ("action_chunk", self.action_chunk),
            ("ensemble_heads", self.ensemble_heads),
            ("proprio_dim", self.proprio_dim),
            ("phase_dim", self.phase_dim),
            ("text_dim", self.text_dim),
            ("text_max_length", self.text_max_length),
        ):
            if value <= 0:
                raise ValueError(
                    f"vla_config {name} must be positive, got {value}"
                )
        if self.proprio_dim != 25:
            raise ValueError(
                f"FLock proprio_dim must be 25, got {self.proprio_dim}"
            )
        if self.text_max_length != SIGLIP2_TEXT_MAX_LENGTH:
            raise ValueError(
                "SigLIP2 text_max_length must be 64, got "
                f"{self.text_max_length}"
            )
        if self.text_lowercase is not True:
            raise ValueError("SigLIP2 vla_config text_lowercase must be true")
        if not math.isfinite(self.temporal_decay) or self.temporal_decay < 0.0:
            raise ValueError(
                "temporal_ensemble_decay must be finite and non-negative"
            )
        if (
            not math.isfinite(self.gripper_hysteresis)
            or self.gripper_hysteresis < 0.0
        ):
            raise ValueError(
                "gripper_hysteresis must be finite and non-negative"
            )

        self._state: _RuntimeState | None = None
        self._text_cache: dict[str, torch.Tensor] = {}
        self.model.to(device=self.device, dtype=self.dtype)
        self.model.eval()

    def reset(self) -> None:
        """Explicit reset for manual callers; validator reset uses step zero."""
        self._state = None

    @torch.inference_mode()
    def act(self, obs: Mapping[str, Any]) -> np.ndarray:
        parsed = _read_observation(obs, self.proprio_dim)
        committed = self._state

        # A validator retry must be read-only. Digesting canonical image and
        # proprio content distinguishes a retry from a corrected same-step
        # observation or a new episode with different initial state.
        if (
            committed is not None
            and parsed.episode_key == committed.episode_key
            and parsed.step == committed.last_step
            and parsed.digest == committed.last_obs_digest
        ):
            return committed.last_action.copy()

        continuous = (
            committed is not None
            and parsed.step != 0
            and parsed.episode_key == committed.episode_key
            and parsed.step == committed.last_step + 1
        )
        base = committed if continuous else None

        image_tensor = (
            torch.from_numpy(parsed.image.copy()).to(self.device).unsqueeze(0)
        )
        visual, rgb = self.model.encode_images(image_tensor)
        visual = _single_encoded_frame(
            visual, "visual", self.device, self.dtype
        )
        rgb = _single_encoded_frame(rgb, "rgb", self.device, self.dtype)

        proprio_tensor = torch.as_tensor(
            parsed.proprio, device=self.device, dtype=self.dtype
        )
        phase_array = _fixed_feature_array(
            self.phase_vector_fn(parsed.step, parsed.horizon),
            self.phase_dim,
            "phase",
        )
        phase_tensor = torch.as_tensor(
            phase_array, device=self.device, dtype=self.dtype
        )

        visual_history = _append_history(
            base.visual_history if base else (), visual, self.history
        )
        rgb_history = _append_history(
            base.rgb_history if base else (), rgb, self.history
        )
        proprio_history = _append_history(
            base.proprio_history if base else (), proprio_tensor, self.history
        )
        phase_history = _append_history(
            base.phase_history if base else (), phase_tensor, self.history
        )

        # SigLIP2 was trained with lowercased SentencePiece text and fixed
        # length 64. The normalized string is also the cache key so case-only
        # instruction variants share exactly the same embedding.
        text_key = parsed.instruction.strip().lower()
        pending_text_cache: torch.Tensor | None = None
        cached_text = self._text_cache.get(text_key)
        if cached_text is None:
            encoded = self.tokenizer(
                text_key,
                padding="max_length",
                truncation=True,
                max_length=self.text_max_length,
                return_attention_mask=True,
                return_tensors="pt",
            )
            if "input_ids" not in encoded or "attention_mask" not in encoded:
                raise RuntimeError(
                    "SigLIP2 tokenizer did not return input_ids and attention_mask"
                )
            text_features = self.model.encode_text(
                encoded["input_ids"], encoded["attention_mask"]
            )
            text_features = _single_text_feature(
                text_features, self.text_dim, self.device, self.dtype
            )
            pending_text_cache = text_features.detach().clone()
        else:
            text_features = cached_text

        task_id = self.task_to_id.get(parsed.task, len(self.task_to_id))
        difficulty_id = self.difficulty_to_id.get(
            parsed.difficulty, len(self.difficulty_to_id)
        )
        raw = self.model.forward_cached(
            _stack_left_padded(visual_history, self.history),
            _stack_left_padded(rgb_history, self.history),
            _stack_left_padded(proprio_history, self.history),
            _stack_left_padded(phase_history, self.history),
            text_features,
            torch.tensor([task_id], device=self.device, dtype=torch.long),
            torch.tensor(
                [difficulty_id], device=self.device, dtype=torch.long
            ),
        )
        current_chunk = _mean_head_chunks(
            raw,
            ensemble_heads=self.ensemble_heads,
            action_chunk=self.action_chunk,
        )

        chunks = dict(base.prediction_chunks) if base else {}
        chunks[parsed.step] = current_chunk
        action, retained_chunks = _temporal_ensemble(
            chunks,
            step=parsed.step,
            action_chunk=self.action_chunk,
            decay=self.temporal_decay,
        )

        gripper_state = base.gripper_state if base else -1.0
        gripper_score = float(action[6])
        if gripper_score > self.gripper_hysteresis:
            gripper_state = 1.0
        elif gripper_score < -self.gripper_hysteresis:
            gripper_state = -1.0
        action = np.clip(action, -1.0, 1.0).astype(np.float32, copy=True)
        action[6] = np.float32(gripper_state)
        if action.shape != (ACTION_DIM,) or not np.isfinite(action).all():
            raise RuntimeError("policy produced an invalid final action")

        cached_action = action.copy()
        cached_action.setflags(write=False)
        # This is the only state mutation in act(). Any exception above leaves
        # the last successfully committed state byte-for-byte intact.
        self._state = _RuntimeState(
            episode_key=parsed.episode_key,
            last_step=parsed.step,
            last_obs_digest=parsed.digest,
            visual_history=visual_history,
            rgb_history=rgb_history,
            proprio_history=proprio_history,
            phase_history=phase_history,
            prediction_chunks=retained_chunks,
            gripper_state=gripper_state,
            last_action=cached_action,
        )
        if pending_text_cache is not None:
            self._text_cache[text_key] = pending_text_cache
        return action.copy()


def _fixed_feature_array(
    value: Any, expected_dim: int, name: str
) -> np.ndarray:
    try:
        array = np.asarray(value, dtype=np.float32)
    except (TypeError, ValueError) as exc:
        raise RuntimeError(
            f"{name} feature function returned non-numeric data"
        ) from exc
    if array.shape != (expected_dim,) or not np.isfinite(array).all():
        raise RuntimeError(
            f"{name} feature must have shape ({expected_dim},) with finite "
            f"values; got {array.shape}"
        )
    return np.ascontiguousarray(array)


def _single_encoded_frame(
    value: Any, name: str, device: torch.device, dtype: torch.dtype
) -> torch.Tensor:
    if (
        not isinstance(value, torch.Tensor)
        or value.ndim != 3
        or value.shape[0] != 1
    ):
        shape = (
            tuple(value.shape)
            if isinstance(value, torch.Tensor)
            else type(value).__name__
        )
        raise RuntimeError(
            f"image encoder {name} output must be [1,tokens,dim], got {shape}"
        )
    result = value.squeeze(0).detach().to(device=device, dtype=dtype)
    if result.numel() == 0 or not bool(
        torch.isfinite(result.float()).all().item()
    ):
        raise RuntimeError(
            f"image encoder {name} output contains no finite tokens"
        )
    return result


def _single_text_feature(
    value: Any,
    expected_dim: int,
    device: torch.device,
    dtype: torch.dtype,
) -> torch.Tensor:
    if (
        not isinstance(value, torch.Tensor)
        or value.ndim != 2
        or tuple(value.shape) != (1, expected_dim)
    ):
        shape = (
            tuple(value.shape)
            if isinstance(value, torch.Tensor)
            else type(value).__name__
        )
        raise RuntimeError(
            f"text encoder output must be [1,{expected_dim}], got {shape}"
        )
    result = value.detach().to(device=device, dtype=dtype)
    if not bool(torch.isfinite(result.float()).all().item()):
        raise RuntimeError("text encoder output contains NaN or Inf")
    return result


def _validate_id_map(value: Any, name: str) -> dict[str, int]:
    if not isinstance(value, Mapping):
        raise ValueError(f"vla_config {name} must be an object")
    result: dict[str, int] = {}
    for key, item in value.items():
        if (
            not isinstance(key, str)
            or isinstance(item, bool)
            or not isinstance(item, numbers.Integral)
        ):
            raise ValueError(
                f"vla_config {name} must map strings to integer IDs"
            )
        result[key] = int(item)
    if set(result.values()) != set(range(len(result))):
        raise ValueError(
            f"vla_config {name} IDs must be contiguous from zero"
        )
    return result


def _build_local_siglip2_tokenizer(
    tokenizer_cls: Any,
    model_root: Path,
    siglip_config: Mapping[str, Any],
) -> Any:
    """Load and validate the bundled SigLIP2 tokenizer without Hub access."""
    text_config = siglip_config.get("text_config")
    if not isinstance(text_config, Mapping):
        raise ValueError("siglip_config.text_config must be an object")
    try:
        expected_vocab_size = int(text_config["vocab_size"])
    except (KeyError, TypeError, ValueError) as exc:
        raise ValueError(
            "siglip_config.text_config has an invalid vocab_size"
        ) from exc
    try:
        tokenizer = tokenizer_cls.from_pretrained(
            str(model_root),
            local_files_only=True,
            use_fast=True,
        )
        observed_vocab_size = int(tokenizer.vocab_size)
        observed_full_vocab_size = len(tokenizer.get_vocab())
    except (AttributeError, OSError, TypeError, ValueError) as exc:
        raise ValueError(
            "Could not construct the bundled SigLIP2 tokenizer"
        ) from exc
    if (
        observed_vocab_size != expected_vocab_size
        or observed_full_vocab_size != expected_vocab_size
    ):
        raise ValueError(
            "Bundled SigLIP2 tokenizer vocabulary does not match "
            "siglip_config.text_config: "
            f"expected={expected_vocab_size}, "
            f"vocab_size={observed_vocab_size}, "
            f"get_vocab_size={observed_full_vocab_size}"
        )

    # Exercise the exact runtime contract now rather than discovering a broken
    # tokenizer only after the simulator has started an episode.
    try:
        probe = tokenizer(
            "mixed case siglip2 probe",
            padding="max_length",
            truncation=True,
            max_length=SIGLIP2_TEXT_MAX_LENGTH,
            return_attention_mask=True,
            return_tensors="pt",
        )
        input_ids = probe["input_ids"]
        attention_mask = probe["attention_mask"]
    except (KeyError, OSError, TypeError, ValueError) as exc:
        raise ValueError(
            "Bundled SigLIP2 tokenizer failed its fixed-length probe"
        ) from exc
    expected_shape = (1, SIGLIP2_TEXT_MAX_LENGTH)
    if (
        not isinstance(input_ids, torch.Tensor)
        or not isinstance(attention_mask, torch.Tensor)
        or tuple(input_ids.shape) != expected_shape
        or tuple(attention_mask.shape) != expected_shape
    ):
        raise ValueError(
            "Bundled SigLIP2 tokenizer must produce input_ids and "
            f"attention_mask with shape {expected_shape}"
        )
    return tokenizer


def load_policy(
    model_dir: str, device: str, dtype: str
) -> SigLIP2ActionChunkPolicy:
    """Load a complete local SigLIP2 checkpoint with no network fallback."""
    model_root = Path(model_dir).expanduser().resolve()
    config_path = model_root / "vla_config.json"
    weights_path = model_root / "model.safetensors"
    runtime_path = model_root / "siglip2_action_chunk_model.py"
    required_paths = [config_path, weights_path, runtime_path]
    required_paths.extend(
        model_root / asset for asset in REQUIRED_TOKENIZER_ASSETS
    )
    for path in required_paths:
        if not path.is_file():
            raise FileNotFoundError(
                f"self-contained SigLIP2 submission is missing {path.name}"
            )

    config = json.loads(config_path.read_text(encoding="utf-8"))
    if not isinstance(config, dict) or not isinstance(
        config.get("siglip_config"), dict
    ):
        raise ValueError(
            "vla_config.json must contain an inline siglip_config object"
        )
    if str(model_root) not in sys.path:
        sys.path.insert(0, str(model_root))

    # Imports are intentionally local. Production has no model-Hub loading
    # path, while unit tests can replace the runtime model and tokenizer.
    from safetensors.torch import load_file
    from transformers import AutoTokenizer
    from siglip2_action_chunk_model import (
        SigLIP2ActionChunkModel,
        phase_vector,
    )

    torch_device = _resolve_device(device)
    torch_dtype = _resolve_dtype(dtype, torch_device)
    model = SigLIP2ActionChunkModel(config)
    state = load_file(str(weights_path), device="cpu")
    model.load_state_dict(state, strict=True)
    del state
    tokenizer = _build_local_siglip2_tokenizer(
        AutoTokenizer,
        model_root,
        config["siglip_config"],
    )
    return SigLIP2ActionChunkPolicy(
        model=model,
        config=config,
        device=torch_device,
        dtype=torch_dtype,
        tokenizer=tokenizer,
        phase_vector_fn=phase_vector,
    )


__all__ = ["SigLIP2ActionChunkPolicy", "load_policy"]