tubular/src/config.rs

use std::{
    collections::{BTreeMap, BTreeSet},
    ffi::OsString,
    fs::{self, OpenOptions},
    io::{self, Write},
    path::{Path, PathBuf},
    sync::atomic::{AtomicU64, Ordering},
};

use toml_edit::{Array, DocumentMut, Item, Table, value};

use crate::keys;

#[derive(Default, Debug)]
pub struct Config {
    pub channels: BTreeMap<String, String>,
    pub categories: BTreeMap<String, Vec<String>>,
    pub pinned_categories: Vec<String>,
    pub keys: keys::Bindings,
}

impl Config {
    pub fn subscriptions(&self) -> Vec<String> {
        let mut channels: Vec<_> = self
            .channels
            .values()
            .chain(self.categories.values().flatten())
            .collect();
        channels.sort_unstable();
        channels.dedup();
        channels.into_iter().cloned().collect()
    }

    pub fn saved_channels(&self) -> Vec<(String, String)> {
        let mut labels = BTreeMap::new();
        // Mappings are sorted by handle, so aliases use the first readable label.
        for (handle, identity) in &self.channels {
            labels.entry(identity).or_insert(handle);
        }
        for identity in self.categories.values().flatten() {
            labels.entry(identity).or_insert(identity);
        }
        let mut channels: Vec<_> = labels
            .into_iter()
            .map(|(identity, label)| (label.clone(), identity.clone()))
            .collect();
        channels.sort_unstable();
        channels
    }
}

pub fn config_path() -> io::Result<PathBuf> {
    resolve_path(
        std::env::var_os("XDG_CONFIG_HOME"),
        std::env::var_os("HOME"),
    )
}

fn resolve_path(xdg: Option<OsString>, home: Option<OsString>) -> io::Result<PathBuf> {
    let root = match xdg.filter(|path| !path.is_empty() && Path::new(path).is_absolute()) {
        Some(path) => PathBuf::from(path),
        None => PathBuf::from(
            home.filter(|path| !path.is_empty())
                .ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, "HOME is unavailable"))?,
        )
        .join(".config"),
    };
    Ok(root.join("tubular/config.toml"))
}

pub fn load(path: &Path) -> io::Result<Config> {
    let (_, document) = read_document(path)?;
    let channels = parse_channels(&document)?;
    let categories = parse_categories(&document, &channels)?;
    let pinned_categories = parse_pins(&document)?;
    let keys = keys::Bindings::from_toml(document.get("keys"))?;
    Ok(Config {
        channels,
        categories,
        pinned_categories,
        keys,
    })
}

pub fn set_channel_tags(
    path: &Path,
    tags: &BTreeMap<String, bool>,
    handle: &str,
    channel: &str,
) -> io::Result<Config> {
    let normalized_tags = normalize_tags(tags)?;
    let handle = normalize_handle(handle)?;
    let channel = normalize_channel(channel)?;
    if !channel.starts_with("https://www.youtube.com/channel/") {
        return Err(invalid(
            "Channel mappings require a canonical UC channel ID",
        ));
    }
    let (original, mut document) = read_document(path)?;
    let mut mappings = parse_channels(&document)?;
    // Validate the whole file before changing it or creating directories.
    parse_categories(&document, &mappings)?;
    let pinned_categories = parse_pins(&document)?;
    let keys = keys::Bindings::from_toml(document.get("keys"))?;
    match mappings.get(&handle) {
        Some(existing) if existing != &channel => {
            return Err(invalid(format!(
                "Handle {handle:?} already maps to another channel"
            )));
        }
        Some(_) => {}
        None => {
            if !document.contains_key("channels") {
                document.insert("channels", Item::Table(Table::new()));
            }
            document["channels"]
                .as_table_like_mut()
                .expect("channel mappings were validated")
                .insert(&handle, value(&channel));
            mappings.insert(handle.clone(), channel.clone());
        }
    }
    if !document.contains_key("categories") {
        document.insert("categories", Item::Table(Table::new()));
    }
    let categories = document["categories"]
        .as_table_like_mut()
        .expect("categories were validated");
    for (name, checked) in normalized_tags {
        // Keep manually written keys and their formatting, including surrounding spaces.
        let key = categories
            .iter()
            .find(|(key, _)| key.trim() == name)
            .map(|(key, _)| key.to_owned())
            .unwrap_or_else(|| name.to_owned());
        if !categories.contains_key(&key) {
            categories.insert(&key, value(Array::new()));
        }
        let channels = categories
            .get_mut(&key)
            .and_then(Item::as_array_mut)
            .expect("category was validated");
        if checked {
            let mut found = false;
            for entry in channels.iter_mut() {
                if resolve_channel(entry.as_str().expect("channel was validated"), &mappings)?
                    == channel
                {
                    // Keep manually written duplicates and entry comments.
                    if entry.as_str() != Some(&handle) {
                        let decor = entry.decor().clone();
                        *entry = handle.as_str().into();
                        *entry.decor_mut() = decor;
                    }
                    found = true;
                }
            }
            if !found {
                channels.push(handle.as_str());
            }
        } else {
            remove_category_channels(channels, &mappings, |identity| identity == channel)?;
        }
    }
    let categories = parse_categories(&document, &mappings)?;
    save(path, original.as_deref(), document.to_string().as_bytes())?;
    Ok(Config {
        channels: mappings,
        categories,
        pinned_categories,
        keys,
    })
}

