//! Runtime configuration: the rust-server boot knobs
//! ([`RustServerServerArgs`]), the typed view of the scheduler's `server_args`
//! dump ([`ServerArgs`] / [`ModelConfig`]), and the [`RuntimeConfig`] pairing
//! them for `runtime::start`.

use std::net::SocketAddr;
use std::sync::Arc;

/// Boot knobs specific to the embedded rust server — none of these exist in
/// the Python `server_args` dump (see [`ServerArgs`]); they arrive as explicit
/// `Server::start` parameters.
#[derive(Clone, Debug)]
pub struct RustServerServerArgs {
    pub http_addr: SocketAddr,
    pub api_worker_num: usize,
    pub ingress_ring_cap: usize,
    pub egress_ring_cap: usize,
    pub channel_cap: usize,
    /// CPU core ids the pools pin to (e.g. this rank's NUMA-local cores minus
    /// the scheduler's reserved launch cores). `None` → run unpinned.
    pub cores: Option<Vec<usize>>,
}

impl Default for RustServerServerArgs {
    fn default() -> Self {
        Self {
            http_addr: "127.0.0.1:30000".parse().unwrap(),
            api_worker_num: 2,
            ingress_ring_cap: 8192,
            egress_ring_cap: 8192,
            channel_cap: 8192,
            cores: None,
        }
    }
}

#[derive(Clone, Debug)]
pub struct RuntimeConfig {
    /// Rust-server-only boot knobs (listen address, pool/ring sizes, pinning).
    pub rust_server_args: RustServerServerArgs,
    /// The scheduler's `server_args` dump (worker counts, tokenizer source,
    /// config-endpoint metadata). `Arc` so cloning the config (and, downstream,
    /// each `AppState`) is cheap; immutable after construction.
    pub server_args: Arc<ServerArgs>,
}

impl Default for RuntimeConfig {
    fn default() -> Self {
        Self {
            rust_server_args: RustServerServerArgs::default(),
            server_args: Arc::new(
                ServerArgs::from_json("{}").expect("empty server_args blob parses"),
            ),
        }
    }
}

