karnaugh/src/themes.rs

use std::{
    collections::{BTreeSet, HashMap, HashSet},
    fs,
    path::Path,
    sync::OnceLock,
};

use toml::{Table, Value};
use tracing::{info, warn};
use typst::foundations::{Func, NativeFunc, func};

use crate::AppConfig;

static REGISTRY: OnceLock<ThemeRegistry> = OnceLock::new();

#[derive(Debug)]
struct ThemeRegistry {
    themes: HashMap<String, Theme>,
}

impl ThemeRegistry {
    fn new(theme_root: &Path) -> Self {
        let mut helix_themes = HashMap::new();
        for entry in fs::read_dir(theme_root).into_iter().flatten().flatten() {
            let name = entry
                .path()
                .file_prefix()
                .unwrap()
                .to_string_lossy()
                .into_owned();
            if let Ok(text) = fs::read_to_string(entry.path())
                && let Ok(toml) = toml::from_str::<Table>(&text)
            {
                helix_themes.insert(name, toml);
            }
        }

        // resolve inheritance
        while helix_themes
            .values()
            .any(|table| table.contains_key("inherits"))
        {
            let old_themes = helix_themes.clone();
            old_themes
                .keys()
                .filter(|name| old_themes.get(*name).unwrap().contains_key("inherits"))
                .for_each(|name| {
                    let parent_name = old_themes
                        .get(name)
                        .unwrap()
                        .get("inherits")
                        .unwrap()
                        .as_str()
                        .unwrap();
                    let parent = old_themes.get(parent_name).unwrap();
                    if !parent.contains_key("inherits") {
                        let table = helix_themes.get_mut(name).unwrap();
                        parent.into_iter().for_each(|(k, v)| {
                            if k == "palette" {
                                let parent_palette = v.as_table();
                                let child_palette = table
                                    .entry(k.clone())
                                    .or_insert_with(|| Value::Table(Table::new()))
                                    .as_table_mut();
                                if let (Some(parent_palette), Some(child_palette)) =
                                    (parent_palette, child_palette)
                                {
                                    for (pk, pv) in parent_palette {
                                        child_palette
                                            .entry(pk.clone())
                                            .or_insert_with(|| pv.clone());
                                    }
                                }
                            } else if !table.contains_key(k) {
                                table.insert(k.clone(), v.clone());
                            }
                        });
                        table.remove("inherits");
                    }
                });
        }

        let themes = helix_themes
            .into_iter()
            .flat_map(|(name, table)| match Theme::from_toml(table) {
                Some(theme) => Some((name, theme)),
                None => {
                    warn!("Theme {name} is missing an `ui.background`, skipping it");
                    None
                }
            })
            .collect();

        Self { themes }
    }
}

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
struct Color(String);

impl Color {
    fn from_str(color: &str) -> Option<Self> {
        if color.starts_with('#') {
            Some(Self(color.into()))
        } else {
            None
        }
    }

