#include <iostream>
#include <memory>
#include <string>
#include <vector>
#include <algorithm>
#include "../nodes/vp_file_src_node.h"
#include "../nodes/infers/vp_yunet_face_detector_node.h"
#include "../nodes/infers/vp_face_swap_node.h"
#include "../nodes/osd/vp_face_osd_node.h"
#include "../nodes/vp_screen_des_node.h"
#include "../nodes/vp_file_des_node.h"
#include "../utils/analysis_board/vp_analysis_board.h"
#include <experimental/filesystem> // C++17 必须

/*
* ## face_swap_sample ##
* swap face for any video/images, no training need before running.
*/

std::string get_filename_without_extension(const std::string& path) {
    std::experimental::filesystem::path fs_path(path);
    return fs_path.stem().string(); // stem() 获取不带扩展名的文件名
}
// 解析命令行参数的函数
bool has_argument(const std::vector<std::string>& args, const std::string& option) {
    return std::find(args.begin(), args.end(), option) != args.end();
}

void ensure_directory_exists(const std::string& save_dir) {
    std::experimental::filesystem::path dir_path(save_dir);
    
    // 检查目录是否存在
    if (!std::experimental::filesystem::exists(dir_path)) {
        try {
            // 创建目录
            std::experimental::filesystem::create_directories(dir_path);  // 创建父目录及子目录
            std::cout << "Directory created: " << save_dir << std::endl;
        } catch (const std::experimental::filesystem::filesystem_error& e) {
            std::cerr << "Error creating directory: " << e.what() << std::endl;
        }
    } else {
        std::cout << "Directory already exists: " << save_dir << std::endl;
    }
}

int main(int argc, char* argv[]) {
    VP_SET_LOG_LEVEL(vp_utils::vp_log_level::INFO);
    VP_SET_LOG_INCLUDE_CODE_LOCATION(false);
    VP_SET_LOG_INCLUDE_THREAD_ID(false);
    VP_LOGGER_INIT();

    // 默认的命令行参数
    std::string video_path = "./vp_data/test_video/face.mp4";
    std::string image_path = "./vp_data/inswapper/data/mans1.jpeg";
    bool show_screen = false;
    bool show_board = false;

    float resize_ratio = 0.5;

    // 将命令行参数存储为 vector
    std::vector<std::string> args(argv + 1, argv + argc);

    // 如果提供了第一个和第二个参数，则覆盖默认的路径
    if (args.size() >= 2) {
        video_path = args[0];
        image_path = args[1];
        if (args.size() >= 3){
            resize_ratio =std::stof(args[2]);
        }
    } else {
        std::cout << "Usage: " << argv[0] << " <video_path> <image_path> <resize_ration> [--show_screen] [--show_board]\n";
        std::cout << "Example: " << argv[0] << " ./vp_data/test_video/face.mp4 ./vp_data/inswapper/data/mans1.jpeg 0.5 --show_screen --show_board\n";
        return 1;
    }
    std::string video_filename = get_filename_without_extension(video_path);
    std::string image_filename = get_filename_without_extension(image_path);

    ensure_directory_exists(video_filename);


    // 检测可选参数
    show_screen = has_argument(args, "--show_screen");
    show_board = has_argument(args, "--show_board");

    std::cout << "Video path: " << video_path << "\n";
    std::cout << "Image path: " << image_path << "\n";
    std::cout << "Video Resize ratio: " << resize_ratio << "\n";
    std::cout << "Show screen: " << (show_screen ? "true" : "false") << "\n";
    std::cout << "Show board: " << (show_board ? "true" : "false") << "\n";

    // 创建节点
    //auto file_src_0 = std::make_shared<vp_nodes::vp_file_src_node>("file_src_0", 0, video_path, 1.0, true, "avdec_h264", 4);
    auto file_src_0 = std::make_shared<vp_nodes::vp_file_src_node>("file_src_0", 0, video_path, resize_ratio, false, "avdec_h264", 4);
    
    auto yunet_face_detector = std::make_shared<vp_nodes::vp_yunet_face_detector_node>("yunet_face_detector", "./vp_data/models/face/face_detection_yunet_2022mar.onnx");
    auto face_swap = std::make_shared<vp_nodes::vp_face_swap_node>(
        "face_swap", 
        "./vp_data/models/face/face_detection_yunet_2022mar.onnx", 
        "./vp_data/models/face/swap/w600k_r50.onnx", 
        "./vp_data/models/face/swap/emap.txt", 
        "./vp_data/models/face/swap/inswapper_128.onnx", 
        image_path
    );

    
    auto file_des_node_0 = std::make_shared<vp_nodes::vp_file_des_node>("file_des_0", 0, video_filename, image_filename, 1);

    // 构建管道
    yunet_face_detector->attach_to({file_src_0});
    face_swap->attach_to({yunet_face_detector});

    if (show_screen) {
        auto screen_des_0_ori = std::make_shared<vp_nodes::vp_screen_des_node>("screen_des_0_ori", 0, false);
        auto screen_des_0_osd = std::make_shared<vp_nodes::vp_screen_des_node>("screen_des_0_osd", 0);
        screen_des_0_ori->attach_to({face_swap});
        screen_des_0_osd->attach_to({face_swap});
    }

    file_des_node_0->attach_to({face_swap});

    file_src_0->start();

    // 如果用户要求显示 board，则显示
    if (show_board) {
        vp_utils::vp_analysis_board board({file_src_0});
        board.display(1, false);
    }

    std::string wait;
    std::getline(std::cin, wait);
    file_src_0->detach_recursively();

    return 0;
}