//! Operator-controlled branding overrides.
//!
//! v0.5.2 splits the single brand mark into **three independent slots**,
//! FleetDM-style:
//!
//! * `light` — shown in the WebUI top-left and on the form-login page
//!   when the active theme is light.
//! * `dark` — same surfaces, when the active theme is dark.
//! * `client` — the raster painted above the iPXE boot menu entries
//!   (`/branding/pxe-logo`), i.e. what a PXE client sees on the screen.
//!
//! Each slot lives at `<work_dir>/branding/logo-<slot>.<ext>` and is
//! served in preference to the bundled rainbow-horizon mark when present.
//! Borrowed-from-FleetDM: tenant chrome, same product.
//!
//! Legacy continuity: a pre-v0.5.2 single `logo.<ext>` (recorded under
//! the old `logo_filename`/`logo_mime` keys) is migrated on first load
//! into both the `dark` and `client` slots — that preserves the previous
//! behaviour (one mark fed both the dark WebUI and the PXE screen) until
//! the operator uploads dedicated variants.
//!
//! Storage policy mirrors `HostBindings` / `BootLog`: in-memory cache is
//! authoritative for the current process, disk is the source of truth on
//! restart, and a corrupt cache file falls back to the bundled default
//! rather than blocking startup.

use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use std::path::{Path, PathBuf};
use std::sync::Arc;

/// Allowed MIME types for an uploaded logo. We deliberately keep this
/// narrow — anything that can be `<img src="...">`'d into the brand
/// block, no scripts. SVG carries the obvious XSS risk for raw inline
/// HTML; we always serve the bytes as a separate asset with a strict
/// content-type rather than inlining, so SVG is safe.
pub const ALLOWED_LOGO_MIMES: &[&str] = &[
    "image/svg+xml",
    "image/png",
    "image/jpeg",
    "image/webp",
    "image/gif",
];

/// Disk cap for an uploaded logo. PXE WebUIs are operator-facing — even
/// a generous 2 MB cap is comfortable for any reasonable brand mark and
/// puts a clear bound on memory + serialization cost.
pub const MAX_LOGO_BYTES: usize = 2 * 1024 * 1024;

/// Which branded surface a logo upload targets.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum LogoSlot {
    /// WebUI + form-login page, light theme.
    Light,
    /// WebUI + form-login page, dark theme.
    Dark,
    /// iPXE boot-menu background seen by PXE clients.
    Client,
}

impl LogoSlot {
    #[must_use]
    pub fn as_str(self) -> &'static str {
        match self {
            LogoSlot::Light => "light",
            LogoSlot::Dark => "dark",
            LogoSlot::Client => "client",
        }
    }

    /// Parse a slot name from the URL path segment. Case-insensitive.
    #[must_use]
    pub fn parse(s: &str) -> Option<Self> {
        match s.trim().to_ascii_lowercase().as_str() {
            "light" => Some(LogoSlot::Light),
            "dark" => Some(LogoSlot::Dark),
            "client" => Some(LogoSlot::Client),
            _ => None,
        }
    }
}

/// One brand-mark slot: a filename (relative to the branding dir) plus
/// the MIME we cached at upload time so the HTTP layer can set the
/// Content-Type without re-sniffing.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
struct Slot {
    #[serde(default, skip_serializing_if = "Option::is_none")]
    filename: Option<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    mime: Option<String>,
}

impl Slot {
    fn clear_file(&mut self, dir: &Path) {
        if let Some(name) = self.filename.take() {
            let _ = std::fs::remove_file(dir.join(name));
        }
        self.mime = None;
    }
}

#[derive(Debug, Clone, Default, Serialize, Deserialize)]
struct Inner {
    #[serde(default)]
    light: Slot,
    #[serde(default)]
    dark: Slot,
    #[serde(default)]
    client: Slot,
    /// Monotonic counter bumped on every set/clear (any slot). Surfaces
    /// as a cache-bust token (`/assets/logo.svg?r=<rev>`) so the browser
    /// fetches the new bytes the moment the operator swaps a logo — the
    /// app version alone can't do this since it doesn't change on upload.
    /// Persisted so the token stays stable across restarts and keeps
    /// climbing across multiple swaps.
    #[serde(default)]
    rev: u64,
    // ── Legacy (pre-v0.5.2) single-logo keys ──────────────────────────
    // Read on load for one-way migration into `dark` + `client`, then
    // dropped from the persisted form (skip_serializing_if).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    logo_filename: Option<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    logo_mime: Option<String>,
}

