"""Hugging Face ZeroGPU compatibility helpers."""

from __future__ import annotations

from typing import Protocol


class CodonUtilsModule(Protocol):
    ORGANISM2ID: dict[str, int]
    ID2ORGANISM: dict[int, str]
    NUM_ORGANISMS: int


ADP1_NAME = "Acinetobacter baylyi ADP1"
ADP1_ID = 164


def register_adp1_organism(codon_utils: CodonUtilsModule) -> None:
    """Expose the trained ADP1 embedding in the upstream 164-organism package."""
    current_id = codon_utils.ORGANISM2ID.get(ADP1_NAME)
    current_name = codon_utils.ID2ORGANISM.get(ADP1_ID)
    if current_id not in (None, ADP1_ID):
        raise RuntimeError(f"{ADP1_NAME} is unexpectedly assigned ID {current_id}.")
    if current_name not in (None, ADP1_NAME):
        raise RuntimeError(
            f"Organism ID {ADP1_ID} is already assigned to {current_name}."
        )

    codon_utils.ORGANISM2ID[ADP1_NAME] = ADP1_ID
    codon_utils.ID2ORGANISM[ADP1_ID] = ADP1_NAME
    codon_utils.NUM_ORGANISMS = max(codon_utils.NUM_ORGANISMS, ADP1_ID + 1)