discord-proxy 0.1.0

use super::DiscordChannel;
use anyhow::{Context, Result, bail};
use std::{
    collections::HashSet,
    env, fs,
    path::{Path, PathBuf},
    process::Command as StdCommand,
};
use tokio::process::{Child, Command as TokioCommand};

#[derive(Debug, Clone)]
pub struct DiscordInstall {
    source: String,
    channel: DiscordChannel,
    root: PathBuf,
    app_dir: PathBuf,
    update_exe: Option<PathBuf>,
    exe_path: PathBuf,
    exe_name: String,
}

impl DiscordInstall {
    pub fn source(&self) -> &str {
        &self.source
    }

    pub fn channel(&self) -> DiscordChannel {
        self.channel
    }

    pub fn root(&self) -> &Path {
        &self.root
    }

    pub fn app_dir(&self) -> &Path {
        &self.app_dir
    }

    pub fn update_exe(&self) -> Option<&Path> {
        self.update_exe.as_deref()
    }

    pub fn exe_path(&self) -> &Path {
        &self.exe_path
    }

    pub fn exe_name(&self) -> &str {
        &self.exe_name
    }
}

#[derive(Debug, Clone)]
pub struct InstallCandidateReport {
    source: String,
    path: PathBuf,
    normalized_root: Option<PathBuf>,
    install: Option<DiscordInstall>,
    error: Option<String>,
}

impl InstallCandidateReport {
    pub fn source(&self) -> &str {
        &self.source
    }

    pub fn path(&self) -> &Path {
        &self.path
    }

    pub fn normalized_root(&self) -> Option<&Path> {
        self.normalized_root.as_deref()
    }

    pub fn install(&self) -> Option<&DiscordInstall> {
        self.install.as_ref()
    }

    pub fn error(&self) -> Option<&str> {
        self.error.as_deref()
    }
}

#[derive(Debug, Clone)]
pub struct LaunchPlan {
    install: DiscordInstall,
    proxy_url: String,
}

impl LaunchPlan {
    pub fn new(install: DiscordInstall, proxy_url: String) -> Self {
        Self { install, proxy_url }
    }

    pub fn describe(&self) -> String {
        let command = if let Some(update_exe) = self.install.update_exe() {
            format!(
                "{} --processStart {} --a=--proxy-server={}",
                update_exe.display(),
                self.install.exe_name,
                self.proxy_url
            )
        } else {
            format!(
                "{} --proxy-server={}",
                self.install.exe_path.display(),
                self.proxy_url
            )
        };

        format!(
            "source: {}\nworking dir: {}\ncommand: {command}\nenv: HTTP_PROXY={}\nenv: HTTPS_PROXY={}",
            self.install.source,
            self.install.app_dir.display(),
            self.proxy_url,
            self.proxy_url
        )
    }

    pub fn spawn(&self) -> Result<Child> {
        let mut command = if let Some(update_exe) = self.install.update_exe() {
            let mut command = TokioCommand::new(update_exe);
            command
                .current_dir(&self.install.app_dir)
                .arg("--processStart")
                .arg(&self.install.exe_name)
                .arg(format!("--a=--proxy-server={}", self.proxy_url));
            command
        } else {
            let mut command = TokioCommand::new(&self.install.exe_path);
            command
                .current_dir(&self.install.app_dir)
                .arg(format!("--proxy-server={}", self.proxy_url));
            command
        };

        apply_proxy_environment(&mut command, &self.proxy_url);
        command.spawn().with_context(|| self.describe())
    }
}

pub fn discover_install(
    channel: DiscordChannel,
    override_dir: Option<&Path>,
) -> Result<DiscordInstall> {
    inspect_candidates(channel, override_dir)
        .into_iter()
        .find_map(|report| report.install)
        .with_context(|| discovery_error(channel, override_dir))
}

pub fn inspect_candidates(
    channel: DiscordChannel,
    override_dir: Option<&Path>,
) -> Vec<InstallCandidateReport> {
    candidate_roots(channel, override_dir)
        .into_iter()
        .map(|candidate| inspect_candidate(channel, candidate))
        .collect()
}