impl Inner {
    fn slot(&self, slot: LogoSlot) -> &Slot {
        match slot {
            LogoSlot::Light => &self.light,
            LogoSlot::Dark => &self.dark,
            LogoSlot::Client => &self.client,
        }
    }

    fn slot_mut(&mut self, slot: LogoSlot) -> &mut Slot {
        match slot {
            LogoSlot::Light => &mut self.light,
            LogoSlot::Dark => &mut self.dark,
            LogoSlot::Client => &mut self.client,
        }
    }
}

/// In-memory + on-disk override registry. Cheap to clone; locks are
/// brief. The `branding.json` cache lives alongside the active assets
/// inside `<work_dir>/branding/`.
#[derive(Debug, Clone)]
pub struct BrandingStore {
    /// Root directory: `<work_dir>/branding/`. Created on first write.
    dir: Arc<PathBuf>,
    inner: Arc<RwLock<Inner>>,
}

impl BrandingStore {
    /// Load (or initialise empty) from `<work_dir>/branding/`. Tolerates
    /// missing directories, partial state, and corrupt JSON — a bad
    /// cache should never block PXE for the network. Migrates a legacy
    /// single-logo file into the dark + client slots.
    #[must_use]
    pub fn load_or_default(work_dir: &Path) -> Self {
        let dir = work_dir.join("branding");
        let path = dir.join("branding.json");
        let mut inner = Inner::default();
        if let Ok(text) = std::fs::read_to_string(&path) {
            match serde_json::from_str::<Inner>(&text) {
                Ok(parsed) => inner = parsed,
                Err(e) => {
                    tracing::warn!(
                        target: "openpxe::branding",
                        "branding.json present but unreadable ({e}); starting empty"
                    );
                }
            }
        }
        let store = Self {
            dir: Arc::new(dir),
            inner: Arc::new(RwLock::new(inner)),
        };
        store.migrate_legacy();
        store.prune_missing();
        store
    }

    /// One-way migration: a pre-v0.5.2 `logo.<ext>` becomes the dark +
    /// client slots (the old single mark fed both the dark WebUI and the
    /// PXE screen). Best-effort; failures leave the legacy file in place
    /// rather than blocking startup.
    fn migrate_legacy(&self) {
        let (legacy_name, legacy_mime) = {
            let g = self.inner.read();
            (g.logo_filename.clone(), g.logo_mime.clone())
        };
        let Some(name) = legacy_name else { return };
        let src = self.dir.join(&name);
        if !src.is_file() {
            // Legacy pointer is stale — just drop it.
            let mut g = self.inner.write();
            g.logo_filename = None;
            g.logo_mime = None;
            drop(g);
            self.persist();
            return;
        }
        let mime = legacy_mime.unwrap_or_else(|| "image/svg+xml".to_string());
        let ext = ext_for_mime(&mime).unwrap_or("bin");
        if let Ok(bytes) = std::fs::read(&src) {
            // Seed dark + client only when those slots are still empty so
            // a re-run (or a manual edit) never clobbers operator intent.
            let needs_dark = self.inner.read().dark.filename.is_none();
            let needs_client = self.inner.read().client.filename.is_none();
            if needs_dark {
                let _ = self.write_slot(LogoSlot::Dark, &mime, ext, &bytes);
            }
            if needs_client {
                let _ = self.write_slot(LogoSlot::Client, &mime, ext, &bytes);
            }
        }
        let _ = std::fs::remove_file(&src);
        {
            let mut g = self.inner.write();
            g.logo_filename = None;
            g.logo_mime = None;
        }
        self.persist();
        tracing::info!(
            target: "openpxe::branding",
            "migrated legacy single logo into dark + client slots"
        );
    }

