use mithril_client::models::NewSshKeyModel;
use std::path::{Path, PathBuf};

fn find_skypilot_private_keys(home: &Path) -> Vec<PathBuf> {
    let clients_dir = home.join(".sky").join("clients");
    let Ok(entries) = std::fs::read_dir(clients_dir) else {
        return Vec::new();
    };

    let mut key_paths = Vec::new();

    for entry in entries.flatten() {
        let Ok(file_type) = entry.file_type() else {
            continue;
        };
        if !file_type.is_dir() {
            continue;
        }

        let key_path = entry.path().join("ssh").join("sky-key");
        if key_path.exists() {
            key_paths.push(key_path);
        }
    }

    key_paths.sort();
    key_paths
}

fn local_private_key_candidates(home: &Path) -> Vec<PathBuf> {
    let ssh_dir = home.join(".ssh");
    let mut candidates = Vec::new();

    for file_name in ["id_ed25519", "id_rsa", "id_ecdsa", "id_dsa"] {
        let key_path = ssh_dir.join(file_name);
        if key_path.exists() {
            candidates.push(key_path);
        }
    }

    candidates.extend(find_skypilot_private_keys(home));
    candidates
}

pub fn normalize_pub_key(key: &str) -> String {
    // Extract just the key type and key data, ignoring comments.
    let parts: Vec<&str> = key.split_whitespace().collect();
    if parts.len() >= 2 {
        format!("{} {}", parts[0], parts[1])
    } else {
        key.trim().to_string()
    }
}

pub fn find_matching_ssh_key(
    platform_keys: &[NewSshKeyModel],
) -> Option<(&NewSshKeyModel, PathBuf)> {
    let home = std::env::var("HOME").ok()?;
    let home_path = Path::new(&home);

    let normalized_platform_keys: Vec<(&NewSshKeyModel, String)> = platform_keys
        .iter()
        .map(|k| (k, normalize_pub_key(&k.public_key)))
        .collect();

    for private_key_path in local_private_key_candidates(home_path) {
        let pub_key_path = PathBuf::from(format!("{}.pub", private_key_path.display()));
        let Ok(local_pub_key) = std::fs::read_to_string(pub_key_path) else {
            continue;
        };

        let local_pub_key = normalize_pub_key(&local_pub_key);
        for (platform_key, platform_pub_key) in &normalized_platform_keys {
            if local_pub_key == *platform_pub_key {
                return Some((*platform_key, private_key_path));
            }
        }
    }

    None
}