pub fn set_category_tags(
    path: &Path,
    source: &str,
    expected_channels: &[String],
    tags: &BTreeMap<String, bool>,
) -> io::Result<Config> {
    let source = category_name(source)?;
    let tags = normalize_tags(tags)?;
    let (original, mut document) = read_document(path)?;
    let mappings = parse_channels(&document)?;
    let parsed_categories = parse_categories(&document, &mappings)?;
    let pinned_categories = parse_pins(&document)?;
    let keys = keys::Bindings::from_toml(document.get("keys"))?;
    let targets = parsed_categories
        .get(source)
        .ok_or_else(|| invalid(format!("Unknown source category: {source:?}")))?;
    if targets.is_empty() {
        return Err(invalid("Cannot edit tags for an empty source category"));
    }
    let expected: BTreeSet<_> = expected_channels.iter().map(String::as_str).collect();
    if targets.len() != expected.len()
        || targets
            .iter()
            .any(|channel| !expected.contains(channel.as_str()))
    {
        return Err(invalid(
            "Source category changed since opening the editor; please reopen it",
        ));
    }
    if !tags.values().any(|checked| *checked)
        && targets
            .iter()
            .any(|channel| !mappings.values().any(|mapped| mapped == channel))
    {
        return Err(invalid(
            "Cannot clear every category: a source channel is not registered in [channels]",
        ));
    }

    let categories = document["categories"]
        .as_table_like_mut()
        .expect("source category was validated");
    let source_entries = categories
        .iter()
        .find(|(name, _)| name.trim() == source)
        .and_then(|(_, item)| item.as_array())
        .expect("source category was validated");
    let mut additions = BTreeMap::new();
    for entry in source_entries.iter() {
        let raw = entry.as_str().expect("channel was validated").trim();
        let identity = resolve_channel(raw, &mappings)?;
        let readable = raw
            .strip_prefix("https://www.youtube.com/")
            .or_else(|| raw.strip_prefix("https://youtube.com/"))
            .filter(|address| address.starts_with('@'))
            .unwrap_or(raw);
        let addition = additions
            .entry(identity)
            .or_insert_with(|| readable.to_owned());
        if readable.starts_with('@') && !addition.starts_with('@') {
            addition.clear();
            addition.push_str(readable);
        }
    }
    for (handle, identity) in &mappings {
        if let Some(addition) = additions.get_mut(identity)
            && !addition.starts_with('@')
        {
            addition.clear();
            addition.push_str(handle);
        }
    }
    // The checked set is exact, including categories added by another editor.
    for (name, item) in categories.iter_mut() {
        let entries = item.as_array_mut().expect("category was validated");
        if tags.get(name.trim()).copied().unwrap_or(false) {
            let present: BTreeSet<_> = entries
                .iter()
                .map(|entry| {
                    resolve_channel(entry.as_str().expect("channel was validated"), &mappings)
                })
                .collect::<io::Result<_>>()?;
            for (identity, readable) in &additions {
                if !present.contains(identity) {
                    entries.push(readable.as_str());
                }
            }
        } else {
            remove_category_channels(entries, &mappings, |identity| {
                additions.contains_key(identity)
            })?;
        }
    }
    for (name, checked) in tags {
        if checked && !parsed_categories.contains_key(name) {
            let mut entries = Array::new();
            for readable in additions.values() {
                entries.push(readable.as_str());
            }
            categories.insert(name, value(entries));
        }
    }
    let categories = parse_categories(&document, &mappings)?;
    save(path, original.as_deref(), document.to_string().as_bytes())?;
    Ok(Config {
        channels: mappings,
        categories,
        pinned_categories,
        keys,
    })
}