    /// Drop in-memory slot pointers whose backing file vanished from disk
    /// so the HTTP layer falls back to the bundled mark instead of 500ing.
    fn prune_missing(&self) {
        let mut changed = false;
        {
            let mut g = self.inner.write();
            for slot in [LogoSlot::Light, LogoSlot::Dark, LogoSlot::Client] {
                let present = g
                    .slot(slot)
                    .filename
                    .as_deref()
                    .is_some_and(|n| self.dir.join(n).is_file());
                if !present && g.slot(slot).filename.is_some() {
                    g.slot_mut(slot).filename = None;
                    g.slot_mut(slot).mime = None;
                    changed = true;
                }
            }
        }
        if changed {
            self.persist();
        }
    }

    /// Absolute path to the logo for `slot`, if set and present on disk.
    #[must_use]
    pub fn slot_path(&self, slot: LogoSlot) -> Option<PathBuf> {
        let g = self.inner.read();
        g.slot(slot).filename.as_deref().map(|n| self.dir.join(n))
    }

    /// MIME of the logo for `slot`, if any.
    #[must_use]
    pub fn slot_mime(&self, slot: LogoSlot) -> Option<String> {
        self.inner.read().slot(slot).mime.clone()
    }

    /// Resolve the WebUI logo for a theme, with fallback: light falls
    /// back to dark and vice-versa, so a single uploaded variant still
    /// shows on both themes. Returns `(path, mime)` or `None` (→ bundled).
    #[must_use]
    pub fn web_logo(&self, theme_is_light: bool) -> Option<(PathBuf, String)> {
        let (primary, secondary) = if theme_is_light {
            (LogoSlot::Light, LogoSlot::Dark)
        } else {
            (LogoSlot::Dark, LogoSlot::Light)
        };
        let g = self.inner.read();
        let chosen = if g.slot(primary).filename.is_some() {
            primary
        } else {
            secondary
        };
        let s = g.slot(chosen);
        s.filename.as_deref().map(|n| {
            (
                self.dir.join(n),
                s.mime
                    .clone()
                    .unwrap_or_else(|| "image/svg+xml".to_string()),
            )
        })
    }

    /// Resolve the PXE client logo (no theme fallback — the PXE screen
    /// has a single mark). Returns `(path, mime)` or `None` (→ default
    /// composed background).
    #[must_use]
    pub fn client_logo(&self) -> Option<(PathBuf, String)> {
        let g = self.inner.read();
        let s = &g.client;
        s.filename.as_deref().map(|n| {
            (
                self.dir.join(n),
                s.mime
                    .clone()
                    .unwrap_or_else(|| "application/octet-stream".to_string()),
            )
        })
    }

    /// Replace the logo for `slot`. Returns the chosen on-disk filename so
    /// the caller can echo it back in the API response.
    pub fn set_logo(
        &self,
        slot: LogoSlot,
        mime: &str,
        ext: &str,
        bytes: &[u8],
    ) -> std::io::Result<String> {
        let filename = self.write_slot(slot, mime, ext, bytes)?;
        self.persist();
        tracing::info!(
            target: "openpxe::branding",
            slot = slot.as_str(), file = %filename, mime = %mime, size = bytes.len(),
            "custom logo installed"
        );
        Ok(filename)
    }

    /// Write the bytes for a slot and update the in-memory pointer + rev,
    /// without persisting (the caller decides when to flush). Cleans up
    /// any sibling `logo-<slot>.*` so there's exactly one file per slot.
    fn write_slot(
        &self,
        slot: LogoSlot,
        mime: &str,
        ext: &str,
        bytes: &[u8],
    ) -> std::io::Result<String> {
        std::fs::create_dir_all(self.dir.as_path())?;
        let safe_ext = sanitize_ext(ext);
        let stem = format!("logo-{}", slot.as_str());
        let filename = format!("{stem}.{safe_ext}");
        let final_path = self.dir.join(&filename);
        // Atomic write: tmp -> rename.
        let tmp = final_path.with_extension(format!("{safe_ext}.tmp"));
        std::fs::write(&tmp, bytes)?;
        std::fs::rename(&tmp, &final_path)?;
        // Clean up any sibling `logo-<slot>.<otherext>`.
        if let Ok(entries) = std::fs::read_dir(self.dir.as_path()) {
            for e in entries.flatten() {
                let p = e.path();
                let name = p.file_name().and_then(|s| s.to_str()).unwrap_or("");
                if name.starts_with(&format!("{stem}.")) && name != filename {
                    let _ = std::fs::remove_file(&p);
                }
            }
        }
        let mut g = self.inner.write();
        let s = g.slot_mut(slot);
        s.filename = Some(filename.clone());
        s.mime = Some(mime.to_string());
        g.rev = g.rev.wrapping_add(1);
        Ok(filename)
    }

