//! Minimal OpenSSH `known_hosts` support.
//!
//! Behaviour is deliberately conservative:
//! - match entries from the user's `~/.ssh/known_hosts` and meatshell's own
//!   config-local `known_hosts`;
//! - trust-on-first-use into meatshell's file when no entry exists;
//! - reject the connection when a matching host entry exists with a different
//!   key.

use std::fs::{self, OpenOptions};
use std::io::Write;
use std::path::{Path, PathBuf};

use anyhow::{Context, Result};
use directories::UserDirs;
use ssh_key::PublicKey;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum HostKeyStatus {
    Verified,
    Added,
    Changed,
}

/// Check `host:port` against OpenSSH known_hosts files, adding a new entry to
/// `app_known_hosts` on first use.
pub fn check_or_trust_in(
    host: &str,
    port: u16,
    key: &PublicKey,
    app_known_hosts: &Path,
) -> Result<HostKeyStatus> {
    let key = public_key_parts(key)?;
    let targets = host_patterns(host, port);

    let mut saw_same_key_type = false;
    for (path, can_reject) in known_hosts_files(app_known_hosts) {
        let Ok(raw) = fs::read_to_string(&path) else {
            continue;
        };
        for line in raw.lines() {
            let Some(entry) = parse_line(line) else {
                continue;
            };
            if !entry.hosts.iter().any(|h| targets.iter().any(|t| host_matches(h, t))) {
                continue;
            }
            if entry.key_type == key.0 {
                if can_reject {
                    saw_same_key_type = true;
                }
                if entry.key_data == key.1 {
                    return Ok(HostKeyStatus::Verified);
                }
            }
        }
    }

    if saw_same_key_type {
        return Ok(HostKeyStatus::Changed);
    }

    append_app_known_host(app_known_hosts, host, port, &key.0, &key.1)?;
    Ok(HostKeyStatus::Added)
}

fn known_hosts_files(app_known_hosts: &Path) -> Vec<(PathBuf, bool)> {
    let mut out = Vec::new();
    if let Some(home) = UserDirs::new() {
        out.push((home.home_dir().join(".ssh").join("known_hosts"), false));
    }
    out.push((app_known_hosts.to_path_buf(), true));
    out
}

fn append_app_known_host(
    path: &Path,
    host: &str,
    port: u16,
    key_type: &str,
    key_data: &str,
) -> Result<()> {
    if let Some(parent) = path.parent() {
        fs::create_dir_all(parent)
            .with_context(|| format!("failed to create {}", parent.display()))?;
    }
    #[cfg(unix)]
    let existed = path.exists();
    let mut file = OpenOptions::new()
        .create(true)
        .append(true)
        .open(&path)
        .with_context(|| format!("failed to open {}", path.display()))?;
    #[cfg(unix)]
    if !existed {
        use std::os::unix::fs::PermissionsExt;
        let _ = fs::set_permissions(&path, fs::Permissions::from_mode(0o600));
    }

    let host = if port == 22 {
        host.to_string()
    } else {
        format!("[{host}]:{port}")
    };
    writeln!(file, "{host} {key_type} {key_data} meatshell")
        .with_context(|| format!("failed to write {}", path.display()))?;
    Ok(())
}

fn public_key_parts(key: &PublicKey) -> Result<(String, String)> {
    let encoded = key.to_openssh().context("encode server public key")?;
    let mut fields = encoded.split_whitespace();
    let key_type = fields
        .next()
        .context("server public key missing type")?
        .to_string();
    let key_data = fields
        .next()
        .context("server public key missing data")?
        .to_string();
    Ok((key_type, key_data))
}

fn host_patterns(host: &str, port: u16) -> Vec<String> {
    let mut out = Vec::new();
    out.push(host.to_ascii_lowercase());
    out.push(format!("[{}]:{}", host.to_ascii_lowercase(), port));
    if port == 22 {
        out.push(format!("[{}]:22", host.to_ascii_lowercase()));
    }
    out
}

#[derive(Debug, PartialEq, Eq)]
struct KnownHostEntry {
    hosts: Vec<String>,
    key_type: String,
    key_data: String,
}

fn parse_line(line: &str) -> Option<KnownHostEntry> {
    let trimmed = line.trim();
    if trimmed.is_empty() || trimmed.starts_with('#') {
        return None;
    }

    let mut fields = trimmed.split_whitespace();
    let first = fields.next()?;
    let hosts = if first.starts_with('@') {
        // Marker such as @cert-authority or @revoked. We only handle ordinary
        // host-key entries here; markers are skipped rather than misread.
        return None;
    } else {
        first
    };
    if hosts.starts_with("|1|") {
        // Hashed OpenSSH entries need the per-line HMAC salt. Keep support
        // explicit by ignoring them instead of producing false matches.
        return None;
    }
    let key_type = fields.next()?.to_string();
    let key_data = fields.next()?.to_string();
    Some(KnownHostEntry {
        hosts: hosts.split(',').map(|h| h.to_ascii_lowercase()).collect(),
        key_type,
        key_data,
    })
}

