"""A TypeSafe System One endpoint on Modal, served by an open kev checkpoint. See README.md."""

import os

import modal

# read at deploy time and baked into the image env, so the container resolves the same values
SETTINGS = {
    "SYSTEMONE_RUN": "jaredpalmer/kev-4b",
    "SYSTEMONE_HF_SECRET": "",
    "KEV_REF": "90990a5fac2995b9faa3190f7d437e84f2067768",
    "SYSTEMONE_NAME": "systemone",
    "SYSTEMONE_GPU": "L4",
    "SYSTEMONE_REGION": "us",
    "SYSTEMONE_GATHER_MS": "10",
    "SYSTEMONE_MIN_CONTAINERS": "0",
    "SYSTEMONE_SCALEDOWN_SECONDS": "300",
}
SETTINGS = {k: os.environ.get(k, v) for k, v in SETTINGS.items()}

KEV_REPO = "https://github.com/jaredpalmer/kev.git"
# transformers runs Qwen3.5's causal conv in reference torch without this kernel; the wheel matches torch 2.8 / cu12 / cp312
CAUSAL_CONV1D_WHEEL = (
    "https://github.com/Dao-AILab/causal-conv1d/releases/download/v1.7.0/"
    "causal_conv1d-1.7.0+cu12torch2.8cxx11abiTRUE-cp312-cp312-linux_x86_64.whl"
)
KEV_ROOT = "/kev"
HF = "/hf"

app = modal.App(SETTINGS["SYSTEMONE_NAME"])

image = (
    modal.Image.debian_slim(python_version="3.12")
    .apt_install("git")
    .run_commands(f"git clone {KEV_REPO} {KEV_ROOT} && git -C {KEV_ROOT} checkout --quiet {SETTINGS['KEV_REF']}")
    .uv_pip_install(f"kev[serve] @ file://{KEV_ROOT}")
    .uv_pip_install("flash-linear-attention", "triton>=3.7.1")
    .uv_pip_install(CAUSAL_CONV1D_WHEEL)
    .env(
        {
            "HF_HOME": HF,
            "HF_XET_HIGH_PERFORMANCE": "1",
            "HF_HUB_DISABLE_PROGRESS_BARS": "1",
            "TOKENIZERS_PARALLELISM": "false",
            "PYTHONUNBUFFERED": "1",
            "TRITON_CACHE_DIR": f"{HF}/triton-cache",
            "TRITON_CACHE_AUTOTUNING": "1",
            "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True",
        }
    )
    .env(SETTINGS)
    .add_local_python_source("systemone")
)

hf_cache = modal.Volume.from_name("kev-hf-cache", create_if_missing=True)

with image.imports():
    from systemone import serving


@app.cls(
    image=image,
    gpu=SETTINGS["SYSTEMONE_GPU"],
    region=[r for r in SETTINGS["SYSTEMONE_REGION"].split(",") if r] or None,
    cpu=2,
    memory=(16384, 65536),
    volumes={HF: hf_cache},
    secrets=[modal.Secret.from_name(SETTINGS["SYSTEMONE_HF_SECRET"])] if SETTINGS["SYSTEMONE_HF_SECRET"] else [],
    min_containers=int(SETTINGS["SYSTEMONE_MIN_CONTAINERS"]),
    scaledown_window=int(SETTINGS["SYSTEMONE_SCALEDOWN_SECONDS"]),
    timeout=600,
    startup_timeout=900,
)
@modal.concurrent(max_inputs=8)
class Serve:
    @modal.enter()
    def load(self):
        self.api = serving.build(SETTINGS["SYSTEMONE_RUN"], float(SETTINGS["SYSTEMONE_GATHER_MS"]) / 1000)
        hf_cache.commit()

    @modal.asgi_app(label=SETTINGS["SYSTEMONE_NAME"], requires_proxy_auth=True)
    def web(self):
        return self.api