"""Public Gradio API for private-model PFF-PK inference."""

from __future__ import annotations

import os
from pathlib import Path

import gradio as gr
from huggingface_hub import hf_hub_download

ROOT = Path(__file__).resolve().parent
MODEL_REPO_ID = os.environ.get("MODEL_REPO_ID", "").strip()
MODEL_REVISION = os.environ.get("MODEL_REVISION", "").strip()
if not MODEL_REPO_ID or not MODEL_REVISION:
    raise RuntimeError("MODEL_REPO_ID and MODEL_REVISION must be configured")

token = os.environ.get("HF_TOKEN") or True
checkpoint = hf_hub_download(
    repo_id=MODEL_REPO_ID,
    filename="models/pythia-dose/model.ckpt",
    revision=MODEL_REVISION,
    token=token,
)
config = hf_hub_download(
    repo_id=MODEL_REPO_ID,
    filename="models/pythia-dose/config.yaml",
    revision=MODEL_REVISION,
    token=token,
)
pythia_checkpoint = hf_hub_download(
    repo_id=MODEL_REPO_ID,
    filename="models/pythia/model.ckpt",
    revision=MODEL_REVISION,
    token=token,
)
pythia_config = hf_hub_download(
    repo_id=MODEL_REPO_ID,
    filename="models/pythia/config.yaml",
    revision=MODEL_REVISION,
    token=token,
)
covariates_checkpoint = hf_hub_download(
    repo_id=MODEL_REPO_ID,
    filename="models/pythia-covariates/model.ckpt",
    revision=MODEL_REVISION,
    token=token,
)
covariates_config = hf_hub_download(
    repo_id=MODEL_REPO_ID,
    filename="models/pythia-covariates/config.yaml",
    revision=MODEL_REVISION,
    token=token,
)

os.environ["PFF_REPO"] = str(ROOT)
os.environ["PFF_DOSE_CHECKPOINT"] = checkpoint
os.environ["PFF_DOSE_CONFIG"] = config
os.environ["PFF_PYTHIA_CHECKPOINT"] = pythia_checkpoint
os.environ["PFF_PYTHIA_CONFIG"] = pythia_config
os.environ["PFF_COVARIATES_CHECKPOINT"] = covariates_checkpoint
os.environ["PFF_COVARIATES_CONFIG"] = covariates_config
os.environ.setdefault("PFF_CACHE_ROOT", "/tmp/pff-inference-cache")
os.environ.setdefault("PFF_CPU_THREADS", "2")

from services.inference.pff_service import cached_inference, service_status  # noqa: E402
from services.inference.synthetic_service import synthetic_request  # noqa: E402

PUBLIC_SOLVER = {"method": "heun", "steps": 8}


def health() -> dict:
    """Return only non-sensitive deployment metadata."""
    return service_status()


def synthetic(payload: dict) -> dict:
    """Generate with the bundled, provenance-recorded synthetic_priors package."""
    try:
        return synthetic_request(payload)
    except Exception as error:
        raise gr.Error(str(error)) from error


def inference(payload: dict) -> dict:
    """Generate a PK cohort from a validated dashboard request."""
    if not isinstance(payload, dict):
        raise gr.Error("The inference request must be a JSON object")
    try:
        solver = payload.get("solver") or PUBLIC_SOLVER
        if not isinstance(solver, dict):
            raise ValueError("The public demo uses fixed Heun integration with 8 steps")
        if (
            str(solver.get("method", "")).lower() != PUBLIC_SOLVER["method"]
            or float(solver.get("steps", 0)) != PUBLIC_SOLVER["steps"]
        ):
            raise ValueError("The public demo uses fixed Heun integration with 8 steps")
        return cached_inference({**payload, "solver": PUBLIC_SOLVER.copy()})
    except (ValueError, KeyError, TypeError, FileNotFoundError) as error:
        raise gr.Error(str(error)) from error
    except Exception as error:
        raise gr.Error("PFF inference failed") from error


with gr.Blocks(title="PFF-PK inference API") as demo:
    gr.Markdown(
        "# PFF-PK inference API\n"
        "Server-side zero-shot PK cohort generation for the "
        "[PFF dashboard](https://dfaroughy.github.io/pff-pk-dashboard/empirical/)."
    )
    request = gr.JSON(label="Inference request", value={})
    result = gr.JSON(label="Inference response")
    run = gr.Button("Run inference", variant="primary")
    run.click(
        inference,
        inputs=request,
        outputs=result,
        api_name="inference",
        concurrency_limit=1,
        concurrency_id="pff-inference",
    )
    status = gr.JSON(label="Service status")
    check = gr.Button("Check service")
    check.click(health, inputs=None, outputs=status, api_name="health")
    synthetic_run = gr.Button("Generate synthetic cohort")
    synthetic_run.click(synthetic, inputs=request, outputs=result, api_name="synthetic",
                        concurrency_limit=1, concurrency_id="pff-inference")

demo.queue(max_size=32, default_concurrency_limit=1)

if __name__ == "__main__":
    demo.launch()