ImageApi/src/enhance/config.rs

//! Enhance configuration (#192): one `enhance.toml`, next to `llm.toml`.
//!
//! Path from `ENHANCE_CONFIG_PATH`, else `./enhance.toml`. A missing file at
//! the default path means Enhance is simply off (no `url`); a missing
//! explicit path is an error. Nothing here is a secret: ComfyUI has no auth
//! and stays LAN-only. `enhance.example.toml` is the annotated schema (a
//! test parses it).

use std::path::{Path, PathBuf};

use anyhow::{Context as _, Result, bail};
use serde::Deserialize;

use super::path_map::PathMap;

const DEFAULT_PATH: &str = "enhance.toml";

/// How the app's text combines with the template's own prompt.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize, serde::Serialize, Default)]
#[serde(rename_all = "lowercase")]
pub enum PromptMode {
    /// The template's prompt is a full base instruction; the app's text is
    /// appended after `prompt_join`. Empty text leaves the prompt as is.
    #[default]
    Append,
    /// The app's text *is* the prompt (and is required).
    Replace,
}

#[derive(Debug, Clone, PartialEq, Deserialize)]
#[serde(default, deny_unknown_fields)]
pub struct EnhanceConfig {
    /// ComfyUI base URL, e.g. `http://192.168.10.218:8188`. Unset = Enhance
    /// is disabled (`/enhance/templates` reports `available: false`).
    pub url: Option<String>,
    /// Per-request timeout for ComfyUI calls.
    pub timeout_secs: u64,
    /// How long the `available` probe result is cached.
    pub probe_ttl_secs: u64,
    /// Job poll interval.
    pub poll_secs: u64,
    /// Idle time (empty queue, newest history entry this old) before
    /// ImageApi calls `POST /free` once. 0 disables it.
    pub idle_free_secs: u64,
    /// An active job whose ComfyUI stays unreachable this long becomes
    /// `unavailable`.
    pub unreachable_grace_secs: u64,
    /// Template preselected by clients; default = the first one.
    pub default_template: Option<String>,
    /// Clients warn above this many megapixels in the original.
    pub warn_above_megapixels: f64,
    /// Library path prefix <-> ComfyUI path prefix. Longest prefix wins.
    pub path_map: Vec<PathMapping>,
    #[serde(rename = "template")]
    pub templates: Vec<TemplateConfig>,
}

impl Default for EnhanceConfig {
    fn default() -> Self {
        Self {
            url: None,
            timeout_secs: 15,
            probe_ttl_secs: 10,
            poll_secs: 2,
            idle_free_secs: 300,
            unreachable_grace_secs: 120,
            default_template: None,
            warn_above_megapixels: 1.5,
            path_map: Vec::new(),
            templates: Vec::new(),
        }
    }
}

#[derive(Debug, Clone, PartialEq, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct PathMapping {
    /// Prefix as ImageApi sees it, e.g. `/media/nas/`.
    pub from: String,
    /// Prefix as ComfyUI sees it, e.g. `Z:\` or `\\nas\`.
    pub to: String,
}

#[derive(Debug, Clone, PartialEq, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct TemplateConfig {
    /// Stable id clients send (`POST /enhance { template }`).
    pub name: String,
    pub label: String,
    #[serde(default)]
    pub description: String,
    /// API-format workflow JSON (ComfyUI -> Workflow -> Export (API)).
    /// Relative paths resolve against the config file's directory.
    pub file: PathBuf,
    #[serde(default)]
    pub prompt_mode: PromptMode,
    /// Text between the template's prompt and the app's text (append mode).
    #[serde(default = "default_join")]
    pub prompt_join: String,
    /// Approximate output size, for the client's large-original warning.
    #[serde(default = "default_output_mp")]
    pub output_megapixels: f64,
    /// Optional LoRAs the app can switch on or off (#218), each matching the
    /// workflow nodes titled `imageapi:lora:<id>`.
    #[serde(rename = "lora", default)]
    pub loras: Vec<LoraConfig>,
}

