"""Versioned image input normalization shared by training and served model encoding."""

FORMAT = "image-text-envelope-v1"


def normalize(state, image_count, mode="legacy"):
    if mode == "legacy":
        return state
    if mode != FORMAT:
        raise ValueError("unknown input format")
    # Only unwrap the documented text/images envelope produced by media.resolve.
    # Preserve arbitrary user structured data and all semantic fields.
    if (image_count and isinstance(state, dict) and set(state) == {"text", "images"}
            and isinstance(state["text"], str) and isinstance(state["images"], list)
            and len(state["images"]) == image_count
            and all(isinstance(item, dict) and set(item) == {"id", "index"}
                    and item["index"] == index + 1
                    for index, item in enumerate(state["images"]))):
        return state["text"]
    return state