'Hugging Face AutoProcessor integration.'

from __future__ import annotations

import json
from pathlib import Path
from typing import TYPE_CHECKING, Any

from transformers import AutoConfig, AutoProcessor, PretrainedConfig

from .configuration import CONFIG_NAME, MODEL_TYPE

if TYPE_CHECKING:
    from .processor import ActionCodec2Processor
    from .routing import ActionCodec2Router

HF_CONFIG_NAME = "config.json"
PROCESSOR_CONFIG_NAME = "processor_config.json"
PROCESSOR_CLASS_NAME = "ActionCodec2Processor"
ROUTER_MODEL_TYPE = "actioncodec2_router"
ROUTER_PROCESSOR_CLASS_NAME = "ActionCodec2Router"
ROUTER_CONFIG_NAME = "router_config.yaml"
GENERATED_CARD_MARKER = (
    "<!-- generated by actioncodec2; delete this line to keep your own edits -->"
)


ROUTER_MANIFEST_PATTERNS = [
    ROUTER_CONFIG_NAME,
    HF_CONFIG_NAME,
    PROCESSOR_CONFIG_NAME,
    f"profiles/*/{CONFIG_NAME}",
    "manifest.json",
    f"joint/**/{CONFIG_NAME}",
    f"eef/**/{CONFIG_NAME}",
]
ROUTER_ALLOW_PATTERNS = [
    *ROUTER_MANIFEST_PATTERNS,
    "profiles/**",
    "joint/**",
    "eef/**",
    "fit_report.json",
    "README.md",
]


class ActionCodec2FrontendHFConfig(PretrainedConfig):
    """AutoConfig discovery key for the action-space-bound public processor."""

    model_type = "actioncodec2"

    def __init__(self, **kwargs: Any) -> None:
        kwargs.setdefault("processor_class", "ActionCodec2")
        super().__init__(**kwargs)


class ActionCodec2HFConfig(PretrainedConfig):
    """Minimal AutoConfig key that dispatches to :class:`ActionCodec2Processor`."""

    model_type = MODEL_TYPE

    def __init__(self, **kwargs: Any) -> None:
        kwargs.setdefault("processor_class", PROCESSOR_CLASS_NAME)
        super().__init__(**kwargs)


class ActionCodec2RouterHFConfig(PretrainedConfig):
    """Minimal AutoConfig key that dispatches to :class:`ActionCodec2Router`."""

    model_type = ROUTER_MODEL_TYPE

    def __init__(self, **kwargs: Any) -> None:
        kwargs.setdefault("processor_class", ROUTER_PROCESSOR_CLASS_NAME)
        super().__init__(**kwargs)


def save_huggingface_metadata(
    directory: str | Path,
    *,
    model_type: str = MODEL_TYPE,
    processor_class: str = PROCESSOR_CLASS_NAME,
    **processor_config: Any,
) -> tuple[Path, Path]:
    """Write standard local/Hub discovery files for ``AutoProcessor``."""
    root = Path(directory)
    root.mkdir(parents=True, exist_ok=True)
    config_path = root / HF_CONFIG_NAME
    processor_path = root / PROCESSOR_CONFIG_NAME
    config_path.write_text(
        json.dumps(
            {
                "model_type": model_type,
                "processor_class": processor_class,
            },
            indent=2,
        )
        + "\n",
        encoding="utf-8",
    )
    processor_path.write_text(
        json.dumps({"processor_class": processor_class, **processor_config}, indent=2) + "\n",
        encoding="utf-8",
    )
    return config_path, processor_path


def resolve_pretrained_directory(
    pretrained_model_name_or_path: str | Path,
    *,
    allow_patterns: list[str] | None = None,
    **kwargs: Any,
) -> Path:
    """Resolve a local directory or download one immutable Hub snapshot."""
    source = Path(pretrained_model_name_or_path)
    kwargs.pop("_from_auto", None)
    kwargs.pop("trust_remote_code", None)
    subfolder = str(kwargs.pop("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 source.is_dir():
        return source / subfolder

    try:
        from huggingface_hub import snapshot_download
    except ImportError as error:
        raise ImportError(
            "loading a Hub artifact requires huggingface_hub (installed with Transformers)"
        ) from error
    patterns = allow_patterns or [
        "config.yaml",
        "primitives.npz",
        "bpe.json",
        HF_CONFIG_NAME,
        PROCESSOR_CONFIG_NAME,
    ]
    if subfolder:
        prefix = subfolder.rstrip("/")
        patterns = [f"{prefix}/{pattern}" for pattern in patterns]
    snapshot = snapshot_download(
        repo_id=str(pretrained_model_name_or_path),
        allow_patterns=patterns,
        **kwargs,
    )
    return Path(snapshot) / subfolder


def upload_pretrained_directory(directory: Path, **kwargs: Any) -> None:
    """Upload a complete, already-saved processor artifact to the Hub."""
    from huggingface_hub import HfApi

    repo_id = kwargs.pop("repo_id", directory.name)
    commit_message = kwargs.pop("commit_message", "Upload ActionCodec2 processor")
    api = HfApi(token=kwargs.pop("token", None))
    api.create_repo(repo_id=repo_id, exist_ok=True, **kwargs)
    api.upload_folder(
        repo_id=repo_id,
        folder_path=str(directory),
        commit_message=commit_message,
    )


def register_actioncodec2_auto_processors(
    processor_class: type[ActionCodec2Processor],
    router_class: type[ActionCodec2Router],
) -> bool:
    'Register both installed ActionCodec2 artifact types with Transformers.'
    AutoConfig.register(MODEL_TYPE, ActionCodec2HFConfig, exist_ok=True)
    AutoProcessor.register(ActionCodec2HFConfig, processor_class, exist_ok=True)
    AutoConfig.register(ROUTER_MODEL_TYPE, ActionCodec2RouterHFConfig, exist_ok=True)
    AutoProcessor.register(ActionCodec2RouterHFConfig, router_class, exist_ok=True)
    return True


def save_model_card(directory: str | Path, description: str) -> Path:
    'Write the generated Hub card, refreshing it whenever the artifact is re-saved.'
    path = Path(directory) / "README.md"
    if not path.exists() or GENERATED_CARD_MARKER in path.read_text(encoding="utf-8"):
        path.parent.mkdir(parents=True, exist_ok=True)
        guide = (
            "---\ntags:\n- robotics\n- action-tokenizer\n---\n\n"
            f"{GENERATED_CARD_MARKER}\n\n"
            "# ActionCodec2\n\n"
            "This artifact contains physical vocabularies, BPE rules and action-space layouts.\n"
            "Install actioncodec2 and load the saved processor locally:\n\n"
            "```python\nfrom actioncodec2 import ActionCodec2\n"
            "from transformers import AutoProcessor\n\n"
            "print(ActionCodec2.describe_pretrained('./my_codec'))\n"
            "codec = AutoProcessor.from_pretrained(\n"
            "    './my_codec', action_space='single_eef_delta',\n)\n"
            "print(codec)\n```\n\n"
            "Choose action_space from your data's physical meaning, units and column order;\n"
            "the example layout is not an automatic selection. For a bound artifact, omit\n"
            "action_space to restore its saved binding. Pass your actual sampling rate as\n"
            "fps to encode/decode. Absolute layouts require current_state; second-order\n"
            "profiles require previous_first_order. See the artifact's layout guide below.\n"
        )
        path.write_text(
            guide + "\n\n## 当前 artifact 的实际布局与要求\n\n```text\n" + description + "\n```\n",
            encoding="utf-8",
        )
    return path