// Copyright (C) 2023-2026 Intel Corporation
// SPDX-License-Identifier: Apache-2.0

#include <cxxopts.hpp>
#include <cstdint>
#include <optional>
#include <filesystem>

#include "load_image.hpp"
#include <openvino/genai/visual_language/pipeline.hpp>
#include "../text_generation/read_prompt_from_file.h"

int main(int argc, char* argv[]) try {
    cxxopts::Options options("benchmark_vlm", "Help command");

    options.add_options()
    ("m,model", "Path to model and tokenizers base directory", cxxopts::value<std::string>()->default_value(""))
    ("p,prompt", "Prompt", cxxopts::value<std::string>()->default_value(""))
    ("F,prompt_file", "Read prompt from file", cxxopts::value<std::string>())
    ("i,image", "Path to image. Can be a single image or a directory of images.", cxxopts::value<std::string>()->default_value("image.jpg"))
    ("H,image_height", "Target image height (if resizing is needed)", cxxopts::value<int32_t>())
    ("W,image_width", "Target image width (if resizing is needed)", cxxopts::value<int32_t>())
    ("N,num_warmup", "Number of warmup iterations", cxxopts::value<size_t>()->default_value(std::to_string(1)))
    ("n,num_iter", "Number of iterations", cxxopts::value<size_t>()->default_value(std::to_string(3)))
    ("M,max_new_tokens", "Maximal number of new tokens", cxxopts::value<size_t>()->default_value(std::to_string(20)))
    ("d,device", "Device to run the model on", cxxopts::value<std::string>()->default_value("CPU"))
    ("P,pruning_ratio", "(optional): Percentage of visual tokens to prune (valid range: 0-100); if this option is not provided, pruning is disabled.", cxxopts::value<size_t>())
    ("R,relevance_weight", "(optional): Float value from 0 to 1, controls the trade-off between diversity and relevance for visual tokens pruning; a value of 0 disables relevance weighting, while higher values (up to 1.0) emphasize relevance, making pruning more conservative on borderline tokens.", cxxopts::value<float>())
    ("h,help", "Print usage");

    cxxopts::ParseResult result;
    try {
        result = options.parse(argc, argv);
    } catch (const cxxopts::exceptions::exception& e) {
        std::cout << e.what() << "\n\n";
        std::cout << options.help() << std::endl;
        return EXIT_FAILURE;
    }

    if (result.count("help")) {
        std::cout << options.help() << std::endl;
        return EXIT_SUCCESS;
    }

    std::string prompt;
    if (result.count("prompt") && result.count("prompt_file")) {
        std::cout << "Prompt and prompt file should not exist together!" << std::endl;
        return EXIT_FAILURE;
    } else {
        if (result.count("prompt_file")) {
            prompt = utils::read_prompt(result["prompt_file"].as<std::string>());
        } else {
            prompt = result["prompt"].as<std::string>().empty() ? "What is on the image?" : result["prompt"].as<std::string>();
        }
    }
    if (prompt.empty()) {
        std::cout << "Prompt is empty!" << std::endl;
        return EXIT_FAILURE;
    } 

    std::string models_path;
    if (result.count("model")) {
        models_path = result["model"].as<std::string>();
    } else {
        std::cout << "Model path is not provided!" << std::endl;
        return EXIT_FAILURE;
    }

    const std::string image_path = result["image"].as<std::string>();
    std::string device = result["device"].as<std::string>();
    size_t num_warmup = result["num_warmup"].as<size_t>();
    size_t num_iter = result["num_iter"].as<size_t>();

    if (result.count("image_height") != result.count("image_width")) {
        std::cout << "image_height and image_width must be provided together!" << std::endl;
        return EXIT_FAILURE;
    }
    const std::optional<utils::ImageSize> image_size = 
        (result.count("image_width") && result.count("image_height"))
        ? std::optional<utils::ImageSize>{utils::ImageSize{result["image_width"].as<int32_t>(), result["image_height"].as<int32_t>()}}
        : std::nullopt;
    
    std::vector<ov::Tensor> images = utils::load_images(image_path, image_size);

    ov::genai::GenerationConfig config;
    if (result.count("pruning_ratio")) {
        config.pruning_ratio = result["pruning_ratio"].as<size_t>();
    }
    if (result.count("relevance_weight")) {
        config.relevance_weight = result["relevance_weight"].as<float>();
    }
    config.max_new_tokens = result["max_new_tokens"].as<size_t>();
    config.ignore_eos = true;

    std::cout << ov::get_openvino_version() << std::endl;

    std::unique_ptr<ov::genai::VLMPipeline> pipe;
    if (device == "NPU")
        pipe = std::make_unique<ov::genai::VLMPipeline>(models_path, device);
    else {
        // Setting of Scheduler config will trigger usage of ContinuousBatching pipeline, which is not default for Qwen2VL, Qwen2.5VL, Gemma3 due to accuracy issues.
        ov::genai::SchedulerConfig scheduler_config;
        scheduler_config.enable_prefix_caching = false;
        scheduler_config.max_num_batched_tokens = std::numeric_limits<std::size_t>::max();
        pipe = std::make_unique<ov::genai::VLMPipeline>(models_path, device, ov::genai::scheduler_config(scheduler_config));
    }

    auto input_data = pipe->get_tokenizer().encode(prompt);
    size_t prompt_token_size = input_data.input_ids.get_shape()[1];
    std::cout << "Number of images: " << images.size() << ", Prompt token size: " << prompt_token_size << std::endl;

    for (size_t i = 0; i < num_warmup; i++)
        pipe->generate(prompt, ov::genai::images(images), ov::genai::generation_config(config));
    
    auto res = pipe->generate(prompt, ov::genai::images(images), ov::genai::generation_config(config));
    auto metrics = res.perf_metrics;
    for (size_t i = 0; i < num_iter - 1; i++) {
        res = pipe->generate(prompt, ov::genai::images(images), ov::genai::generation_config(config));
        metrics = metrics + res.perf_metrics;
    }

    std::cout << std::fixed << std::setprecision(2);
    if (image_size.has_value()) {
        std::cout << "Image is resized to: " << image_size->width << "x" << image_size->height << std::endl; 
    }
    std::cout << "Input token size: " << res.perf_metrics.get_num_input_tokens() << std::endl;
    std::cout << "Output token size: " << res.perf_metrics.get_num_generated_tokens() << std::endl;
    std::cout << "Load time: " << metrics.get_load_time() << " ms" << std::endl;
    std::cout << "Generate time: " << metrics.get_generate_duration().mean << " ± " << metrics.get_generate_duration().std << " ms" << std::endl;
    std::cout << "Tokenization time: " << metrics.get_tokenization_duration().mean << " ± " << metrics.get_tokenization_duration().std << " ms" << std::endl;
    std::cout << "Detokenization time: " << metrics.get_detokenization_duration().mean << " ± " << metrics.get_detokenization_duration().std << " ms" << std::endl;
    std::cout << "Embeddings preparation time: " << metrics.get_prepare_embeddings_duration().mean << " ± " << metrics.get_prepare_embeddings_duration().std << " ms" << std::endl;
    std::cout << "TTFT: " << metrics.get_ttft().mean  << " ± " << metrics.get_ttft().std << " ms" << std::endl;
    std::cout << "TPOT: " << metrics.get_tpot().mean  << " ± " << metrics.get_tpot().std << " ms/token " << std::endl;
    std::cout << "Throughput: " << metrics.get_throughput().mean  << " ± " << metrics.get_throughput().std << " tokens/s" << std::endl;

    return 0;
} catch (const std::exception& error) {
    try {
        std::cerr << error.what() << '\n';
    } catch (const std::ios_base::failure&) {}
    return EXIT_FAILURE;
} catch (...) {
    try {
        std::cerr << "Non-exception object thrown\n";
    } catch (const std::ios_base::failure&) {}
    return EXIT_FAILURE;
}