#include "key_store.h"
#include <fstream>
#include <sstream>
#include <cstdlib>
#include <cstdio>
#include <cstring>

static std::string trim(const std::string& s) {
	size_t a = s.find_first_not_of(" \t\r\n");
	if (a == std::string::npos) return "";
	size_t b = s.find_last_not_of(" \t\r\n");
	return s.substr(a, b - a + 1);
}

static std::string get_fingerprint(const std::string& pubkey_path) {
	std::string cmd = "ssh-keygen -lf \"" + pubkey_path + "\" 2>/dev/null";
	FILE* pipe = popen(cmd.c_str(), "r");
	if (!pipe) return "";
	char buf[512];
	std::string result;
	if (fgets(buf, sizeof(buf), pipe)) {
		result = trim(buf);
		/* Remove the key type prefix like "(RSA)" */
		auto paren = result.find('(');
		if (paren != std::string::npos) {
			result = trim(result.substr(0, paren));
		}
	}
	pclose(pipe);
	return result;
}

KeyStore::KeyStore() {
	const char* home = getenv("HOME");
	if (!home) home = ".";
	config_path_ = std::string(home) + "/.config/sshclient/keys.conf";
	keys_dir_ = std::string(home) + "/.config/sshclient/keys";
	std::string cmd = "mkdir -p \"" + keys_dir_ + "\"";
	system(cmd.c_str());
	load();
}

void KeyStore::load() {
	keys_.clear();
	std::ifstream f(config_path_);
	if (!f.is_open()) return;

	SshKey current;
	bool in_key = false;
	std::string line;

	while (std::getline(f, line)) {
		line = trim(line);
		if (line.empty() || line[0] == '#') {
			if (in_key && !current.name.empty()) {
				keys_.push_back(current);
				current = SshKey();
				in_key = false;
			}
			continue;
		}
		if (line[0] == '[' && line.back() == ']') {
			if (in_key && !current.name.empty()) {
				keys_.push_back(current);
				current = SshKey();
			}
			current.name = trim(line.substr(1, line.size() - 2));
			in_key = true;
			continue;
		}
		if (!in_key) continue;
		auto eq = line.find('=');
		if (eq == std::string::npos) continue;
		std::string key = trim(line.substr(0, eq));
		std::string val = trim(line.substr(eq + 1));
		if (key == "pubkey_path") current.pubkey_path = val;
		else if (key == "privkey_path") current.privkey_path = val;
		else if (key == "fingerprint") current.fingerprint = val;
	}
	if (in_key && !current.name.empty()) {
		keys_.push_back(current);
	}
}

void KeyStore::save() const {
	std::ofstream f(config_path_);
	if (!f.is_open()) return;

	f << "# SSH Client - Stored Keys\n\n";
	for (const auto& k : keys_) {
		f << "[" << k.name << "]\n";
		f << "pubkey_path=" << k.pubkey_path << "\n";
		f << "privkey_path=" << k.privkey_path << "\n";
		f << "fingerprint=" << k.fingerprint << "\n\n";
	}
}

void KeyStore::add(const SshKey& k) {
	keys_.push_back(k);
	save();
}

void KeyStore::remove(int index) {
	if (index >= 0 && index < (int)keys_.size()) {
		keys_.erase(keys_.begin() + index);
		save();
	}
}

bool KeyStore::generate_keypair(const std::string& name, int bits) {
	std::string priv_path = keys_dir_ + "/" + name;
	std::string pub_path = priv_path + ".pub";

	std::string cmd = "ssh-keygen -t rsa -b " + std::to_string(bits) +
		" -f \"" + priv_path + "\" -N \"\" -C \"" + name + "\" 2>/dev/null";
	int rc = system(cmd.c_str());
	if (rc != 0) return false;

	SshKey k;
	k.name = name;
	k.privkey_path = priv_path;
	k.pubkey_path = pub_path;
	k.fingerprint = get_fingerprint(pub_path);

	add(k);
	return true;
}

bool KeyStore::import_key(const std::string& name, const std::string& privkey_path) {
	/* Validate key by running ssh-keygen -l */
	std::string check_cmd = "ssh-keygen -l -f \"" + privkey_path + "\" >/dev/null 2>&1";
	int rc = system(check_cmd.c_str());
	if (rc != 0) return false;

	std::string dest = keys_dir_ + "/" + name;
	std::string cp_cmd = "cp \"" + privkey_path + "\" \"" + dest + "\" && chmod 600 \"" + dest + "\"";
	rc = system(cp_cmd.c_str());
	if (rc != 0) return false;

	std::string pub_src = privkey_path + ".pub";
	std::string pub_dest = dest + ".pub";
	std::string cp_pub = "cp \"" + pub_src + "\" \"" + pub_dest + "\" 2>/dev/null";
	system(cp_pub.c_str());

	SshKey k;
	k.name = name;
	k.privkey_path = dest;
	k.pubkey_path = pub_dest;
	k.fingerprint = get_fingerprint(pub_dest);

	add(k);
	return true;
}

SshKey KeyStore::find_by_name(const std::string& name) const {
	for (const auto& k : keys_) {
		if (k.name == name) return k;
	}
	return {};
}