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

// 获取当前时间戳
std::string get_timestamp() {
    std::time_t t = std::time(nullptr);
    std::tm* now = std::localtime(&t);
    char buffer[20];
    std::strftime(buffer, sizeof(buffer), "%Y%m%d_%H%M%S", now);
    return std::string(buffer);
}

// 判断文件是否存在
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 pg_dump_path = "/usr/local/postgresql/bin/pg_dump";
    std::string user = "postgres";
    std::string host = "127.0.0.1";
    int port = 5432;

    // 创建备份目录
    std::string timestamp = get_timestamp();
    std::string backup_dir = "/var/postgresql_backups/backups" + timestamp;

    std::string mkdir_cmd = "mkdir -p \"" + backup_dir + "\"";
    if (std::system(mkdir_cmd.c_str()) != 0) {
        std::cerr << "创建备份目录失败: " << backup_dir << "\n";
        return 1;
    }
    std::cout << "创建备份目录: " << backup_dir << "\n";

    // 复制密钥文件和映射文件
    std::string key_file_src = "/new_enc/kms_test/output.json";
    std::string map_file_src = "/etc/encryptsql/map.json";
    std::string key_file_dst = backup_dir + "/output.json";
    std::string map_file_dst = backup_dir + "/map.json";

    if (file_exists(key_file_src)) {
        std::string cp_cmd = "cp \"" + key_file_src + "\" \"" + key_file_dst + "\"";
        std::system(cp_cmd.c_str());
        std::cout << "已复制密钥文件到: " << key_file_dst << "\n";
    } else {
        std::cerr << "密钥文件不存在: " << key_file_src << "\n";
    }

    if (file_exists(map_file_src)) {
        std::string cp_cmd = "cp \"" + map_file_src + "\" \"" + map_file_dst + "\"";
        std::system(cp_cmd.c_str());
        std::cout << "已复制映射文件到: " << map_file_dst << "\n";
    } else {
        std::cerr << "映射文件不存在: " << map_file_src << "\n";
    }

    // 收集数据库列表
    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());

    std::string password;

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

    // 设置环境变量 PGPASSWORD
    setenv("PGPASSWORD", password.c_str(), 1);  // 1表示强制覆盖已有变量

    bool all_success = true;

    // 开始备份
    for (const auto& dbname : databases) {
        std::string outfile = backup_dir + "/" + dbname + "_backup.sql";
        std::string used_user = (dbname == "dekmaster") ? "dekmaster" : user;

        std::ostringstream cmd;
        cmd << pg_dump_path
            << " -h " << host
            << " -p " << port
            << " -U " << used_user
            << " " << dbname
            << " > \"" << outfile << "\" 2> \"" << outfile << ".err\"";

        std::cout << "正在备份数据库: " << dbname << " （用户：" << used_user << "）\n";

        int result = std::system(cmd.str().c_str());

        if (result == 0) {
            std::cout << "备份成功: " << outfile << "\n";
        } else {
            all_success = false;

            // 读取错误输出文件判断是否是密码错误
            std::ifstream errfile(outfile + ".err");
            std::string line, all_err;
            while (std::getline(errfile, line)) {
                all_err += line + "\n";
            }
            errfile.close();

            if (all_err.find("FATAL:  password authentication failed") != std::string::npos) {
                std::cerr << "密码错误，无法连接数据库: " << dbname << "\n";
            } else {
                std::cerr << "备份失败: " << dbname << "\n";
                std::cerr << "错误信息：\n" << all_err;
            }

            // 删除已复制的密钥和映射文件以及备份目录
            std::cout << "正在清理已创建的文件...\n";
            std::string cleanup_cmd = "rm -rf \"" + backup_dir + "\"";
            std::system(cleanup_cmd.c_str());
            std::cout << "已删除备份目录: " << backup_dir << "\n";

            return 1;  // 立即中止程序
        }
    }

    std::cout << "\n所有数据库备份完成。\n";
    return 0;
}

