//! The `Recipe` type (parsed from `recipes/<name>/recipe.toml`) and
//! [`EnvValue`] — the `[env]` value type: a literal string, or a reference
//! into a recipe field (`{ from = "port" }`) resolved at serve time.

use crate::config::types::{
    ComposedModel, DeployTarget, ExtraMount, RecipeBuild, default_engine, default_port,
    default_shm_size, default_sync_image, default_tensor_parallel_size,
};
use crate::pre_serve::{DistributedBackend, ModsTable, PreServeCommand, RecipeArgs, VenvSetup};
use anyhow::{Result, bail};
use serde::Deserialize;
use std::collections::HashMap;

/// An `[env]` entry value: either a literal string, or a reference into a
/// recipe field resolved at serve time (`{ from = "port" }`). Untagged so
/// both forms mix freely in one `[env]` table.
#[derive(Clone, Debug, PartialEq, Deserialize)]
#[serde(untagged)]
pub enum EnvValue {
    Plain(String),
    Ref { from: String },
}

impl From<String> for EnvValue {
    fn from(s: String) -> Self {
        EnvValue::Plain(s)
    }
}

impl From<&str> for EnvValue {
    fn from(s: &str) -> Self {
        EnvValue::Plain(s.to_string())
    }
}

impl EnvValue {
    /// Resolve to a literal value. A `Ref` entry pulls its value from a
    /// recipe field (`from = "port"` → the recipe's `port`), so the value is
    /// declared once instead of being duplicated between `port = ...` and
    /// `[env]`. Unknown targets are an error listing the supported ones.
    pub fn resolve(&self, key: &str, recipe: &Recipe) -> Result<String> {
        match self {
            EnvValue::Plain(s) => Ok(s.clone()),
            EnvValue::Ref { from } => match from.as_str() {
                "port" => Ok(recipe.port.to_string()),
                _ => {
                    bail!("recipe [env] '{key}': unknown 'from' target '{from}' (supported: port)")
                }
            },
        }
    }
}