fn normalize_tags(tags: &BTreeMap<String, bool>) -> io::Result<BTreeMap<&str, bool>> {
    let mut normalized = BTreeMap::new();
    for (name, checked) in tags {
        let name = category_name(name)?;
        if normalized.insert(name, *checked).is_some() {
            return Err(invalid(format!(
                "Duplicate category name after trimming: {name:?}"
            )));
        }
    }
    Ok(normalized)
}

fn remove_category_channels(
    channels: &mut Array,
    mappings: &BTreeMap<String, String>,
    remove: impl Fn(&str) -> bool,
) -> io::Result<()> {
    for index in (0..channels.len()).rev() {
        if !remove(&resolve_channel(
            channels
                .get(index)
                .unwrap()
                .as_str()
                .expect("channel was validated"),
            mappings,
        )?) {
            continue;
        }
        let removed = channels.remove(index);
        // Array comments can be attached to the next entry, not their own line.
        let prefix = removed
            .decor()
            .prefix()
            .and_then(|raw| raw.as_str())
            .unwrap_or("");
        let suffix = removed
            .decor()
            .suffix()
            .and_then(|raw| raw.as_str())
            .unwrap_or("");
        if prefix.contains('#') || suffix.contains('#') {
            if let Some(next) = channels.get_mut(index) {
                let next_prefix = next
                    .decor()
                    .prefix()
                    .and_then(|raw| raw.as_str())
                    .unwrap_or("");
                let preserved = format!("{prefix}{suffix}{next_prefix}");
                next.decor_mut().set_prefix(preserved);
            } else {
                let trailing = channels.trailing().as_str().unwrap_or("");
                let preserved = format!("{prefix}{suffix}{trailing}");
                channels.set_trailing(preserved);
            }
        }
    }
    Ok(())
}

pub fn set_category_pinned(path: &Path, name: &str, pinned: bool) -> io::Result<Config> {
    let name = category_name(name)?;
    let (original, mut document) = read_document(path)?;
    let channels = parse_channels(&document)?;
    let categories = parse_categories(&document, &channels)?;
    let pins = parse_pins(&document)?;
    let keys = keys::Bindings::from_toml(document.get("keys"))?;
    if !categories.contains_key(name) {
        return Err(invalid(format!("Unknown category: {name:?}")));
    }
    if !document.contains_key("settings") {
        document.insert("settings", Item::Table(Table::new()));
    }
    let settings = document["settings"]
        .as_table_like_mut()
        .expect("settings were validated");
    if !settings.contains_key("pinned_categories") {
        settings.insert("pinned_categories", value(Array::new()));
    }
    let array = settings
        .get_mut("pinned_categories")
        .and_then(Item::as_array_mut)
        .expect("pins were validated");
    if let Some(index) = pins.iter().position(|pin| pin == name) {
        if !pinned {
            array.remove(index);
        }
    } else if pinned {
        array.push(name);
    }
    let pinned_categories = parse_pins(&document)?;
    save(path, original.as_deref(), document.to_string().as_bytes())?;
    Ok(Config {
        channels,
        categories,
        pinned_categories,
        keys,
    })
}