pub fn is_process_running(exe_name: &str) -> bool {
    let Ok(output) = StdCommand::new("tasklist")
        .args(["/FI", &format!("IMAGENAME eq {exe_name}"), "/NH"])
        .output()
    else {
        return false;
    };

    if !output.status.success() {
        return false;
    }

    String::from_utf8_lossy(&output.stdout)
        .to_ascii_lowercase()
        .contains(&exe_name.to_ascii_lowercase())
}

fn inspect_candidate(channel: DiscordChannel, candidate: CandidateRoot) -> InstallCandidateReport {
    let normalized_root = normalize_candidate_path(&candidate.path);
    match normalized_root {
        Ok(root) => match inspect_root(channel, &root, &candidate.source) {
            Ok(install) => InstallCandidateReport {
                source: candidate.source,
                path: candidate.path,
                normalized_root: Some(root),
                install: Some(install),
                error: None,
            },
            Err(error) => InstallCandidateReport {
                source: candidate.source,
                path: candidate.path,
                normalized_root: Some(root),
                install: None,
                error: Some(error.to_string()),
            },
        },
        Err(error) => InstallCandidateReport {
            source: candidate.source,
            path: candidate.path,
            normalized_root: None,
            install: None,
            error: Some(error.to_string()),
        },
    }
}

fn inspect_root(channel: DiscordChannel, root: &Path, source: &str) -> Result<DiscordInstall> {
    if !root.is_dir() {
        bail!("root does not exist or is not a directory");
    }

    let update_exe = root
        .join("Update.exe")
        .is_file()
        .then(|| root.join("Update.exe"));
    let (app_dir, exe_name) = find_launchable_app(channel, root)?;
    let exe_path = app_dir.join(&exe_name);

    Ok(DiscordInstall {
        source: source.to_string(),
        channel,
        root: root.to_path_buf(),
        app_dir,
        update_exe,
        exe_path,
        exe_name,
    })
}

fn find_launchable_app(channel: DiscordChannel, root: &Path) -> Result<(PathBuf, String)> {
    for app_dir in sorted_app_dirs(root)? {
        if let Some(exe_name) = find_discord_exe(channel, &app_dir) {
            return Ok((app_dir, exe_name));
        }
    }

    if let Some(exe_name) = find_discord_exe(channel, root) {
        return Ok((root.to_path_buf(), exe_name));
    }

    bail!("no launchable Discord executable found");
}

fn sorted_app_dirs(root: &Path) -> Result<Vec<PathBuf>> {
    let mut app_dirs = fs::read_dir(root)
        .with_context(|| format!("failed to read {}", root.display()))?
        .filter_map(|entry| entry.ok())
        .map(|entry| entry.path())
        .filter(|path| path.is_dir())
        .filter_map(|path| {
            let version = path
                .file_name()
                .and_then(|name| name.to_str())
                .and_then(parse_app_version)?;
            Some((version, path))
        })
        .collect::<Vec<_>>();

    app_dirs.sort_by(|(left, _), (right, _)| right.cmp(left));
    Ok(app_dirs.into_iter().map(|(_, path)| path).collect())
}

#[derive(Debug, Clone)]
struct CandidateRoot {
    source: String,
    path: PathBuf,
}

