#include <iostream>
#include <vector>
#include <string>
#include <sstream>
#include <cstdlib>
#include <ctime>
#include <dirent.h>
#include <algorithm>
#include <sys/stat.h>
#include <unistd.h>

// 检查目录是否存在
bool dir_exists(const std::string& path) {
    struct stat info;
    return stat(path.c_str(), &info) == 0 && S_ISDIR(info.st_mode);
}

// 找到最新备份目录（按时间排序取最大）
std::string find_latest_backup_dir(const std::string& base_path) {
    DIR* dir = opendir(base_path.c_str());
    if (!dir) return "";

    std::vector<std::string> candidates;
    struct dirent* entry;
    while ((entry = readdir(dir)) != nullptr) {
        std::string name = entry->d_name;
        if (name.find("backups") == 0) {
            std::string full_path = base_path + "/" + name;
            if (dir_exists(full_path)) {
                candidates.push_back(name);
            }
        }
    }
    closedir(dir);

    if (candidates.empty()) return "";

    std::sort(candidates.begin(), candidates.end()); // 按名字排序（时间戳在名字中）
    return base_path + "/" + candidates.back();      // 返回最新的备份目录完整路径
}

// 判断文件是否存在
bool file_exists(const std::string& path) {
    return access(path.c_str(), F_OK) == 0;
}

int main(int argc, char* argv[]) {
    if (argc < 2) {
        std::cerr << "用法: " << argv[0] << " dbname1 [dbname2 ...]\n";
        return 1;
    }

    // 参数配置
    std::string base_backup_dir = "/var/postgresql_backups";
    std::string psql_path = "/usr/local/postgresql/bin/psql";  // 或实际路径
    std::string user = "postgres";
    std::string host = "127.0.0.1";
    int port = 5432;

    // 查找最新备份目录
    std::string backup_dir = find_latest_backup_dir(base_backup_dir);
    if (backup_dir.empty()) {
        std::cerr << "未找到任何备份目录。\n";
        return 1;
    }

    std::cout << "使用备份目录: " << backup_dir << "\n";

    // 获取用户输入密码（隐藏）
    std::string password;
    char* pw = getpass("请输入数据库密码: ");
    if (pw) {
        password = pw;
    } else {
        std::cerr << "无法读取密码\n";
        return 1;
    }
    setenv("PGPASSWORD", password.c_str(), 1);

    // 收集数据库名并自动追加 dekmaster
    std::vector<std::string> databases;
    for (int i = 1; i < argc; ++i) {
        databases.push_back(argv[i]);
    }
    databases.push_back("dekmaster");

    // 去重
    std::sort(databases.begin(), databases.end());
    databases.erase(std::unique(databases.begin(), databases.end()), databases.end());

    // 恢复每个数据库
    for (const auto& dbname : databases) {
        std::string sql_file = backup_dir + "/" + dbname + "_backup.sql";
        std::string used_user = (dbname == "dekmaster") ? "dekmaster" : user;

        if (!file_exists(sql_file)) {
            std::cerr << "找不到备份文件: " << sql_file << "\n";
            continue;
        }

        // 执行 psql 恢复
        std::ostringstream cmd;
        cmd << psql_path
            << " -h " << host
            << " -p " << port
            << " -U " << used_user
            << " -d " << dbname
            << " -f \"" << sql_file << "\"";

        std::cout << "正在恢复数据库: " << dbname << " （用户：" << used_user << "）\n";
        std::cout << "执行命令: " << cmd.str() << "\n";
        int result = std::system(cmd.str().c_str());

        if (result == 0) {
            std::cout << "恢复成功: " << dbname << "\n";
        } else {
            std::cerr << "恢复失败: " << dbname << "\n";
        }
    }

        // 恢复密钥文件和映射文件
    std::string key_file_backup = backup_dir + "/output.json";
    std::string map_file_backup = backup_dir + "/map.json";

    std::string key_file_target = "/new_enc/kms_test/output.json";
    std::string map_file_target = "/etc/encryptsql/map.json";

    // 恢复密钥文件
    if (file_exists(key_file_backup)) {
        std::string cp_cmd = "cp \"" + key_file_backup + "\" \"" + key_file_target + "\"";
        int result = std::system(cp_cmd.c_str());
        if (result == 0)
            std::cout << "密钥文件恢复成功: " << key_file_target << "\n";
        else
            std::cerr << "密钥文件恢复失败。\n";
    } else {
        std::cerr << "备份中未找到密钥文件。\n";
    }

    // 恢复映射文件
    if (file_exists(map_file_backup)) {
        std::string cp_cmd = "cp \"" + map_file_backup + "\" \"" + map_file_target + "\"";
        int result = std::system(cp_cmd.c_str());
        if (result == 0)
            std::cout << "映射文件恢复成功: " << map_file_target << "\n";
        else
            std::cerr << "映射文件恢复失败。\n";
    } else {
        std::cerr << "备份中未找到映射文件。\n";
    }


    std::cout << "所有恢复任务完成。\n";
    return 0;
}