pub fn create_category(path: &Path, name: &str, channels: &[String]) -> io::Result<Config> {
    let name = category_name(name)?;
    let (original, mut document) = read_document(path)?;
    let mappings = parse_channels(&document)?;
    let mut categories = parse_categories(&document, &mappings)?;
    let pinned_categories = parse_pins(&document)?;
    let keys = keys::Bindings::from_toml(document.get("keys"))?;
    if categories.contains_key(name) {
        return Err(invalid(format!("Category already exists: {name:?}")));
    }

    let mut labels = BTreeMap::new();
    // As in saved_channels, use the first mapped handle for each identity.
    for (handle, identity) in &mappings {
        labels.entry(identity.as_str()).or_insert(handle.as_str());
    }
    let identities: BTreeSet<_> = channels
        .iter()
        .map(|channel| resolve_channel(channel, &mappings))
        .collect::<io::Result<_>>()?;
    let mut entries = Array::new();
    for identity in &identities {
        let readable = labels.get(identity.as_str()).copied().unwrap_or_else(|| {
            identity
                .strip_prefix("https://www.youtube.com/")
                .filter(|address| address.starts_with('@'))
                .unwrap_or(identity)
        });
        entries.push(readable);
    }
    if !document.contains_key("categories") {
        document.insert("categories", Item::Table(Table::new()));
    }
    document["categories"]
        .as_table_like_mut()
        .expect("categories were validated")
        .insert(name, value(entries));
    categories.insert(name.to_owned(), identities.into_iter().collect());
    save(path, original.as_deref(), document.to_string().as_bytes())?;
    Ok(Config {
        channels: mappings,
        categories,
        pinned_categories,
        keys,
    })
}

pub fn rename_category(path: &Path, old: &str, new: &str) -> io::Result<Config> {
    let old = category_name(old)?;
    let new = category_name(new)?;
    let (original, mut document) = read_document(path)?;
    let channels = parse_channels(&document)?;
    let categories = parse_categories(&document, &channels)?;
    let pinned_categories = parse_pins(&document)?;
    let keys = keys::Bindings::from_toml(document.get("keys"))?;
    if !categories.contains_key(old) {
        return Err(invalid(format!("Unknown source category: {old:?}")));
    }
    if old == new {
        return Ok(Config {
            channels,
            categories,
            pinned_categories,
            keys,
        });
    }
    if categories.contains_key(new) {
        return Err(invalid(format!("Category already exists: {new:?}")));
    }

    // KeyMut only edits formatting. Move entries through formatted insertion to
    // rename the key without changing its position or any membership formatting.
    match &mut document["categories"] {
        Item::Table(table) => {
            let names: Vec<_> = table.iter().map(|(name, _)| name.to_owned()).collect();
            for name in names {
                let (mut key, item) = table.remove_entry(&name).expect("category was validated");
                if key.get().trim() == old {
                    key = toml_edit::Key::new(new)
                        .with_leaf_decor(key.leaf_decor().clone())
                        .with_dotted_decor(key.dotted_decor().clone());
                }
                table.insert_formatted(&key, item);
            }
        }
        Item::Value(toml_edit::Value::InlineTable(table)) => {
            let names: Vec<_> = table.iter().map(|(name, _)| name.to_owned()).collect();
            for name in names {
                let (mut key, item) = table.remove_entry(&name).expect("category was validated");
                if key.get().trim() == old {
                    key = toml_edit::Key::new(new)
                        .with_leaf_decor(key.leaf_decor().clone())
                        .with_dotted_decor(key.dotted_decor().clone());
                }
                table.insert_formatted(&key, item);
            }
        }
        _ => unreachable!("categories were validated"),
    }
    if let Some(array) = document
        .get_mut("settings")
        .and_then(Item::as_table_like_mut)
        .and_then(|settings| settings.get_mut("pinned_categories"))
        .and_then(Item::as_array_mut)
    {
        for pin in array.iter_mut() {
            if pin.as_str().expect("pin was validated").trim() == old {
                let decor = pin.decor().clone();
                *pin = new.into();
                *pin.decor_mut() = decor;
            }
        }
    }
    let categories = parse_categories(&document, &channels)?;
    let pinned_categories = parse_pins(&document)?;
    save(path, original.as_deref(), document.to_string().as_bytes())?;
    Ok(Config {
        channels,
        categories,
        pinned_categories,
        keys,
    })
}