#[derive(Clone, Deserialize)]
pub struct Recipe {
    #[serde(default = "default_engine")]
    pub engine: String,
    pub name: String,
    #[serde(default)]
    pub docker_image: String,
    #[serde(default)]
    pub model: String,
    /// Parent model this recipe's category ranks under (medals, tree
    /// grouping). Optional — defaults to `model` itself. Set it when a
    /// quantized/forked repo (`local-inference-lab/GLM-5.3-Flash-NVFP4-Spark`)
    /// should compete inside its base model's category
    /// (`zai-org/GLM-5.3-Flash`).
    #[serde(default)]
    pub parent_model: Option<String>,
    #[serde(default = "default_port")]
    pub port: u16,
    #[serde(default)]
    pub required_nodes: Option<u32>,
    #[serde(default)]
    pub pipeline_parallel_size: Option<u32>,
    #[serde(default = "default_tensor_parallel_size")]
    pub tensor_parallel_size: u32,
    #[serde(default)]
    pub args: RecipeArgs,
    #[serde(default)]
    pub env: HashMap<String, EnvValue>,
    #[serde(default)]
    pub cluster: Option<String>,
    #[serde(default)]
    pub entrypoint: Vec<String>,
    pub distributed_backend: Option<DistributedBackend>,
    #[serde(default = "default_shm_size")]
    pub shm_size: String,
    #[serde(default)]
    pub before_serve: Vec<PreServeCommand>,
    /// Commands run per node **after** stopping containers (`sm121 stop`);
    /// same shape as [`before_serve`] (raw shell or `Script` aliases),
    /// parallel across nodes. E.g. `after_stop = [{ script =
    /// "restore_dgx_services" }]`.
    #[serde(default)]
    pub after_stop: Vec<PreServeCommand>,
    /// Whether this recipe should always be served single-node on the first
    /// free node found (solo mode). When set, `sm121 serve` ignores
    /// `required_nodes` / multi-node cluster selection and launches exactly
    /// one container with no distributed flags. The CLI `--solo`/`--no-solo`
    /// flag overrides this per-invocation.
    #[serde(default)]
    pub solo: Option<bool>,
    /// Deployment target — `docker` (default) or `tmux`. When `tmux`, the
    /// `serve` command launches via detached `tmux` sessions instead of
    /// `docker run`. Docker-only features (mods, composed_model, image
    /// sync) are rejected in tmux mode.
    #[serde(default)]
    pub deploy: DeployTarget,
    /// Python virtualenv path for tmux deployment. When set, the tmux
    /// launch command prepends `source <venv>/bin/activate &&` before
    /// the entrypoint, so the engine runs from the specified venv instead
    /// of the system Python. Docker-only (ignored in Docker mode — the
    /// container has its own Python).
    ///
    /// Example: `venv = "~/venvs/vllm025"`
    #[serde(default)]
    pub venv: Option<String>,
    /// Shell commands to create and initialize the venv at `venv` if it
    /// does not already exist on a node. Either an inline multi-line string
    /// (one command per line; blank lines and `#` comments are skipped) or
    /// a `Vec<String>` list. Only used in tmux deployment mode.
    ///
    /// The launcher checks `test -f <venv>/bin/activate` on each node
    /// before serve; if absent, it runs these commands in order, then
    /// re-checks. A failure aborts the serve.
    ///
    /// Example (inline): `uv venv ~/venvs/vllm025 -p 3.12 && uv pip install ...`
    /// Example (list): `["uv venv ~/venvs/vllm025 -p 3.12", "..."]`
    #[serde(default)]
    pub venv_setup: VenvSetup,
    /// Path to a shell script (relative to the recipe directory) that creates
    /// and initializes the venv. When set, the launcher reads the file from
    /// `recipes/<recipe_name>/<venv_setup_script>` and runs it on each node
    /// if the venv doesn't exist. Mutually exclusive with `venv_setup` — if
    /// both are set, `venv_setup_script` takes precedence.
    ///
    /// Example: `venv_setup_script = "venv_setup.sh"`
    #[serde(default)]
    pub venv_setup_script: Option<String>,
    #[serde(default)]
    pub mods: ModsTable,
    #[serde(default = "default_sync_image")]
    pub sync_image: bool,
    /// Build directory for docker-images/ (default: derived from recipe name).
    #[serde(default)]
    pub docker_image_dir: Option<String>,
    /// Recipe-driven `versions.env` overrides merged at build time (recipe
    /// wins, after the profile) and folded into the build checksum, so a pin
    /// change here triggers a rebuild without editing the image tree. Keys
    /// must match the `versions.env` names the `Dockerfile` declares as ARG.
    #[serde(default)]
    pub build: RecipeBuild,
    /// Named version profile: `docker-images/<dir>/profiles/<name>.env` (same
    /// `KEY=VAL` format as `versions.env`). Merged over the image's pins
    /// (profile wins) then recipe `[build.versions]` (recipe wins) before
    /// build, folded into the build checksum. Lets a recipe reference a full
    /// pin set (an engine generation: vLLM + b12x + LMCache cut together) by
    /// name instead of copy-pasting every pin. Docker-only.
    #[serde(default)]
    pub build_profile: Option<String>,
    /// When set, the launcher builds a /tmp composed-model dir per-node
    /// (symlinks into the HF cache + generated files), bind-mounts it at
    /// `mount`, and rewrites `model` to that path. Lets a recipe combine
    /// multiple HF repos (e.g. a base checkpoint + an expert overlay)
    /// without mutating the HF cache. See [`ComposedModel`].
    ///
    /// Docker-only: rejected when `deploy = "tmux"`.
    #[serde(default)]
    pub composed_model: Option<ComposedModel>,
    /// Extra bind-mounts for `docker run` (e.g. persistent JIT cache). Host
    /// path resolves from `Config.mount_aliases` when `alias` is set (alias
    /// wins); otherwise falls back to `host` (`~` expanded). One of the two
    /// is required.
    #[serde(default)]
    pub extra_mounts: Vec<ExtraMount>,
}

impl Recipe {
    /// All `[env]` entries resolved to literal values — `Ref` entries pull
    /// their value from the recipe's own fields (fail-fast on unknown
    /// targets). Docker `-e` and tmux `export` use this.
    pub fn resolved_env(&self) -> Result<HashMap<String, String>> {
        let mut out = HashMap::new();
        for (k, v) in &self.env {
            out.insert(k.clone(), v.resolve(k, self)?);
        }
        Ok(out)
    }
}