/// The scheduler's startup blob (`RustServer._build_server_args`) parsed once into
/// typed fields: values are post-`__post_init__`; unrelated unknown keys are dropped.
#[derive(Debug, serde::Deserialize)]
pub struct ServerArgs {
    /// HF repo id / local dir of the model, reported by `/get_model_info`.
    #[serde(default)]
    pub model_path: String,
    /// Model name reported by `/v1/models` and `/server_info`.
    #[serde(default)]
    pub served_model_name: String,
    /// Tokenizer source (model dir / `tokenizer.json` / HF repo id). Empty only
    /// in minimal standalone blobs — then boot requires `skip_tokenizer_init`.
    #[serde(default)]
    pub tokenizer_path: String,
    /// HF revision, used only when `tokenizer_path` is a repo id. `None` → main.
    #[serde(default)]
    pub revision: Option<String>,
    /// HTTP bind address (see [`Self::bind`]).
    #[serde(default = "default_host")]
    pub host: String,
    #[serde(default = "default_port")]
    pub port: u16,
    /// Log levels driving the access log — uvicorn runs at
    /// `log_level_http or log_level` (see [`Self::http_access_log_enabled`]).
    #[serde(default = "default_log_level")]
    pub log_level: String,
    #[serde(default)]
    pub log_level_http: Option<String>,
    /// Optional built-in chat-template name or path to a Jinja/legacy JSON
    /// template file. Without an override, uses the tokenizer config template.
    #[serde(default)]
    pub chat_template: Option<String>,
    /// Parser selected by `--tool-call-parser`.
    #[serde(default)]
    pub tool_call_parser: Option<String>,
    /// Reasoning splitter selected by `--reasoning-parser` (e.g. deepseek-r1).
    /// When set, chat completions strip the model's reasoning markers out of
    /// `content` into `reasoning_content` — both unary and streaming.
    #[serde(default)]
    pub reasoning_parser: Option<String>,
    /// Python's global default for whether an SSE stream ends with a usage chunk.
    #[serde(default)]
    pub stream_response_default_include_usage: bool,
    /// Pinned tokenizer threads / detok shards (Python asserts both ≥ 1).
    #[serde(default = "default_worker_num")]
    pub tokenizer_worker_num: usize,
    #[serde(default = "default_worker_num")]
    pub detokenizer_worker_num: usize,
    /// Token-ids-in / token-ids-out mode: no tokenizer load, raw `output_ids`
    /// frames (drives the `Skip` detok backend and the ingress branch).
    #[serde(default)]
    pub skip_tokenizer_init: bool,
    /// Streamed `/generate` frames carry per-step deltas instead of cumulative
    /// text. Matches the Python `TokenizerManager`.
    #[serde(default)]
    pub incremental_streaming_output: bool,
    /// PD-disaggregation role: `"null"` (unified), `"prefill"`, or `"decode"`.
    /// (On prefill, the KV bootstrap registry is mounted on the api router —
    /// see [`Self::enable_pd_bootstrap`].)
    #[serde(default = "default_disaggregation_mode")]
    pub disaggregation_mode: String,
    /// The resolved Python `ModelConfig`, attached to the blob at dump time.
    #[serde(default)]
    pub model_config: ModelConfig,
    /// Default sampling params advertised by `/get_model_info`, verbatim from
    /// `server_args.preferred_sampling_params` (a JSON object or null).
    #[serde(default)]
    pub preferred_sampling_params: Option<serde_json::Value>,
    /// Over-long inputs are truncated to fit the context instead of 400ing, and
    /// `max_new_tokens` is clamped rather than rejected (Python
    /// `TokenizerManager._validate_one_request`).
    #[serde(default)]
    pub allow_auto_truncate: bool,
    /// `return_hidden_states` is refused unless the server was launched with it:
    /// the scheduler simply won't produce them, so the request would 200 with the
    /// field silently missing.
    #[serde(default)]
    pub enable_return_hidden_states: bool,
    /// Output slots reserved per request on top of its input (eagle stores draft
    /// tokens there). Not a `server_args` field — `TokenizerManager` derives it and
    /// `RustServer._build_server_args` stamps it in, so both sides count alike.
    #[serde(default)]
    pub num_reserved_tokens: u64,
    /// Launch-time stamps (not `server_args` fields): sglang package version
    /// and the scheduler-derived KV token capacity, reported by `/server_info`.
    #[serde(default)]
    pub version: Option<String>,
    #[serde(default)]
    pub max_total_num_tokens: Option<u64>,
}

/// The slice of the resolved Python `ModelConfig` the rust server reads.
#[derive(Debug, Default, serde::Deserialize)]
pub struct ModelConfig {
    /// Resolved context length (`max_model_len` in `/v1/models`); mandatory at
    /// boot ([`ServerArgs::validate_mandatory`]).
    #[serde(default)]
    pub context_len: Option<u64>,
    /// Bounds client-supplied token ids — ingress 400s out-of-vocab ids before
    /// they crash the scheduler's embedding lookup;  mandatory at
    /// boot ([`ServerArgs::validate_mandatory`]).
    #[serde(default)]
    pub vocab_size: Option<u64>,
    /// Whether the model accepts multimodal inputs. Gates the MM Encoding branch
    /// in tm-ingress; `false` silently ignores mm fields, as the Python
    /// `TokenizerManager` does with `mm_processor is None`.
    #[serde(default)]
    pub is_multimodal: bool,
    /// Resolved default sampling parameters, stamped by
    /// `RustServer._build_server_args` from Python's
    /// `ModelConfig.get_default_sampling_params()`. Already gated on
    /// `--sampling-defaults`: holds the model's generation_config.json values
    /// in "model" mode, and is empty in "openai" mode. Consumed when a chat
    /// request omits `temperature`/`top_p` — the conversion must not skip
    /// straight to the OpenAI terminal defaults.
    #[serde(default)]
    pub default_sampling_params: DefaultSamplingParams,
}