fn parse_pins(document: &DocumentMut) -> io::Result<Vec<String>> {
    if document.contains_key("pinned_categories") {
        return Err(invalid("Move top-level pinned_categories into [settings]"));
    }
    let Some(settings) = document.get("settings") else {
        return Ok(Vec::new());
    };
    let settings = settings
        .as_table_like()
        .ok_or_else(|| invalid("settings must be a TOML table"))?;
    let Some(item) = settings.get("pinned_categories") else {
        return Ok(Vec::new());
    };
    let array = item
        .as_array()
        .ok_or_else(|| invalid("pinned_categories must be an array"))?;
    let mut pins = Vec::with_capacity(array.len());
    for item in array.iter() {
        let name = category_name(
            item.as_str()
                .ok_or_else(|| invalid("Pinned categories must be strings"))?,
        )?;
        if pins.iter().any(|pin| pin == name) {
            return Err(invalid(format!("Duplicate pinned category: {name:?}")));
        }
        pins.push(name.to_owned());
    }
    Ok(pins)
}

fn invalid(message: impl Into<String>) -> io::Error {
    io::Error::new(io::ErrorKind::InvalidData, message.into())
}

fn category_name(name: &str) -> io::Result<&str> {
    if name.chars().any(char::is_control) || name.trim().is_empty() {
        return Err(invalid(
            "Category names must be nonempty and contain no control characters",
        ));
    }
    Ok(name.trim())
}

pub(crate) fn normalize_channel(input: &str) -> io::Result<String> {
    if input.chars().any(char::is_control) {
        return Err(invalid(
            "Channel addresses must contain no control characters",
        ));
    }
    let input = input.trim();
    let url_address = input
        .strip_prefix("https://www.youtube.com/")
        .or_else(|| input.strip_prefix("https://youtube.com/"));
    let address = url_address.unwrap_or(input);
    if let Some(handle) = address.strip_prefix('@') {
        // Validate address safety, not creation rules: existing handles can have edge separators.
        let length = handle.chars().count();
        if length == 0
            || length > 30
            || !handle
                .chars()
                .all(|c| c.is_alphanumeric() || matches!(c, '_' | '-' | '.'))
        {
            return Err(invalid("Invalid YouTube channel handle"));
        }
        return Ok(format!("https://www.youtube.com/@{handle}"));
    }
    let id = match url_address {
        Some(address) => address
            .strip_prefix("channel/")
            .ok_or_else(|| invalid("Expected a YouTube channel URL"))?,
        None => address,
    };
    if id.len() != 24
        || !id.starts_with("UC")
        || !id
            .bytes()
            .all(|c| c.is_ascii_alphanumeric() || c == b'_' || c == b'-')
    {
        return Err(invalid(
            "Expected a YouTube @handle or canonical UC channel ID",
        ));
    }
    Ok(format!("https://www.youtube.com/channel/{id}"))
}

fn normalize_handle(input: &str) -> io::Result<String> {
    let mut address = normalize_channel(input)?;
    let prefix = "https://www.youtube.com/";
    if !address[prefix.len()..].starts_with('@') {
        return Err(invalid("Expected a YouTube @handle"));
    }
    address.drain(..prefix.len());
    Ok(address)
}

fn parse_channels(document: &DocumentMut) -> io::Result<BTreeMap<String, String>> {
    let mut mappings = BTreeMap::new();
    let Some(item) = document.get("channels") else {
        return Ok(mappings);
    };
    let channels = item
        .as_table_like()
        .ok_or_else(|| invalid("channels must be a TOML table"))?;
    for (handle, address) in channels.iter() {
        let handle = normalize_handle(handle)?;
        let address = normalize_channel(
            address
                .as_str()
                .ok_or_else(|| invalid("Channel mappings must be strings"))?,
        )?;
        if !address.starts_with("https://www.youtube.com/channel/") {
            return Err(invalid(
                "Channel mappings require a canonical UC channel ID",
            ));
        }
        if mappings.insert(handle, address).is_some() {
            return Err(invalid("Duplicate channel handle after normalization"));
        }
    }
    Ok(mappings)
}

