rubot/src/bot/modules/weather.rs

use std::collections::{BTreeMap, HashMap};
use std::fs;
use std::io;
use std::iter::Peekable;
use std::path::{Path, PathBuf};
use std::str::Chars;
use std::time::{Duration, Instant};

use crate::bot::module::{Action, Command, CommandSpec, Module};

const CACHE_TTL: Duration = Duration::from_secs(600);

const COMMANDS: &[CommandSpec] = &[
    CommandSpec {
        name: "add",
        usage: "add <location>",
        about: "save your weather location",
    },
    CommandSpec {
        name: "w",
        usage: "w [location]",
        about: "weather for your saved location",
    },
];

pub struct Weather {
    store: Store,
    path: PathBuf,
    cache: HashMap<String, (Instant, String)>,
}

impl Weather {
    pub fn new(network: &str) -> Weather {
        let dir = std::env::var("RUSTBOT_DATA_DIR").unwrap_or_else(|_| ".".to_string());
        let safe: String = network
            .chars()
            .map(|c| if c.is_ascii_alphanumeric() { c } else { '_' })
            .collect();
        Weather::with_path(PathBuf::from(dir).join(format!("weather.{safe}.json")))
    }

    pub fn with_path(path: PathBuf) -> Weather {
        Weather {
            store: Store::load(&path),
            path,
            cache: HashMap::new(),
        }
    }

    pub fn location_of(&self, nick: &str) -> Option<&str> {
        self.store.get(nick).map(String::as_str)
    }
}

impl Module for Weather {
    fn name(&self) -> &'static str {
        "weather"
    }

    fn commands(&self) -> &'static [CommandSpec] {
        COMMANDS
    }

    fn on_command(&mut self, cmd: &Command) -> Vec<Action> {
        match cmd.name {
            "add" => {
                if cmd.args.is_empty() {
                    return vec![Action::reply(format!(
                        "usage: {}add <location>",
                        cmd.prefix
                    ))];
                }
                let loc = cmd.args.join(" ");
                self.store.set(cmd.sender, &loc);
                if let Err(e) = self.store.save(&self.path) {
                    crate::log(&format!(
                        "weather: could not save {}: {e}",
                        self.path.display()
                    ));
                    return vec![Action::reply(format!(
                        "{}: could not save location",
                        cmd.sender
                    ))];
                }
                vec![Action::reply(format!(
                    "{}: saved {loc} — try {}w",
                    cmd.sender, cmd.prefix
                ))]
            }
            "w" => {
                let loc = if !cmd.args.is_empty() {
                    cmd.args.join(" ")
                } else if let Some(l) = self.store.get(cmd.sender) {
                    l.clone()
                } else {
                    return vec![Action::reply(format!(
                        "{}: no location set — {}add <location> first",
                        cmd.sender, cmd.prefix
                    ))];
                };
                let key = loc.to_lowercase();
                if let Some((at, text)) = self.cache.get(&key) {
                    if at.elapsed() < CACHE_TTL {
                        return vec![Action::reply(text.clone())];
                    }
                }
                match fetch_weather(&loc) {
                    Ok(text) => {
                        self.cache.insert(key, (Instant::now(), text.clone()));
                        vec![Action::reply(text)]
                    }
                    Err(e) => vec![Action::reply(format!("weather unavailable: {e}"))],
                }
            }
            _ => vec![],
        }
    }
}

struct Store {
    map: BTreeMap<String, String>,
}

impl Store {
    fn load(path: &Path) -> Store {
        match fs::read_to_string(path) {
            Ok(s) => Store { map: parse_db(&s) },
            Err(_) => Store {
                map: BTreeMap::new(),
            },
        }
    }

    fn save(&self, path: &Path) -> io::Result<()> {
        let body = to_json(&self.map) + "\n";
        let mut tmp = path.to_path_buf().into_os_string();
        tmp.push(".tmp");
        let tmp = PathBuf::from(tmp);
        fs::write(&tmp, body)?;
        fs::rename(&tmp, path)
    }

    fn get(&self, nick: &str) -> Option<&String> {
        self.map.get(&nick.to_lowercase())
    }