fn host_matches(pattern: &str, target: &str) -> bool {
    if pattern == target {
        return true;
    }
    if let Some(suffix) = pattern.strip_prefix("*.") {
        return target
            .strip_suffix(suffix)
            .map(|prefix| prefix.ends_with('.') && !prefix.trim_end_matches('.').is_empty())
            .unwrap_or(false);
    }
    false
}

#[cfg(test)]
mod tests {
    use super::{
        check_or_trust_in, host_matches, known_hosts_files, parse_line, HostKeyStatus,
        KnownHostEntry,
    };
    use std::fs;
    use std::time::{SystemTime, UNIX_EPOCH};

    fn fixture_key(key_type: &str, key_data: &str) -> ssh_key::PublicKey {
        format!("{key_type} {key_data}")
            .parse()
            .expect("valid public key fixture")
    }

    fn temp_known_hosts_path(name: &str) -> std::path::PathBuf {
        let unique = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_nanos();
        std::env::temp_dir()
            .join("meatshell-known-hosts-tests")
            .join(format!("{name}-{unique}"))
    }

    #[test]
    fn parses_plain_known_hosts_line() {
        assert_eq!(
            parse_line("example.com,192.0.2.1 ssh-ed25519 AAAA comment"),
            Some(KnownHostEntry {
                hosts: vec!["example.com".into(), "192.0.2.1".into()],
                key_type: "ssh-ed25519".into(),
                key_data: "AAAA".into(),
            })
        );
    }

    #[test]
    fn skips_hashed_and_marker_entries() {
        assert!(parse_line("|1|salt|hash ssh-ed25519 AAAA").is_none());
        assert!(parse_line("@cert-authority *.example.com ssh-ed25519 AAAA").is_none());
    }

    #[test]
    fn matches_exact_and_wildcard_hosts() {
        assert!(host_matches("example.com", "example.com"));
        assert!(host_matches("*.example.com", "db.prod.example.com"));
        assert!(!host_matches("*.example.com", "example.com"));
    }

    #[test]
    fn accepts_additional_key_type_for_known_host() {
        let path = temp_known_hosts_path("additional-key-type");
        if let Some(parent) = path.parent() {
            fs::create_dir_all(parent).unwrap();
        }
        fs::write(
            &path,
            "10.31.88.88 ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIHcL1R5w/NbHtGXf+qXCDP1KIXL4SMLJK/7UKcaGEJXC\n",
        )
        .unwrap();

        let key = fixture_key(
            "ecdsa-sha2-nistp256",
            "AAAAE2VjZHNhLXNoYTItbmlzdHAyNTYAAAAIbmlzdHAyNTYAAABBBHwf2HMM5TRXvo2SQJjsNkiDD5KqiiNjrGVv3UUh+mMT5RHxiRtOnlqvjhQtBq0VpmpCV/PwUdhOig4vkbqAcEc=",
        );
        let status = check_or_trust_in("10.31.88.88", 22, &key, &path).unwrap();

        assert_eq!(status, HostKeyStatus::Added);
        let raw = fs::read_to_string(path).unwrap();
        assert!(raw.contains("ssh-ed25519"));
        assert!(raw.contains("ecdsa-sha2-nistp256"));
    }

    #[test]
    fn rejects_changed_key_for_same_key_type() {
        let path = temp_known_hosts_path("changed-key");
        if let Some(parent) = path.parent() {
            fs::create_dir_all(parent).unwrap();
        }
        fs::write(
            &path,
            "10.31.88.88 ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIFgFZtmFxw40sXfFl6kOVb9zVTiDjXCh8y2jhkJ/AH63\n",
        )
        .unwrap();

        let key = fixture_key(
            "ssh-ed25519",
            "AAAAC3NzaC1lZDI1NTE5AAAAIHcL1R5w/NbHtGXf+qXCDP1KIXL4SMLJK/7UKcaGEJXC",
        );
        let status = check_or_trust_in("10.31.88.88", 22, &key, &path).unwrap();

        assert_eq!(status, HostKeyStatus::Changed);
    }

    #[test]
    fn only_app_known_hosts_can_reject_changed_keys() {
        let app_path = temp_known_hosts_path("source-flags");
        let files = known_hosts_files(&app_path);
        assert!(files.len() >= 2);
        assert!(files.iter().any(|(path, can_reject)| path == &app_path && *can_reject));
        assert!(files.iter().any(|(_, can_reject)| !*can_reject));
    }
}
