#!/usr/bin/env python3
"""Validate pinned X-VLA weights and LeRobot training YAML files without loading the model."""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any

import draccus
import numpy as np
import yaml
from safetensors import safe_open

from lerobot.configs import parser
from lerobot.configs.train import TrainPipelineConfig
import lerobot.configs.policies as policy_configs
from lerobot.policies.xvla.configuration_xvla import XVLAConfig  # noqa: F401


MODEL_SHA256 = "f05bc0fab1c9523d7f5d6b41a651313641ca2227a88822249845da8e20e036c9"


def parse_training_config(path: Path) -> TrainPipelineConfig:
    # Config parsing normally replaces unavailable CUDA with CPU. This check is
    # intentionally hardware-independent, so preserve the requested CUDA device.
    original_device_check = policy_configs.is_torch_device_available
    policy_configs.is_torch_device_available = lambda _device: True
    try:
        clean_path = parser.extract_path_fields_from_config(
            str(path), TrainPipelineConfig.__get_path_fields__()
        )
        cfg = draccus.parse(TrainPipelineConfig, config_path=clean_path, args=[])
        cfg.validate()
        return cfg
    finally:
        policy_configs.is_torch_device_available = original_device_check


def validate_config(
    path: Path,
    expected_steps: int,
    expected_batch_size: int,
    expected_save_checkpoint: bool = True,
) -> dict[str, Any]:
    raw_config = yaml.safe_load(path.read_text(encoding="utf-8"))
    raw_policy = raw_config.get("policy") or {}
    for feature_field in ("input_features", "output_features"):
        if feature_field not in raw_policy or raw_policy[feature_field] is not None:
            raise ValueError(
                f"{path}: policy.{feature_field} must be an unquoted YAML null so "
                "the production CLI infers dataset features"
            )

    cfg = parse_training_config(path)
    policy = cfg.policy
    if policy is None or policy.type != "xvla":
        raise ValueError(f"{path}: did not resolve to an XVLA policy")
    expected = {
        "device": "cuda",
        "dtype": "bfloat16",
        "action_mode": "auto",
        "max_action_dim": 20,
        "max_state_dim": 20,
        "num_image_views": 3,
        "freeze_vision_encoder": False,
        "freeze_language_encoder": False,
        "train_policy_transformer": True,
        "train_soft_prompts": True,
    }
    for key, value in expected.items():
        if getattr(policy, key) != value:
            raise ValueError(f"{path}: policy.{key} did not resolve to {value!r}")
    if policy.input_features is not None or policy.output_features is not None:
        raise ValueError(f"{path}: policy feature dictionaries must be inferred from the dataset")
    if cfg.steps != expected_steps or cfg.batch_size != expected_batch_size:
        raise ValueError(f"{path}: unexpected steps or batch size")
    if cfg.save_checkpoint is not expected_save_checkpoint:
        raise ValueError(f"{path}: unexpected checkpoint-saving setting")
    if cfg.tolerance_s != 0.001 or cfg.dataset.eval_split != 0.1:
        raise ValueError(f"{path}: unexpected timestamp tolerance or eval split")
    if cfg.dataset.use_imagenet_stats:
        raise ValueError(
            f"{path}: dataset.use_imagenet_stats must be false because X-VLA normalizes "
            "images in its policy processor and merged camera stats are intentionally absent"
        )
    return {
        "path": str(path),
        "steps": cfg.steps,
        "batch_size_per_process": cfg.batch_size,
        "save_checkpoint": cfg.save_checkpoint,
        "dataset_root": str(cfg.dataset.root),
        "policy_path": str(policy.pretrained_path),
        "dtype": policy.dtype,
        "action_mode": policy.action_mode,
        "num_image_views": policy.num_image_views,
        "use_imagenet_stats": cfg.dataset.use_imagenet_stats,
        "eval_split": cfg.dataset.eval_split,
        "tolerance_s": cfg.tolerance_s,
    }


def main() -> None:
    project_root = Path(__file__).resolve().parents[1]
    model_root = project_root / "models" / "xvla-base"
    model_manifest = json.loads(
        (model_root / "download_manifest.json").read_text(encoding="utf-8")
    )
    if model_manifest["model_sha256"] != MODEL_SHA256:
        raise ValueError("Pinned X-VLA model manifest SHA256 mismatch")

    tensor_count = 0
    parameter_count = 0
    dtype_counts: dict[str, int] = {}
    with safe_open(model_root / "model.safetensors", framework="pt", device="cpu") as stream:
        for key in stream.keys():
            tensor = stream.get_slice(key)
            count = int(np.prod(tensor.get_shape()))
            dtype = str(tensor.get_dtype())
            tensor_count += 1
            parameter_count += count
            dtype_counts[dtype] = dtype_counts.get(dtype, 0) + count

    merged_validation = json.loads(
        (project_root / "data" / "merged" / "validation_report.json").read_text(encoding="utf-8")
    )
    if merged_validation["status"] != "passed" or merged_validation["frames"] != 120_469:
        raise ValueError("Merged training dataset validation report is not usable")

    configs = [
        validate_config(project_root / "configs" / "xvla_smoke.yaml", 20, 1),
        validate_config(
            project_root / "configs" / "xvla_pilot.yaml",
            100,
            16,
            expected_save_checkpoint=False,
        ),
        validate_config(project_root / "configs" / "xvla_full.yaml", 20_000, 4),
        validate_config(project_root / "configs" / "xvla_full_1gpu.yaml", 20_000, 16),
    ]
    report = {
        "status": "passed",
        "lerobot_version": "0.6.0",
        "model": {
            "root": str(model_root),
            "revision": model_manifest["revision"],
            "sha256": model_manifest["model_sha256"],
            "tensor_count": tensor_count,
            "parameter_count": parameter_count,
            "dtype_parameter_counts": dtype_counts,
        },
        "dataset": {
            "root": merged_validation["root"],
            "episodes": merged_validation["episodes"],
            "frames": merged_validation["frames"],
            "tasks": len(merged_validation["tasks"]),
        },
        "configs": configs,
    }
    report_path = project_root / "configs" / "train_setup_validation.json"
    report_path.write_text(
        json.dumps(report, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    print(
        f"[passed] X-VLA setup: {parameter_count:,} parameters, "
        f"{merged_validation['frames']:,} data frames, {len(configs)} configs",
        flush=True,
    )
    print(f"Validation report: {report_path}", flush=True)


if __name__ == "__main__":
    main()