fn resolve_channel(input: &str, mappings: &BTreeMap<String, String>) -> io::Result<String> {
    let address = normalize_channel(input)?;
    match address
        .strip_prefix("https://www.youtube.com/")
        .and_then(|handle| mappings.get(handle))
    {
        Some(mapped) => Ok(mapped.clone()),
        None => Ok(address),
    }
}

fn read_bytes(path: &Path) -> io::Result<Option<Vec<u8>>> {
    match fs::read(path) {
        Ok(bytes) => Ok(Some(bytes)),
        Err(error) if error.kind() == io::ErrorKind::NotFound => Ok(None),
        Err(error) => Err(error),
    }
}

fn read_document(path: &Path) -> io::Result<(Option<Vec<u8>>, DocumentMut)> {
    let bytes = read_bytes(path)?;
    let document = match &bytes {
        Some(bytes) => std::str::from_utf8(bytes)
            .map_err(|error| invalid(format!("Config is not UTF-8: {error}")))?
            .parse::<DocumentMut>()
            .map_err(|error| invalid(format!("Invalid config TOML: {error}")))?,
        None => DocumentMut::new(),
    };
    Ok((bytes, document))
}

fn parse_categories(
    document: &DocumentMut,
    mappings: &BTreeMap<String, String>,
) -> io::Result<BTreeMap<String, Vec<String>>> {
    let mut parsed = BTreeMap::new();
    let Some(item) = document.get("categories") else {
        return Ok(parsed);
    };
    let categories = item
        .as_table_like()
        .ok_or_else(|| invalid("categories must be a TOML table"))?;
    for (name, channels) in categories.iter() {
        let name = category_name(name)?;
        let channels = channels
            .as_array()
            .ok_or_else(|| invalid(format!("Category {name:?} must be an array of channels")))?;
        let mut normalized = Vec::with_capacity(channels.len());
        for channel in channels.iter() {
            let channel = resolve_channel(
                channel
                    .as_str()
                    .ok_or_else(|| invalid("Channels must be strings"))?,
                mappings,
            )?;
            if !normalized.contains(&channel) {
                normalized.push(channel);
            }
        }
        if parsed.insert(name.to_owned(), normalized).is_some() {
            return Err(invalid(format!(
                "Duplicate category name after trimming: {name:?}"
            )));
        }
    }
    Ok(parsed)
}

static NEXT_TEMP: AtomicU64 = AtomicU64::new(0);

fn save(path: &Path, original: Option<&[u8]>, contents: &[u8]) -> io::Result<()> {
    let parent = path
        .parent()
        .filter(|p| !p.as_os_str().is_empty())
        .unwrap_or(Path::new("."));
    let filename = path
        .file_name()
        .ok_or_else(|| invalid("Config path has no filename"))?;
    fs::create_dir_all(parent)?;
    let (temporary, mut file) = loop {
        let mut name = OsString::from(".");
        name.push(filename);
        name.push(format!(
            ".{}.{}.tmp",
            std::process::id(),
            NEXT_TEMP.fetch_add(1, Ordering::Relaxed)
        ));
        let temporary = parent.join(name);
        let mut options = OpenOptions::new();
        options.write(true).create_new(true);
        #[cfg(unix)]
        {
            use std::os::unix::fs::OpenOptionsExt;
            options.mode(0o600);
        }
        match options.open(&temporary) {
            Ok(file) => break (temporary, file),
            Err(error) if error.kind() == io::ErrorKind::AlreadyExists => continue,
            Err(error) => return Err(error),
        }
    };
    let result = (|| {
        if original.is_some() {
            file.set_permissions(fs::metadata(path)?.permissions())?;
        }
        file.write_all(contents)?;
        file.sync_all()?;
        if read_bytes(path)?.as_deref() != original {
            return Err(io::Error::other(
                "Config changed while saving; please retry",
            ));
        }
        if original.is_none() {
            // Unlike rename, hard_link cannot overwrite a file created by a racing editor.
            fs::hard_link(&temporary, path)?;
        } else {
            // std has no compare-and-swap rename; check immediately before replacement.
            fs::rename(&temporary, path)?;
        }
        Ok(())
    })();
    drop(file);
    let _ = fs::remove_file(&temporary);
    result
}

#[cfg(test)]
#[path = "../tests/unit/config.rs"]
mod tests;