"""Load a named ActionCodec2 variant or a standalone saved processor."""

from __future__ import annotations

import hashlib
import importlib
import importlib.util
import json
from pathlib import Path
import re
import sys


PROFILE_FILES = ("config.yaml", "primitives.npz", "bpe.json", "config.json", "processor_config.json")
SHARED_PATTERNS = (
    "config.json", "processor_config.json", "processing_actioncodec2.py", "variants.json",
    "runtime/**", "actioncodec2_kernels/**", "setup.py", "pyproject.toml", "LICENSE",
    "README.md", "requirements.txt",
)


VARIANT_ALIASES = {
    **{f"paper-{order}": f"paper-o{index}-5mm-4k-2bin-1p2ms"
       for index, order in enumerate(("1st", "2nd"), 1)},
    **{f"extended-{order}-{precision}-{size}k":
       f"extended-o{index}-{precision}-{size}k-2bin-1p2ms"
       for index, order in enumerate(("1st", "2nd"), 1)
       for precision in ("1mm", "5mm", "10mm") for size in (1, 2, 4, 8, 16)
       if precision != "1mm" or size != 1},
    **{f"sparse-{dynamics}-{mode}": f"sparse-o1-5mm-4k-{mode}-{speed}"
       for dynamics, speed in (("standard", "1p2ms"), ("high-dynamic", "2ms"))
       for mode in ("2bin", "21bin")},
}


def resolve_variant_name(variant, variants):
    """Resolve either published spelling against an old or new catalog."""
    if variant in variants:
        return variant
    canonical = VARIANT_ALIASES.get(variant, variant)
    if canonical in variants:
        return canonical
    for legacy, name in VARIANT_ALIASES.items():
        if name == canonical and legacy in variants:
            return legacy
    raise ValueError(f"select variant= from: {', '.join(sorted(variants))}")


def asset_inventory(directory: Path) -> dict[str, str]:
    """Hash the complete configuration and fitted assets for one variant."""
    sparse = (directory / "sparse_router_config.yaml").is_file()
    names = ["sparse_router_config.yaml" if sparse else "router_config.yaml", "config.json", "processor_config.json"]
    profiles = [name for name in ("eef", "joint") if (directory / "profiles" / name).is_dir()]
    if not profiles:
        raise ValueError("variant has no fitted profiles")
    names.extend(f"profiles/{profile}/{name}" for profile in profiles
                 for name in (tuple("sparse_bpe.json" if n == "bpe.json" else n for n in PROFILE_FILES) if sparse else PROFILE_FILES))
    names.extend(f"profiles/{profile}/preprocessing.json" for profile in profiles
                 if (directory / "profiles" / profile / "preprocessing.json").is_file())
    return {name: hashlib.sha256((directory / name).read_bytes()).hexdigest()
            for name in sorted(names)}


def inventory_fingerprint(files: dict[str, str]) -> str:
    """Return the stable identity of a fitted asset inventory."""
    return hashlib.sha256(json.dumps(files, sort_keys=True, separators=(",", ":")).encode()).hexdigest()


def _check_inventory(directory: Path, record: dict) -> None:
    files = asset_inventory(directory)
    if files != record.get("files") or inventory_fingerprint(files) != record.get("fingerprint"):
        raise ValueError("ActionCodec2 asset fingerprint mismatch; restore the complete selected variant")


class _VariantBinding:
    @property
    def variant(self):
        return self._variant_origin.get("name")

    @property
    def variant_fingerprint(self):
        return self._variant_origin.get("fingerprint")

    def for_action_space(self, name):
        codec = super().for_action_space(name)
        codec._variant_origin = dict(self._variant_origin)
        return codec

    def fit(self, *args, **kwargs):
        result = super().fit(*args, **kwargs)
        self._variant_origin = {}
        return result

    def save_pretrained(self, save_directory, push_to_hub=False, **kwargs):
        directory = Path(save_directory)
        include_training_report = kwargs.pop("include_training_report", False)
        if (directory / "variants.json").is_file():
            raise ValueError("save a selected variant to a separate directory, outside the catalog root")
        if not push_to_hub and kwargs:
            raise TypeError(f"unused save_pretrained keyword(s): {', '.join(sorted(kwargs))}")
        files = super().save_pretrained(directory, include_training_report=include_training_report)
        inventory = asset_inventory(directory)
        manifest = directory / "artifact_manifest.json"
        manifest.write_text(json.dumps({
            "format_version": 1,
            "origin": self._variant_origin,
            "files": inventory,
            "fingerprint": inventory_fingerprint(inventory),
        }, indent=2) + "\n", encoding="utf-8")
        files.append(str(manifest))
        if push_to_hub:
            package = type(self).__module__.split(".frontend.", 1)[0]
            importlib.import_module(f"{package}.backends.actioncodec2.hf").upload_pretrained_directory(directory, **kwargs)
        return files