    /// Drop the override for `slot` and return to the bundled / default.
    pub fn clear_logo(&self, slot: LogoSlot) -> std::io::Result<()> {
        {
            let mut g = self.inner.write();
            let dir = self.dir.as_path();
            g.slot_mut(slot).clear_file(dir);
            g.rev = g.rev.wrapping_add(1);
        }
        self.persist();
        tracing::info!(target: "openpxe::branding", slot = slot.as_str(), "custom logo cleared");
        Ok(())
    }

    /// True if a custom logo is configured for `slot`.
    #[must_use]
    pub fn has_logo(&self, slot: LogoSlot) -> bool {
        self.inner.read().slot(slot).filename.is_some()
    }

    /// True if either WebUI theme slot has a custom logo — drives the
    /// FleetDM-style full-width brand block (and the `has-custom-logo`
    /// class) on the sidebar + login page.
    #[must_use]
    pub fn has_any_web_logo(&self) -> bool {
        let g = self.inner.read();
        g.light.filename.is_some() || g.dark.filename.is_some()
    }

    /// Presence triple `(light, dark, client)` for the `/api/me` and
    /// `/api/status` bootstrap payloads.
    #[must_use]
    pub fn presence(&self) -> (bool, bool, bool) {
        let g = self.inner.read();
        (
            g.light.filename.is_some(),
            g.dark.filename.is_some(),
            g.client.filename.is_some(),
        )
    }

    /// Cache-bust token for the logo asset URLs. Changes on every
    /// set/clear so `/assets/logo.svg?r=<rev>` resolves to a fresh URL
    /// whenever the operator swaps a brand mark. Stable otherwise.
    #[must_use]
    pub fn logo_rev(&self) -> u64 {
        self.inner.read().rev
    }

    fn persist(&self) {
        let snap = self.inner.read().clone();
        let body = match serde_json::to_vec_pretty(&snap) {
            Ok(b) => b,
            Err(e) => {
                tracing::warn!(target: "openpxe::branding", "serialize branding.json: {e}");
                return;
            }
        };
        if let Err(e) = std::fs::create_dir_all(self.dir.as_path()) {
            tracing::warn!(target: "openpxe::branding", "mkdir branding/: {e}");
            return;
        }
        let path = self.dir.join("branding.json");
        let tmp = path.with_extension("json.tmp");
        if let Err(e) = std::fs::write(&tmp, body) {
            tracing::warn!(target: "openpxe::branding", "write branding.json tmp: {e}");
            return;
        }
        if let Err(e) = std::fs::rename(&tmp, &path) {
            tracing::warn!(target: "openpxe::branding", "rename branding.json: {e}");
        }
    }
}

/// Trim arbitrary operator-supplied extension strings to a small, safe
/// alphanumeric form. Anything weird collapses to `bin`. We never let
/// the extension affect the path beyond the final segment of `logo.<x>`.
fn sanitize_ext(ext: &str) -> String {
    let lc: String = ext
        .chars()
        .filter(char::is_ascii_alphanumeric)
        .map(|c| c.to_ascii_lowercase())
        .collect();
    if lc.is_empty() || lc.len() > 5 {
        "bin".into()
    } else {
        lc
    }
}

