"""Frugal PCD validation/benchmarking and optional authenticated serving on Modal.

modal run lfm25_pcd_modal.py --task prepare
modal run lfm25_pcd_modal.py --task validate
modal run lfm25_pcd_modal.py --task benchmark
Publication is a separate, explicit --task publish operation.
"""

import json
from pathlib import Path

import modal

app = modal.App("lfm25-26b-pcd")
root = Path(__file__).parent
cache = modal.Volume.from_name("huggingface-cache", create_if_missing=True)
results = modal.Volume.from_name("lfm25-26b-pcd-results", create_if_missing=True)
base = modal.Image.debian_slim(python_version="3.11").env(
    {
        "HF_HOME": "/root/.cache/huggingface",
        "TOKENIZERS_PARALLELISM": "false",
        "HF_HUB_DISABLE_TELEMETRY": "1",
        "OMP_NUM_THREADS": "4",
    }
)
cpu_image = base.uv_pip_install("huggingface-hub==1.31.0", "pyyaml==6.0.3").add_local_python_source(
    "pcd"
)
gpu_image = (
    base.pip_install_from_requirements(str(root / "requirements-pcd.txt"))
    .add_local_python_source("pcd")
    .add_local_dir(root / "tests" / "pcd", "/root/tests/pcd")
)


@app.function(
    image=cpu_image,
    secrets=[modal.Secret.from_name("huggingface")],
    volumes={"/root/.cache/huggingface": cache},
    timeout=900,
    memory=4096,
    max_containers=1,
)
def prepare():
    from pcd.publishing import prepare as prepare_snapshot

    info = prepare_snapshot()
    cache.commit()
    return info


@app.cls(
    image=gpu_image,
    gpu="L40S",
    cpu=4,
    memory=16384,
    timeout=900,
    min_containers=0,
    max_containers=1,
    scaledown_window=30,
    volumes={"/root/.cache/huggingface": cache, "/results": results},
)
class PCDModel:
    precision: str = modal.parameter(default="float16")

    @modal.enter()
    def load(self):
        import torch

        from pcd.engine import Engine

        if self.precision not in {"float16", "bfloat16", "float32"}:
            raise ValueError("unsupported precision")
        torch.set_num_threads(4)
        torch.set_float32_matmul_precision("highest")
        self.engine = Engine(
            device="cuda", dtype=self.precision, attention="sdpa", local_files_only=True
        )

    @modal.method()
    def run(self, task: str, suite: str = "diagnostic", repeats: int = 3):
        from pcd.benchmark import execute

        report = execute(self.engine, task=task, suite=suite, repeats=repeats)
        path = Path("/results") / (report["run_id"] + ".json")
        path.write_text(json.dumps(report, indent=2, ensure_ascii=False) + "\n")
        results.commit()
        return report

    @modal.method()
    def infer(self, context: str, schema: dict, mode: str = "token"):
        return self.engine.constrained(context, schema, mode=mode)

    @modal.fastapi_endpoint(method="POST", requires_proxy_auth=True, docs=False)
    def extract(self, payload: dict):
        import jsonschema
        from fastapi import HTTPException

        try:
            if set(payload) - {"context", "schema", "mode"}:
                raise ValueError("unsupported request keys")
            return self.engine.constrained(
                payload["context"], payload["schema"], mode=payload.get("mode", "token")
            )
        except (ValueError, KeyError, TypeError, jsonschema.SchemaError) as exc:
            raise HTTPException(status_code=422, detail=str(exc)) from None


release_dir = root / "release" / "pcd"
publish_image = (
    cpu_image.add_local_dir(release_dir, "/release-src") if release_dir.is_dir() else cpu_image
)


@app.function(
    image=publish_image,
    secrets=[modal.Secret.from_name("huggingface")],
    volumes={"/root/.cache/huggingface": cache, "/results": results},
    timeout=1800,
    cpu=4,
    memory=8192,
    max_containers=1,
)
def publish(public: bool = False):
    from pcd.publishing import publish as publish_bundle

    result = publish_bundle(public=public)
    results.commit()
    return result


@app.local_entrypoint()
def main(
    task: str = "validate",
    suite: str = "diagnostic",
    repeats: int = 3,
    precision: str = "float16",
    public: bool = False,
):
    if task not in {"prepare", "validate", "benchmark", "probe", "publish"}:
        raise ValueError("task must be prepare, validate, benchmark, probe, or publish")
    if repeats < 1 or repeats > 10:
        raise ValueError("repeats must be between 1 and 10")
    if public and task != "publish":
        raise ValueError("--public only applies to --task publish")
    if task == "publish":
        if not release_dir.is_dir():
            raise ValueError("stage the release with scripts/build_pcd_release.py first")
        report = publish.remote(public)
    elif task == "prepare":
        report = prepare.remote()
    else:
        report = PCDModel(precision=precision).run.remote(task, suite, repeats)
    output = root / "results" / "pcd"
    output.mkdir(parents=True, exist_ok=True)
    path = output / (report.get("run_id", task) + ".json")
    path.write_text(json.dumps(report, indent=2, ensure_ascii=False) + "\n")
    print(json.dumps({"report": str(path), "summary": report.get("summary", report)}, indent=2))