fn candidate_roots(channel: DiscordChannel, override_dir: Option<&Path>) -> Vec<CandidateRoot> {
    let mut candidates = Vec::new();
    let mut seen = HashSet::new();

    if let Some(path) = override_dir {
        push_candidate(&mut candidates, &mut seen, "override", path.to_path_buf());
        return candidates;
    }

    for key_name in uninstall_key_names(channel) {
        for hive in ["HKCU", "HKLM"] {
            let key =
                format!(r"{hive}\Software\Microsoft\Windows\CurrentVersion\Uninstall\{key_name}");
            if let Some(path) = read_registry_value(&key, "InstallLocation") {
                push_candidate(
                    &mut candidates,
                    &mut seen,
                    format!("registry uninstall {key}"),
                    PathBuf::from(path),
                );
            }
        }
    }

    if let Some(local_app_data) = env::var_os("LOCALAPPDATA") {
        push_candidate(
            &mut candidates,
            &mut seen,
            "LOCALAPPDATA",
            PathBuf::from(local_app_data).join(channel.install_dir_name()),
        );
    }

    for protocol_name in protocol_key_names(channel) {
        let key = format!(r"HKCU\SOFTWARE\Classes\{protocol_name}\shell\open\command");
        if let Some(command) = read_registry_default_value(&key)
            && let Some(path) = extract_command_exe_path(&command)
        {
            push_candidate(
                &mut candidates,
                &mut seen,
                format!("registry protocol {key}"),
                PathBuf::from(path),
            );
        }
    }

    candidates
}

fn push_candidate(
    candidates: &mut Vec<CandidateRoot>,
    seen: &mut HashSet<String>,
    source: impl Into<String>,
    path: PathBuf,
) {
    let key = path
        .to_string_lossy()
        .replace('/', "\\")
        .to_ascii_lowercase();
    if seen.insert(key) {
        candidates.push(CandidateRoot {
            source: source.into(),
            path,
        });
    }
}

fn uninstall_key_names(channel: DiscordChannel) -> &'static [&'static str] {
    match channel {
        DiscordChannel::Stable => &["Discord"],
        DiscordChannel::Canary => &["DiscordCanary", "Discord Canary"],
        DiscordChannel::Ptb => &["DiscordPTB", "Discord PTB"],
        DiscordChannel::Development => &["DiscordDevelopment", "Discord Development"],
    }
}

fn protocol_key_names(channel: DiscordChannel) -> &'static [&'static str] {
    match channel {
        DiscordChannel::Stable => &["Discord"],
        DiscordChannel::Canary => &["DiscordCanary"],
        DiscordChannel::Ptb => &["DiscordPTB"],
        DiscordChannel::Development => &["DiscordDevelopment"],
    }
}

fn normalize_candidate_path(path: &Path) -> Result<PathBuf> {
    let file_name = path
        .file_name()
        .and_then(|name| name.to_str())
        .unwrap_or_default();

    if file_name.eq_ignore_ascii_case("Update.exe") {
        return path
            .parent()
            .map(Path::to_path_buf)
            .context("Update.exe candidate has no parent directory");
    }

    if file_name.to_ascii_lowercase().ends_with(".exe") {
        let exe_dir = path
            .parent()
            .context("executable candidate has no parent directory")?;
        if is_app_dir(exe_dir) {
            return exe_dir
                .parent()
                .map(Path::to_path_buf)
                .context("app-* executable candidate has no install root");
        }
        return Ok(exe_dir.to_path_buf());
    }

    if is_app_dir(path) {
        return path
            .parent()
            .map(Path::to_path_buf)
            .context("app-* directory candidate has no install root");
    }

    Ok(path.to_path_buf())
}

fn is_app_dir(path: &Path) -> bool {
    path.file_name()
        .and_then(|name| name.to_str())
        .and_then(parse_app_version)
        .is_some()
}

fn read_registry_value(key: &str, value: &str) -> Option<String> {
    let output = StdCommand::new("reg")
        .args(["query", key, "/v", value])
        .output()
        .ok()?;
    parse_registry_output(output.status.success(), &output.stdout)
}

fn read_registry_default_value(key: &str) -> Option<String> {
    let output = StdCommand::new("reg")
        .args(["query", key, "/ve"])
        .output()
        .ok()?;
    parse_registry_output(output.status.success(), &output.stdout)
}

fn parse_registry_output(success: bool, stdout: &[u8]) -> Option<String> {
    if !success {
        return None;
    }

    let output = String::from_utf8_lossy(stdout);
    output
        .lines()
        .find_map(parse_registry_value_line)
        .filter(|value| !value.is_empty())
}