/// One LoRA switch (`[[template.lora]]`).
#[derive(Debug, Clone, PartialEq, Deserialize, serde::Serialize)]
#[serde(deny_unknown_fields)]
pub struct LoraConfig {
    /// Matches the nodes titled `imageapi:lora:<id>` (Wan's high and low
    /// pair share one id); clients send it in `POST /enhance { loras }`.
    pub id: String,
    pub label: String,
    /// On when the request doesn't say.
    #[serde(default)]
    pub default: bool,
}

/// Template names and LoRA ids: non-empty, letters, digits, `_` and `-`.
fn valid_id(s: &str) -> bool {
    !s.is_empty()
        && s.chars()
            .all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-')
}

fn default_join() -> String {
    " ".into()
}

fn default_output_mp() -> f64 {
    1.0
}

impl EnhanceConfig {
    pub fn validate(&self) -> Result<()> {
        if let Some(url) = &self.url
            && !(url.starts_with("http://") || url.starts_with("https://"))
        {
            bail!("url must start with http:// or https://, got {url:?}");
        }
        for (key, v) in [
            ("timeout_secs", self.timeout_secs),
            ("probe_ttl_secs", self.probe_ttl_secs),
            ("poll_secs", self.poll_secs),
            ("unreachable_grace_secs", self.unreachable_grace_secs),
        ] {
            if v == 0 {
                bail!("{key} must be greater than 0");
            }
        }
        if self.warn_above_megapixels.is_nan() || self.warn_above_megapixels <= 0.0 {
            bail!("warn_above_megapixels must be greater than 0");
        }
        let mut seen = std::collections::HashSet::new();
        for t in &self.templates {
            if !valid_id(&t.name) {
                bail!(
                    "template name {:?} must be non-empty and use only letters, digits, '_' and '-'",
                    t.name
                );
            }
            if !seen.insert(t.name.as_str()) {
                bail!("duplicate template name {:?}", t.name);
            }
            if t.label.trim().is_empty() {
                bail!("template {:?} needs a label", t.name);
            }
            if t.output_megapixels.is_nan() || t.output_megapixels <= 0.0 {
                bail!("template {:?}: output_megapixels must be > 0", t.name);
            }
            let mut ids = std::collections::HashSet::new();
            for l in &t.loras {
                if !valid_id(&l.id) {
                    bail!(
                        "template {:?}: lora id {:?} must be non-empty and use only letters, digits, '_' and '-'",
                        t.name,
                        l.id
                    );
                }
                if !ids.insert(l.id.as_str()) {
                    bail!("template {:?}: duplicate lora id {:?}", t.name, l.id);
                }
                if l.label.trim().is_empty() {
                    bail!("template {:?}: lora {:?} needs a label", t.name, l.id);
                }
            }
        }
        if let Some(d) = &self.default_template
            && !self.templates.iter().any(|t| &t.name == d)
        {
            bail!("default_template {d:?} is not a configured template");
        }
        // Parses every mapping (rejects empty prefixes).
        PathMap::new(&self.path_map)?;
        Ok(())
    }

    /// The preselected template's name.
    pub fn default_name(&self) -> Option<&str> {
        self.default_template
            .as_deref()
            .or_else(|| self.templates.first().map(|t| t.name.as_str()))
    }
}

/// Where the config was read from, and what it said.
#[derive(Debug)]
pub struct Loaded {
    pub config: EnhanceConfig,
    /// Directory relative template `file`s resolve against.
    pub base_dir: PathBuf,
    /// The file, when one was read.
    pub path: Option<PathBuf>,
}