/// One `SamplingParams` field per key `get_default_sampling_params()` may emit
/// (`repetition_penalty`, `temperature`, `top_k`, `top_p`, `min_p`), filtered
/// to values the generation config actually sets — hence all `Option`.
///
/// `top_k` / `min_p` / `repetition_penalty` are parsed for parity with the
/// Python dict but not yet consumed: the Dynamo chat request type only carries
/// `temperature` and `top_p`, so the conversion resolves just those two.
#[derive(Debug, Default, serde::Deserialize)]
#[allow(dead_code)]
pub struct DefaultSamplingParams {
    #[serde(default)]
    pub temperature: Option<f64>,
    #[serde(default)]
    pub top_p: Option<f64>,
    #[serde(default)]
    pub top_k: Option<i64>,
    #[serde(default)]
    pub min_p: Option<f64>,
    #[serde(default)]
    pub repetition_penalty: Option<f64>,
}

fn join_host_port(host: &str, port: u16) -> String {
    if host.contains(':') && !host.starts_with('[') {
        format!("[{host}]:{port}") // bare IPv6 (`::`) needs brackets to bind
    } else {
        format!("{host}:{port}")
    }
}

fn default_host() -> String {
    "127.0.0.1".into()
}
fn default_port() -> u16 {
    30000
}
fn default_log_level() -> String {
    "info".into()
}
fn default_disaggregation_mode() -> String {
    "null".into()
}
fn default_worker_num() -> usize {
    1
}

impl ServerArgs {
    /// Parse the blob; errors on malformed JSON or a wrongly-typed field.
    pub fn from_json(s: &str) -> Result<Self, String> {
        serde_json::from_str(s).map_err(|e| e.to_string())
    }

    /// Fail fast at startup if a field an endpoint depends on is missing.
    pub fn validate_mandatory(&self) -> Result<(), String> {
        if self.served_model_name.is_empty() {
            return Err("no 'served_model_name' in server_args".into());
        }
        if self.model_config.context_len.is_none() {
            return Err("no resolvable context length (model_config.context_len)".into());
        }
        if self.model_config.vocab_size.is_none() {
            return Err("no resolvable vocab size (model_config.vocab_size)".into());
        }
        if !matches!(
            self.disaggregation_mode.as_str(),
            "null" | "prefill" | "decode"
        ) {
            return Err(format!(
                "unknown disaggregation_mode '{}' in server_args",
                self.disaggregation_mode
            ));
        }
        Ok(())
    }

    /// True on a prefill or decode node — requests need bootstrap routing.
    pub fn is_disaggregation(&self) -> bool {
        self.disaggregation_mode != "null"
    }

    /// Serve the PD KV bootstrap registry on the api listener: every prefill
    /// rust server hosts it, unconditionally — no extra topology gating. KV
    /// managers and decode nodes reach the registry at the resolved
    /// `disaggregation_bootstrap_port`, which rust-server mode aliases to the
    /// api port, so whichever prefill server that port names is the one that
    /// receives the registrations.
    pub fn enable_pd_bootstrap(&self) -> bool {
        self.disaggregation_mode == "prefill"
    }

    /// Whether the served model is multimodal, from the scheduler's dump. See
    /// [`ModelConfig::is_multimodal`].
    pub fn model_is_multimodal(&self) -> bool {
        self.model_config.is_multimodal
    }

    /// Bind address `host:port`. `host` is expected to be an IP — the result is
    /// parsed as a `SocketAddr`, so a bare IPv6 host gets bracketed.
    pub fn bind(&self) -> String {
        join_host_port(&self.host, self.port)
    }

    /// Whether the HTTP access log is emitted, mirroring the Python server:
    /// uvicorn runs at `log_level_http or log_level` and prints access lines
    /// only at info/debug. `--log-level-http warning` turns them off.
    pub fn http_access_log_enabled(&self) -> bool {
        let level = self
            .log_level_http
            .as_deref()
            .filter(|s| !s.is_empty())
            .unwrap_or(&self.log_level);
        matches!(
            level.to_ascii_lowercase().as_str(),
            "trace" | "debug" | "info"
        )
    }

    /// Pinned API threads for the embedded HTTP api-server. Python `server_args`
    /// has no such field — this is derived: enough to cover the widest pool.
    pub fn api_worker_num(&self) -> usize {
        4.max(self.tokenizer_worker_num)
            .max(self.detokenizer_worker_num)
    }
}