class ActionCodec2:
    """Hugging Face entry point for catalog and standalone artifacts."""

    @classmethod
    def register_for_auto_class(cls, auto_class="AutoProcessor"):
        return cls

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *, variant=None,
                        variant_fingerprint=None, action_space=None, **kwargs):
        source = Path(pretrained_model_name_or_path).expanduser()
        kwargs.pop("_from_auto", None)
        kwargs.pop("trust_remote_code", None)
        subfolder = kwargs.pop("subfolder", "")
        if subfolder:
            raise ValueError("select an ActionCodec2 catalog variant with variant=, not subfolder=")
        allowed = {"cache_dir", "force_download", "local_files_only", "token", "revision", "repo_type"}
        unknown = sorted(set(kwargs) - allowed)
        if unknown:
            raise TypeError(f"unsupported from_pretrained keyword(s): {', '.join(unknown)}")
        if variant is not None and (not isinstance(variant, str) or
                                    re.fullmatch(r"[a-z0-9]+(?:-[a-z0-9]+)*", variant) is None):
            raise ValueError("variant must be a published variant name")
        if variant_fingerprint is not None and (
            not isinstance(variant_fingerprint, str) or
            re.fullmatch(r"[0-9a-f]{64}", variant_fingerprint) is None
        ):
            raise ValueError("variant_fingerprint must be a SHA-256 fingerprint")
        remote = not source.is_dir()
        if remote and str(pretrained_model_name_or_path).startswith(("/", "./", "../", "~")):
            raise FileNotFoundError(f"ActionCodec2 artifact does not exist: {source}")
        revision = None
        if remote:
            from huggingface_hub import hf_hub_download, snapshot_download

            code_revision = Path(__file__).parent.name
            if re.fullmatch(r"[0-9a-f]{40}", code_revision):
                requested = kwargs.get("revision")
                if requested and re.fullmatch(r"[0-9a-f]{40}", requested) and requested != code_revision:
                    raise ValueError("ActionCodec2 code and assets must use the same revision")
                kwargs["revision"] = code_revision
            metadata_file = Path(hf_hub_download(
                str(pretrained_model_name_or_path), "processor_config.json", **kwargs
            ))
            root = metadata_file.parent
            revision = root.name
            kwargs["revision"] = revision
        else:
            root = source
            metadata_file = root / "processor_config.json"
        metadata = json.loads(metadata_file.read_text(encoding="utf-8"))
        if not isinstance(metadata, dict):
            raise ValueError("invalid ActionCodec2 processor configuration")
        record = None
        origin = {}
        if metadata.get("variant_catalog") == "variants.json":
            catalog_path = (Path(hf_hub_download(
                str(pretrained_model_name_or_path), "variants.json", **kwargs
            )) if remote else root / "variants.json")
            catalog = json.loads(catalog_path.read_text(encoding="utf-8"))
            if (not isinstance(catalog, dict) or catalog.get("format_version") != 1 or
                    not isinstance(catalog.get("variants"), dict)):
                raise ValueError("unsupported ActionCodec2 variant catalog")
            variants = catalog["variants"]
            variant = resolve_variant_name(variant, variants)
            record = variants[variant]
            if not isinstance(record, dict):
                raise ValueError("invalid ActionCodec2 variant record")
            directory = root / "variants" / variant
            if remote:
                root = Path(snapshot_download(
                    str(pretrained_model_name_or_path),
                    allow_patterns=[*SHARED_PATTERNS, f"variants/{variant}/**"], **kwargs,
                ))
                directory = root / "variants" / variant
            _check_inventory(directory, record)
            origin = {"name": variant, "fingerprint": record["fingerprint"]}
            if remote:
                origin.update(repo_id=str(pretrained_model_name_or_path), revision=revision)
        else:
            if remote:
                root = Path(snapshot_download(
                    str(pretrained_model_name_or_path),
                    allow_patterns=[*SHARED_PATTERNS, "router_config.yaml", "profiles/**",
                                    "sparse_router_config.yaml", "artifact_manifest.json"], **kwargs,
                ))
            directory = root
            manifest = root / "artifact_manifest.json"
            if manifest.is_file():
                record = json.loads(manifest.read_text(encoding="utf-8"))
                if (not isinstance(record, dict) or record.get("format_version") != 1 or
                        not isinstance(record.get("origin"), dict)):
                    raise ValueError("unsupported ActionCodec2 saved artifact manifest")
                _check_inventory(directory, record)
                origin = record["origin"]
            if variant is not None and VARIANT_ALIASES.get(variant, variant) != VARIANT_ALIASES.get(origin.get("name"), origin.get("name")):
                raise ValueError("saved processor contains a different variant; load the catalog to switch")
        if variant_fingerprint is not None and variant_fingerprint != origin.get("fingerprint"):
            raise ValueError("variant fingerprint differs from the tokenizer bound to this policy")

        family = "actioncodec2_sparse" if (directory / "sparse_router_config.yaml").is_file() else "actioncodec2"
        if __package__ == "actioncodec2.core":
            package_name = __package__
        else:
            runtime = root / "runtime"
            package_file = runtime / "__init__.py"
            if not package_file.is_file():
                raise FileNotFoundError(f"ActionCodec2 runtime is missing from {root}")
            identity = hashlib.sha256(str(runtime.resolve()).encode()).hexdigest()[:16]
            package_name = f"_actioncodec2_artifact_{identity}"
            if package_name not in sys.modules:
                spec = importlib.util.spec_from_file_location(
                    package_name, package_file, submodule_search_locations=[str(runtime)]
                )
                package = importlib.util.module_from_spec(spec)
                sys.modules[package_name] = package
                spec.loader.exec_module(package)
        processor = importlib.import_module(f"{package_name}.frontend.{family}.processor")
        base = processor.ActionCodec2Sparse if family == "actioncodec2_sparse" else processor.ActionCodec2
        if not hasattr(processor, "_catalog_processor_class"):
            processor._catalog_processor_class = type(
                base.__name__, (_VariantBinding, base), {"__module__": processor.__name__}
            )
        codec = processor._catalog_processor_class.from_pretrained(directory, action_space=action_space)
        codec._variant_origin = origin
        codec._source_directory = root
        return codec