# AUTOGENERATED - modify shared_anyscale_util in root directory to make changes
import base64
from dataclasses import asdict, dataclass
import urllib.parse


BYOD_PREFIX = "custom_docker_"
BYOD_DISPLAY_PREFIX = "docker://"


def is_byod_id(id: str) -> bool:
    return id.startswith(BYOD_PREFIX)


@dataclass(frozen=True)
class BYODInfo:
    docker_image_name: str
    python_version: str
    ray_version: str

    def encode(self) -> str:
        return BYOD_PREFIX + base64.b64encode(
            urllib.parse.urlencode(asdict(self)).encode("ascii")
        ).decode("ascii")

    def get_name(self) -> str:
        return f"{BYOD_DISPLAY_PREFIX}{self.docker_image_name}"


def decode_byod(id: str) -> BYODInfo:
    if not is_byod_id(id):
        raise ValueError(f"Expected id to start with {BYOD_PREFIX}, got {id}")
    decoded = urllib.parse.parse_qs(
        base64.b64decode(id[len(BYOD_PREFIX) :]).decode("ascii")
    )
    return BYODInfo(
        decoded["docker_image_name"][0],
        decoded["python_version"][0],
        decoded["ray_version"][0],
    )