    fn set(&mut self, nick: &str, loc: &str) {
        self.map.insert(nick.to_lowercase(), loc.to_string());
    }
}

fn fetch_weather(location: &str) -> Result<String, String> {
    let path = format!(
        "/{}?format=%l|%c|%C|%t|%f|%w|%h&m",
        crate::http::url_encode(location)
    );
    let body = crate::http::get("wttr.in", &path)?;
    let raw = body
        .lines()
        .map(str::trim)
        .find(|l| !l.is_empty())
        .unwrap_or("");
    if raw.is_empty() {
        return Err("no data".to_string());
    }
    Ok(format_weather(raw, location))
}

pub fn format_weather(raw: &str, location: &str) -> String {
    let f: Vec<&str> = raw.split('|').map(str::trim).collect();
    if f.len() >= 7 {
        let temp = f[3].trim_start_matches('+');
        let feels = f[4].trim_start_matches('+');
        let msg = format!(
            "{} {}: {}, {} (feels like {}) · 💨 {} · 💧 {}",
            f[1], f[0], f[2], temp, feels, f[5], f[6]
        );
        return msg.chars().take(250).collect();
    }
    let low = raw.to_lowercase();
    if low.contains("not found") || low.contains("unknown location") {
        return format!("couldn't find '{location}' — try a city name or postcode");
    }
    raw.chars().take(250).collect()
}

pub fn to_json(map: &BTreeMap<String, String>) -> String {
    let items: Vec<String> = map
        .iter()
        .map(|(k, v)| format!("\"{}\":\"{}\"", json_escape(k), json_escape(v)))
        .collect();
    format!("{{{}}}", items.join(","))
}

pub fn parse_db(s: &str) -> BTreeMap<String, String> {
    let mut map = BTreeMap::new();
    let mut chars = s.chars().peekable();
    skip_ws(&mut chars);
    if chars.next() != Some('{') {
        return map;
    }
    loop {
        skip_ws(&mut chars);
        match chars.peek() {
            Some('"') => {}
            Some(',') => {
                chars.next();
                continue;
            }
            _ => break,
        }
        let Some(key) = parse_string(&mut chars) else {
            break;
        };
        skip_ws(&mut chars);
        if chars.next() != Some(':') {
            break;
        }
        skip_ws(&mut chars);
        let Some(value) = parse_string(&mut chars) else {
            break;
        };
        map.insert(key, value);
    }
    map
}

fn skip_ws(chars: &mut Peekable<Chars>) {
    while let Some(c) = chars.peek() {
        if c.is_whitespace() {
            chars.next();
        } else {
            break;
        }
    }
}

fn parse_string(chars: &mut Peekable<Chars>) -> Option<String> {
    if chars.next() != Some('"') {
        return None;
    }
    let mut out = String::new();
    while let Some(c) = chars.next() {
        match c {
            '"' => return Some(out),
            '\\' => match chars.next()? {
                '"' => out.push('"'),
                '\\' => out.push('\\'),
                '/' => out.push('/'),
                'n' => out.push('\n'),
                'r' => out.push('\r'),
                't' => out.push('\t'),
                'b' => out.push('\u{08}'),
                'f' => out.push('\u{0C}'),
                'u' => {
                    let mut code = 0u32;
                    for _ in 0..4 {
                        code = code * 16 + chars.next()?.to_digit(16)?;
                    }
                    if let Some(ch) = char::from_u32(code) {
                        out.push(ch);
                    }
                }
                other => out.push(other),
            },
            c => out.push(c),
        }
    }
    None
}

fn json_escape(s: &str) -> String {
    let mut out = String::with_capacity(s.len() + 2);
    for c in s.chars() {
        match c {
            '"' => out.push_str("\\\""),
            '\\' => out.push_str("\\\\"),
            '\n' => out.push_str("\\n"),
            '\r' => out.push_str("\\r"),
            '\t' => out.push_str("\\t"),
            c if (c as u32) < 0x20 => out.push_str(&format!("\\u{:04x}", c as u32)),
            c => out.push(c),
        }
    }
    out
}