/// Read `enhance.toml` (see module docs). A bad file is an error so a typo
/// fails startup instead of silently disabling the feature.
pub fn load() -> Result<Loaded> {
    let (path, explicit) = match std::env::var("ENHANCE_CONFIG_PATH") {
        Ok(p) if !p.trim().is_empty() => (PathBuf::from(p.trim()), true),
        _ => (PathBuf::from(DEFAULT_PATH), false),
    };
    match std::fs::read_to_string(&path) {
        Ok(text) => {
            let config = parse(&text).with_context(|| format!("in {}", path.display()))?;
            let base_dir = path
                .parent()
                .filter(|p| !p.as_os_str().is_empty())
                .map(Path::to_path_buf)
                .unwrap_or_else(|| PathBuf::from("."));
            Ok(Loaded {
                config,
                base_dir,
                path: Some(path),
            })
        }
        Err(e) if e.kind() == std::io::ErrorKind::NotFound && !explicit => Ok(Loaded {
            config: EnhanceConfig::default(),
            base_dir: PathBuf::from("."),
            path: None,
        }),
        Err(e) => Err(e).with_context(|| format!("reading enhance config {}", path.display())),
    }
}

pub fn parse(text: &str) -> Result<EnhanceConfig> {
    let config: EnhanceConfig = toml::from_str(text)?;
    config.validate()?;
    Ok(config)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn example_file_parses() {
        let cfg = parse(include_str!("../../enhance.example.toml")).unwrap();
        assert!(cfg.url.is_some());
        assert_eq!(cfg.templates[0].name, "restore");
        assert_eq!(cfg.default_name(), Some("restore"));
        assert!(
            cfg.templates
                .iter()
                .any(|t| t.prompt_mode == PromptMode::Replace)
        );
    }

    #[test]
    fn empty_file_is_disabled_defaults() {
        let cfg = parse("").unwrap();
        assert_eq!(cfg, EnhanceConfig::default());
        assert!(cfg.url.is_none());
    }

    #[test]
    fn rejects_unknown_keys_and_bad_values() {
        assert!(parse("nope = 1").is_err());
        assert!(parse("url = \"ftp://x\"").is_err());
        assert!(parse("poll_secs = 0").is_err());
        assert!(parse("default_template = \"x\"").is_err());
        let dup = "[[template]]\nname='a'\nlabel='A'\nfile='a.json'\n[[template]]\nname='a'\nlabel='B'\nfile='b.json'";
        assert!(parse(dup).is_err());
        assert!(parse("[[template]]\nname='a b'\nlabel='A'\nfile='a.json'").is_err());
        assert!(parse("[[path_map]]\nfrom=''\nto='Z:\\\\'").is_err());
    }

    #[test]
    fn lora_switches_parse_and_validate() {
        let ok = "[[template]]\nname='a'\nlabel='A'\nfile='a.json'\n\
                  [[template.lora]]\nid='style'\nlabel='Style'\n\
                  [[template.lora]]\nid='detail'\nlabel='Detail'\ndefault=true";
        let cfg = parse(ok).unwrap();
        let l = &cfg.templates[0].loras;
        assert_eq!((l[0].id.as_str(), l[0].default), ("style", false));
        assert_eq!((l[1].id.as_str(), l[1].default), ("detail", true));

        let base = "[[template]]\nname='a'\nlabel='A'\nfile='a.json'\n";
        let dup = format!(
            "{base}[[template.lora]]\nid='x'\nlabel='X'\n[[template.lora]]\nid='x'\nlabel='Y'"
        );
        assert!(parse(&dup).is_err());
        assert!(parse(&format!("{base}[[template.lora]]\nid='a b'\nlabel='X'")).is_err());
        assert!(parse(&format!("{base}[[template.lora]]\nid='x'\nlabel=' '")).is_err());
        assert!(
            parse(&format!(
                "{base}[[template.lora]]\nid='x'\nlabel='X'\nstrength=1"
            ))
            .is_err()
        );
    }

    #[test]
    fn template_defaults() {
        let cfg = parse("[[template]]\nname='a'\nlabel='A'\nfile='a.json'").unwrap();
        let t = &cfg.templates[0];
        assert_eq!(t.prompt_mode, PromptMode::Append);
        assert_eq!(t.prompt_join, " ");
        assert_eq!(t.output_megapixels, 1.0);
        assert!(t.loras.is_empty());
    }
}