"""Only enable a prebuilt worker after a complete, measured acceptance run."""

import hashlib
import json
import math
import re
from pathlib import Path

from flora_captioner import FLORA
from huggingface_hub.utils import validate_repo_id

IDENTITY_FILES = {
    "worker.py",
    "subset.py",
    "flora_captioner.py",
    "flora_acceleration.py",
    "flora_batcher.py",
    "flora_speed_entry.py",
    "studio_runtime.py",
    "pipeline-version.json",
}
STUDIO_RUNTIME_FILES = {
    "Dockerfile",
    "worker-requirements.txt",
    "flora-requirements.txt",
    "flora-acceleration-requirements.txt",
    "build_causal_conv1d.py",
    "bootstrap.py",
    "flora_studio_entry.py",
}


def validated_worker_image(captioner, features, path=None):
    if captioner != "flora" or set(features) != {"Subtasks"}:
        return None
    path = path or Path(__file__).with_name("flora-worker-image.json")
    if not path.exists():
        return None
    profile = json.loads(path.read_text())
    if type(profile.get("batch_size", 4)) is not int or profile.get("batch_size", 4) not in (1, 2, 4, 8):
        raise ValueError("Accelerated Flora batch size is not validated")
    if (
        type(profile.get("episode_parallelism", 1)) is not int
        or not 1 <= profile.get("episode_parallelism", 1) <= 8
    ):
        raise ValueError("Accelerated Flora episode parallelism is not validated")
    if profile.get("generation_batch_size", profile.get("batch_size", 4)) not in (1, 2, 4, 8, 16, 32, 64):
        raise ValueError("Accelerated Flora GPU batch size is not validated")
    accepted = profile.get("acceptance", {})
    seconds = accepted.get("end_to_end_seconds")
    if (
        accepted.get("episodes") != 25
        or accepted.get("independent") is not True
        or accepted.get("fused_kernels") is not True
        or not isinstance(seconds, (int, float))
        or isinstance(seconds, bool)
        or not math.isfinite(seconds)
        or not 0 < seconds < 1500
        or profile.get("captioner_profile") != FLORA
    ):
        raise ValueError("Accelerated Flora worker has not passed its acceptance run")
    validate_repo_id(profile["repo"])
    source = profile.get("image_source", "external")
    if source not in {"external", "studio"}:
        raise ValueError("Unknown Flora image source")
    if source == "external" and not re.fullmatch(r"[a-f0-9]{40}", profile["revision"]):
        raise ValueError("Accelerated Flora image revision must be pinned")
    hashes = profile.get("expected_code", {})
    expected_files = IDENTITY_FILES | ({"flora_vision.py"} if profile.get("vision_frame_batching") else set())
    if set(hashes) != expected_files or any(
        not re.fullmatch(r"[a-f0-9]{64}", value) for value in hashes.values()
    ):
        raise ValueError("Accelerated Flora code checksums are incomplete")
    if source == "studio":
        runtime_files = profile.get("runtime_files", {})
        if set(runtime_files) != STUDIO_RUNTIME_FILES or any(
            not re.fullmatch(r"[a-f0-9]{64}", value) for value in runtime_files.values()
        ):
            raise ValueError("Studio's CUDA build checksums are incomplete")
        if set(profile.get("kernel_versions", {})) != {
            "torch",
            "transformers",
            "flash-linear-attention",
            "causal-conv1d",
        }:
            raise ValueError("Studio's CUDA dependency versions are incomplete")
        # A self-hosted image cannot embed its own Git SHA. Verify the tested
        # worker files locally; record the actual built revision at launch.
        for name, expected in {**hashes, **runtime_files}.items():
            if (
                Path(name).name != name
                or hashlib.sha256((path.parent / name).read_bytes()).hexdigest() != expected
            ):
                raise ValueError("Studio's Flora runtime changed; validate it before launching")
    return profile