fn parse_registry_value_line(line: &str) -> Option<String> {
    let trimmed = line.trim();
    let registry_type = ["REG_SZ", "REG_EXPAND_SZ"]
        .into_iter()
        .find(|registry_type| trimmed.contains(registry_type))?;
    let (_, value) = trimmed.split_once(registry_type)?;
    Some(value.trim().to_string())
}

fn extract_command_exe_path(command: &str) -> Option<String> {
    let trimmed = command.trim();
    if let Some(rest) = trimmed.strip_prefix('"') {
        let (path, _) = rest.split_once('"')?;
        return Some(path.to_string());
    }

    trimmed.split_whitespace().next().map(str::to_string)
}

fn parse_app_version(name: &str) -> Option<Vec<u64>> {
    let version = name.strip_prefix("app-")?;
    let parts = version
        .split('.')
        .map(str::parse::<u64>)
        .collect::<Result<Vec<_>, _>>()
        .ok()?;
    (!parts.is_empty()).then_some(parts)
}

fn find_discord_exe(channel: DiscordChannel, app_dir: &Path) -> Option<String> {
    channel
        .exe_candidates()
        .iter()
        .copied()
        .chain([
            "Discord.exe",
            "DiscordCanary.exe",
            "DiscordPTB.exe",
            "DiscordDevelopment.exe",
        ])
        .find(|candidate| app_dir.join(candidate).is_file())
        .map(str::to_string)
}

fn discovery_error(channel: DiscordChannel, override_dir: Option<&Path>) -> String {
    let reports = inspect_candidates(channel, override_dir);
    let details = reports
        .iter()
        .map(|report| {
            let root = report
                .normalized_root()
                .map(|path| path.display().to_string())
                .unwrap_or_else(|| "<none>".to_string());
            let error = report.error().unwrap_or("unknown error");
            format!(
                "- {}: {} -> {} ({error})",
                report.source(),
                report.path().display(),
                root
            )
        })
        .collect::<Vec<_>>()
        .join("\n");

    if details.is_empty() {
        format!("Discord install not found for {channel}; no candidate roots were available")
    } else {
        format!("Discord install not found for {channel}:\n{details}")
    }
}

fn apply_proxy_environment(command: &mut TokioCommand, proxy_url: &str) {
    for key in [
        "HTTP_PROXY",
        "HTTPS_PROXY",
        "http_proxy",
        "https_proxy",
        "ALL_PROXY",
        "all_proxy",
    ] {
        command.env(key, proxy_url);
    }
}

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

    #[test]
    fn parses_app_versions() {
        assert_eq!(parse_app_version("app-1.0.9239"), Some(vec![1, 0, 9239]));
        assert_eq!(parse_app_version("packages"), None);
        assert_eq!(parse_app_version("app-not-a-version"), None);
    }

    #[test]
    fn normalizes_executable_inside_app_dir_to_install_root() {
        let path = PathBuf::from(r"C:\Discord\app-1.0.9239\Discord.exe");

        let root = normalize_candidate_path(&path).unwrap();

        assert_eq!(root, PathBuf::from(r"C:\Discord"));
    }

    #[test]
    fn normalizes_update_exe_to_install_root() {
        let path = PathBuf::from(r"C:\Discord\Update.exe");

        let root = normalize_candidate_path(&path).unwrap();

        assert_eq!(root, PathBuf::from(r"C:\Discord"));
    }

    #[test]
    fn extracts_quoted_protocol_command_path() {
        let command = r#""C:\Discord\app-1.0.9238\Discord.exe" --url -- "%1""#;

        let path = extract_command_exe_path(command).unwrap();

        assert_eq!(path, r"C:\Discord\app-1.0.9238\Discord.exe");
    }

    #[test]
    fn parses_registry_value_line() {
        let line = r#"InstallLocation    REG_SZ    C:\Users\me\AppData\Local\Discord"#;

        let value = parse_registry_value_line(line).unwrap();

        assert_eq!(value, r"C:\Users\me\AppData\Local\Discord");
    }
}