"""Load ActionCodec2Sparse release packages into a router."""

from __future__ import annotations

import json
from pathlib import Path

from ..backends.actioncodec2.defaults import builtin_action_spaces
from ..backends.actioncodec2.preprocessing import PreprocessingSpec
from ..backends.actioncodec2.presets import format_action_spaces
from ..backends.actioncodec2.release import (
    check_gripper_consistency,
    check_release_gates,
    selected_variants,
)
from ..backends.actioncodec2.sparse import ActionCodec2SparseProcessor, ActionCodec2SparseRouter

TOKENIZER = "actioncodec2_sparse"


def release_layouts(profiles):
    """Built-in layouts whose components are all available in ``profiles``."""
    return {
        name: preset
        for name, preset in builtin_action_spaces().items()
        if {component.space for component in preset.components} <= set(profiles)
    }


def load_release_router(directory, *, gripper_mode, primitive_order=1, verify_gates, tokenizer=TOKENIZER):
    """Load one gripper-mode/order vocabulary per profile with its cleaning policy."""
    manifest, mode, variants = selected_variants(
        directory, gripper_mode=gripper_mode, order=primitive_order, tokenizer=tokenizer
    )
    from ..backends.actioncodec2 import ActionCodec2Processor, ActionCodec2Router

    processor_class, router_class = ((ActionCodec2SparseProcessor, ActionCodec2SparseRouter)
                                     if tokenizer == TOKENIZER else (ActionCodec2Processor, ActionCodec2Router))
    profiles, preprocessing = {}, {}
    for name, variant in variants.items():
        entry = (Path(directory) / variant["path"]).resolve()
        if not entry.is_relative_to(Path(directory).resolve()):
            raise ValueError("release profile path escapes its artifact directory")
        processor = processor_class.from_pretrained(entry)
        spec = PreprocessingSpec.from_dict(
            json.loads((entry / "preprocessing.json").read_text(encoding="utf-8"))
        )
        check_gripper_consistency(entry, processor, spec, mode)
        if verify_gates:
            if variant.get("gates_passed") is False:
                raise ValueError(f"{entry}: recorded release gates failed")
            check_release_gates(entry)
        profiles[name], preprocessing[name] = processor, spec
    router = router_class.from_processors(
        profiles,
        token_budgets={name: manifest["codebook_size"] for name in profiles},
        presets=list(release_layouts(profiles).values()),
    )
    metadata = {
        "package": manifest.get("package", Path(directory).name),
        "tokenizer": tokenizer,
        "precision": manifest["precision"],
        "codebook_size": manifest["codebook_size"],
        "gripper_mode": mode,
        "primitive_order": primitive_order,
        "variants": list(variants.values()),
    }
    return router, preprocessing, metadata


def describe_release(directory, *, action_dim, gripper_mode, primitive_order=1, tokenizer=TOKENIZER):
    """Report a release's layouts from its manifest alone."""
    _, _, variants = selected_variants(
        directory, gripper_mode=gripper_mode, order=primitive_order, tokenizer=tokenizer
    )
    presets = release_layouts(variants)
    return format_action_spaces(
        presets, action_dim=action_dim, codec_name="ActionCodec2Sparse" if tokenizer == TOKENIZER else "ActionCodec2", chunk_horizon_decode=tokenizer == TOKENIZER
    )