    fn from_palette(color: &str, palette: &HashMap<String, Color>) -> Option<Self> {
        if color.starts_with('#') {
            Self::from_str(color)
        } else if let Some(color) = palette.get(color) {
            Some(color.clone())
        } else {
            // Neither a hex literal nor a palette entry; assume a named CSS
            // colour and pass it through verbatim.
            warn!("Unresolved colour {color}, passing it through as a literal");
            Some(Self(color.into()))
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
struct Style {
    fg: Option<Color>,
    bold: bool,
    italic: bool,
    underlined: bool,
}

impl Style {
    fn from_toml(style: Value, palette: &HashMap<String, Color>) -> Option<Self> {
        let resolve = |s: &str| Color::from_palette(s, palette);
        match style {
            Value::String(s) => Some(Self {
                fg: resolve(&s),
                bold: false,
                italic: false,
                underlined: false,
            }),
            Value::Table(table) => {
                let fg = table.get("fg").and_then(Value::as_str).and_then(resolve);
                let modifiers: HashSet<&str> = table
                    .get("modifiers")
                    .and_then(Value::as_array)
                    .into_iter()
                    .flatten()
                    .flat_map(Value::as_str)
                    .collect();
                Some(Self {
                    fg,
                    bold: modifiers.contains("bold"),
                    italic: modifiers.contains("italic"),
                    underlined: modifiers.contains("underlined"),
                })
            }
            _ => None,
        }
    }
}

#[derive(Debug)]
struct Theme {
    foreground: Option<Color>,
    background: Option<Color>,
    styles: HashMap<String, Style>,
}

static IGNORED_SCOPES: &[&str] = &["warning", "error", "info", "hint"];

impl Theme {
    fn from_toml(theme: Table) -> Option<Self> {
        let palette = theme
            .get("palette")
            .and_then(Value::as_table)
            .cloned()
            .map(|table| {
                table
                    .into_iter()
                    .flat_map(|(k, v)| v.as_str().and_then(Color::from_str).map(|v| (k, v)))
                    .collect::<HashMap<String, Color>>()
            })
            .unwrap_or_default();

        let resolve = |c: &str| Color::from_palette(c, &palette);

        let ui_background = theme.get("ui.background").and_then(Value::as_table)?;
        let foreground = ui_background
            .get("fg")
            .and_then(Value::as_str)
            .and_then(resolve);
        let background = ui_background
            .get("bg")
            .and_then(Value::as_str)
            .and_then(resolve);

        let styles = theme
            .into_iter()
            .filter(|(k, _)| {
                !k.starts_with("ui.")
                    && !k.starts_with("diagnostic.")
                    && !IGNORED_SCOPES.contains(&k.as_str())
            })
            .flat_map(|(k, v)| Style::from_toml(v, &palette).map(|v| (k, v)))
            .collect();

        Some(Self {
            foreground,
            background,
            styles,
        })
    }
}

pub fn init(config: &AppConfig) {
    match &config.themes_root {
        Some(themes_root) => {
            info!("Loading themes from {}", themes_root.display());
            REGISTRY.get_or_init(|| ThemeRegistry::new(themes_root));
        }
        None => {
            warn!("No themes root has been configured, theme-based colouring will not be possible.")
        }
    }
}

/// The intersection of scopes between two themes.
pub fn shared_scopes(light: &str, dark: &str) -> Option<BTreeSet<String>> {
    let registry = REGISTRY.get()?;
    let light = registry.themes.get(light)?;
    let dark = registry.themes.get(dark)?;
    Some(
        light
            .styles
            .keys()
            .filter(|k| dark.styles.contains_key(*k))
            .cloned()
            .collect(),
    )
}

pub fn theme_css_func() -> Func {
    theme_css::func()
}

fn scope_to_selector(scope: &str) -> String {
    format!(".hl-{}", scope.replace('.', "-"))
}

fn style_declarations(style: &Style) -> String {
    let mut parts = Vec::new();
    if let Some(Color(hex)) = &style.fg {
        parts.push(format!("color:{hex}"));
    }
    parts.push(format!(
        "font-weight:{}",
        if style.bold { "bold" } else { "normal" }
    ));
    parts.push(format!(
        "font-style:{}",
        if style.italic { "italic" } else { "normal" }
    ));
    parts.push(format!(
        "text-decoration:{}",
        if style.underlined {
            "underline"
        } else {
            "none"
        }
    ));
    parts.join(";")
}

fn generate_css(scopes: &[&str], theme: &Theme) -> String {
    let mut rules = Vec::new();

    // The theme's editor foreground/background colour the code block itself.
    let mut block_style = Vec::new();
    if let Some(Color(hex)) = &theme.foreground {
        block_style.push(format!("color:{hex}"));
    }
    if let Some(Color(hex)) = &theme.background {
        block_style.push(format!("background:{hex}"));
    }
    if !block_style.is_empty() {
        rules.push(format!(".code-block{{{}}}", block_style.join(";")));
    }

    let mut groups: HashMap<&Style, Vec<&str>> = HashMap::new();
    for &scope in scopes {
        if let Some(style) = theme.styles.get(scope) {
            groups.entry(style).or_default().push(scope);
        }
    }

    let mut group_rules: Vec<String> = groups
        .iter()
        .map(|(style, scopes)| {
            let selectors = scopes
                .iter()
                .map(|s| scope_to_selector(s))
                .collect::<Vec<_>>()
                .join(",");
            format!("{}{{{}}}", selectors, style_declarations(style))
        })
        .collect();
    // Sort so the emitted stylesheet is deterministic across runs.
    group_rules.sort();
    rules.extend(group_rules);

    rules.join("\n")
}

#[func]
fn theme_css(light: String, dark: String) -> String {
    let Some(registry) = REGISTRY.get() else {
        return String::new();
    };
    let Some(light_theme) = registry.themes.get(&light) else {
        warn!("Unknown light theme: {light}");
        return String::new();
    };
    let Some(dark_theme) = registry.themes.get(&dark) else {
        warn!("Unknown dark theme: {dark}");
        return String::new();
    };

    let shared = shared_scopes(&light, &dark).unwrap_or_default();
    let shared: Vec<&str> = shared.iter().map(String::as_str).collect();

    let light_css = generate_css(&shared, light_theme);
    let dark_css = generate_css(&shared, dark_theme);

    format!("{light_css}\n@media (prefers-color-scheme:dark){{{dark_css}}}")
}