/// Pick a safe filesystem extension from a MIME type. Returns `None`
/// if the MIME isn't on the [`ALLOWED_LOGO_MIMES`] allowlist.
#[must_use]
pub fn ext_for_mime(mime: &str) -> Option<&'static str> {
    match mime {
        "image/svg+xml" => Some("svg"),
        "image/png" => Some("png"),
        "image/jpeg" => Some("jpg"),
        "image/webp" => Some("webp"),
        "image/gif" => Some("gif"),
        _ => None,
    }
}

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

    #[test]
    fn empty_after_load_when_no_branding_dir() {
        let dir = tempdir().unwrap();
        let b = BrandingStore::load_or_default(dir.path());
        assert!(!b.has_logo(LogoSlot::Light));
        assert!(!b.has_logo(LogoSlot::Dark));
        assert!(!b.has_logo(LogoSlot::Client));
        assert!(b.web_logo(false).is_none());
        assert!(b.client_logo().is_none());
        assert!(!b.has_any_web_logo());
    }

    #[test]
    fn set_clear_round_trip_persists() {
        let dir = tempdir().unwrap();
        let b = BrandingStore::load_or_default(dir.path());
        let name = b
            .set_logo(LogoSlot::Dark, "image/png", "png", b"\x89PNG\r\n\x1a\nfake")
            .unwrap();
        assert_eq!(name, "logo-dark.png");
        assert!(b.has_logo(LogoSlot::Dark));
        assert_eq!(b.slot_mime(LogoSlot::Dark).as_deref(), Some("image/png"));
        let (p, _) = b.web_logo(false).unwrap();
        assert!(p.is_file());

        // Re-open and confirm the override survives a restart.
        drop(b);
        let b2 = BrandingStore::load_or_default(dir.path());
        assert!(b2.has_logo(LogoSlot::Dark));
        assert_eq!(b2.slot_mime(LogoSlot::Dark).as_deref(), Some("image/png"));

        // Clear; the file goes away and has_logo flips off.
        b2.clear_logo(LogoSlot::Dark).unwrap();
        assert!(!b2.has_logo(LogoSlot::Dark));
        assert!(!p.exists());
    }

    #[test]
    fn web_logo_falls_back_across_themes() {
        let dir = tempdir().unwrap();
        let b = BrandingStore::load_or_default(dir.path());
        // Only dark uploaded — light theme falls back to it.
        b.set_logo(LogoSlot::Dark, "image/png", "png", b"dark")
            .unwrap();
        let (p_light, _) = b.web_logo(true).expect("light falls back to dark");
        assert!(p_light.ends_with("logo-dark.png"));
        // Upload a distinct light — now light theme uses its own.
        b.set_logo(LogoSlot::Light, "image/png", "png", b"light")
            .unwrap();
        let (p_light2, _) = b.web_logo(true).unwrap();
        assert!(p_light2.ends_with("logo-light.png"));
        // Client is independent and still unset.
        assert!(b.client_logo().is_none());
    }

    #[test]
    fn replacing_slot_removes_old_extension_sibling() {
        let dir = tempdir().unwrap();
        let b = BrandingStore::load_or_default(dir.path());
        b.set_logo(
            LogoSlot::Client,
            "image/png",
            "png",
            b"\x89PNG\r\n\x1a\nfake",
        )
        .unwrap();
        b.set_logo(
            LogoSlot::Client,
            "image/svg+xml",
            "svg",
            br#"<svg xmlns="http://www.w3.org/2000/svg"/>"#,
        )
        .unwrap();
        let entries: Vec<_> = std::fs::read_dir(dir.path().join("branding"))
            .unwrap()
            .filter_map(|e| e.ok().map(|e| e.file_name().to_string_lossy().into_owned()))
            .collect();
        assert!(
            entries.iter().any(|n| n == "logo-client.svg"),
            "got {entries:?}"
        );
        assert!(
            !entries.iter().any(|n| n == "logo-client.png"),
            "stale PNG left over: {entries:?}"
        );
    }

    #[test]
    fn logo_rev_bumps_on_each_set_and_clear() {
        let dir = tempdir().unwrap();
        let b = BrandingStore::load_or_default(dir.path());
        assert_eq!(b.logo_rev(), 0);
        b.set_logo(LogoSlot::Light, "image/png", "png", b"a")
            .unwrap();
        assert_eq!(b.logo_rev(), 1);
        b.set_logo(LogoSlot::Dark, "image/png", "png", b"b")
            .unwrap();
        assert_eq!(b.logo_rev(), 2);
        b.clear_logo(LogoSlot::Light).unwrap();
        assert_eq!(b.logo_rev(), 3);
        drop(b);
        let b2 = BrandingStore::load_or_default(dir.path());
        assert_eq!(b2.logo_rev(), 3);
    }

    #[test]
    fn legacy_single_logo_migrates_to_dark_and_client() {
        // A pre-v0.5.2 branding.json + logo.png migrates on load.
        let dir = tempdir().unwrap();
        let brand_dir = dir.path().join("branding");
        std::fs::create_dir_all(&brand_dir).unwrap();
        std::fs::write(brand_dir.join("logo.png"), b"\x89PNG\r\n\x1a\nlegacy").unwrap();
        // Hand-write the old shape (logo_filename/logo_mime, no slots).
        std::fs::write(
            brand_dir.join("branding.json"),
            br#"{"logo_filename":"logo.png","logo_mime":"image/png","rev":4}"#,
        )
        .unwrap();
        let b = BrandingStore::load_or_default(dir.path());
        assert!(b.has_logo(LogoSlot::Dark), "dark seeded from legacy");
        assert!(b.has_logo(LogoSlot::Client), "client seeded from legacy");
        assert!(!b.has_logo(LogoSlot::Light), "light stays empty");
        // The old logo.png is gone; per-slot files exist.
        assert!(!brand_dir.join("logo.png").exists());
        assert!(brand_dir.join("logo-dark.png").is_file());
        assert!(brand_dir.join("logo-client.png").is_file());
        // rev carried over from the legacy file and advanced as the two
        // slots were seeded (each write bumps it), so it never regresses.
        let migrated_rev = b.logo_rev();
        assert!(
            migrated_rev >= 4,
            "rev should not regress below legacy: {migrated_rev}"
        );
        // And the migration is sticky across a restart (no re-migrate, no
        // further rev churn).
        drop(b);
        let b2 = BrandingStore::load_or_default(dir.path());
        assert!(b2.has_logo(LogoSlot::Dark));
        assert!(b2.has_logo(LogoSlot::Client));
        assert!(!b2.has_logo(LogoSlot::Light));
        assert_eq!(b2.logo_rev(), migrated_rev, "restart must not re-migrate");
    }

    #[test]
    fn sanitize_ext_strips_separators_and_path_chars() {
        assert_eq!(sanitize_ext("svg"), "svg");
        assert_eq!(sanitize_ext("../etc/passwd"), "bin");
        assert_eq!(sanitize_ext("../svg"), "svg");
        assert_eq!(sanitize_ext(""), "bin");
        assert_eq!(sanitize_ext("PNG"), "png");
        assert_eq!(sanitize_ext("svgvvvv"), "bin");
    }

    #[test]
    fn missing_file_referenced_by_json_resolves_to_empty() {
        let dir = tempdir().unwrap();
        let brand_dir = dir.path().join("branding");
        std::fs::create_dir_all(&brand_dir).unwrap();
        // branding.json claims a dark slot whose file doesn't exist.
        std::fs::write(
            brand_dir.join("branding.json"),
            br#"{"dark":{"filename":"logo-dark.png","mime":"image/png"},"rev":1}"#,
        )
        .unwrap();
        let b = BrandingStore::load_or_default(dir.path());
        assert!(
            !b.has_logo(LogoSlot::Dark),
            "should fall back when referenced file is missing"
        );
    }

    #[test]
    fn slot_parse_round_trips() {
        assert_eq!(LogoSlot::parse("light"), Some(LogoSlot::Light));
        assert_eq!(LogoSlot::parse("DARK"), Some(LogoSlot::Dark));
        assert_eq!(LogoSlot::parse(" client "), Some(LogoSlot::Client));
        assert_eq!(LogoSlot::parse("nope"), None);
        assert_eq!(LogoSlot::Light.as_str(), "light");
    }

    #[test]
    fn ext_for_mime_only_accepts_known_types() {
        assert_eq!(ext_for_mime("image/png"), Some("png"));
        assert_eq!(ext_for_mime("image/svg+xml"), Some("svg"));
        assert_eq!(ext_for_mime("application/octet-stream"), None);
        assert_eq!(ext_for_mime("text/html"), None);
    }
}
