// -----------------------------------------------------------------------------------------------
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
//
// C++ API Example for Model Question-Answering
// This example demonstrates how to use the C++ API of the ONNX Runtime GenAI library
// to perform model question-answering tasks. It includes functionalities to create a model,
// tokenizer, and generator, and to handle user input for generating responses based on prompts.
// -----------------------------------------------------------------------------------------------

#include <csignal>
#include <iomanip>
#include <string>

#include "common.h"

OgaGenerator* g_generator = nullptr;

void TerminateGeneration(int signum) {
  if (g_generator == nullptr) {
    return;
  }
  g_generator->SetRuntimeOption("terminate_session", "1");
}

void CXX_API(
    GeneratorParamsArgs& generator_params_args,
    GuidanceArgs& guidance_args,
    const std::string& model_path,
    const std::string& ep,
    const std::string& ep_path,
    const std::vector<std::string>& image_paths,
    const std::vector<std::string>& audio_paths,
    const std::string& system_prompt,
    const std::string& user_prompt,
    bool verbose,
    bool debug,
    bool interactive) {
  if (debug) SetLogger();
  RegisterEP(ep, ep_path);

  if (verbose) std::cout << "Creating config..." << std::endl;
  std::unordered_map<std::string, std::string> ep_options;
  auto config = GetConfig(model_path, ep, ep_options, generator_params_args);

  if (verbose) std::cout << "Creating model..." << std::endl;
  auto model = OgaModel::Create(*config);

  if (verbose) std::cout << "Creating tokenizer..." << std::endl;
  auto tokenizer = OgaTokenizer::Create(*model);
  auto stream = OgaTokenizerStream::Create(*tokenizer);

  if (verbose) std::cout << "Creating processor..." << std::endl;
  auto processor = OgaMultiModalProcessor::Create(*model);

  // Create running list of messages
  std::vector<nlohmann::ordered_json> input_list;
  nlohmann::ordered_json system_message = nlohmann::ordered_json{{"role", "system"}, {"content", system_prompt}};
  input_list.push_back(system_message);

  // Get and set guidance info if requested
  std::string guidance_type, guidance_data, tools;
  if (!guidance_args.response_format.empty()) {
    std::cout << "Make sure your tool call start id and tool call end id are marked as special in tokenizer.json" << std::endl;
    std::tie(guidance_type, guidance_data, tools) = GetGuidance(
        guidance_args.response_format,
        guidance_args.tools_file,
        "",       // tools_str
        nullptr,  // tools
        guidance_args.text_output,
        guidance_args.tool_output,
        guidance_args.tool_call_start,
        guidance_args.tool_call_end);

    input_list[0]["tools"] = tools;
  }

  // Keep asking for input prompts in a loop
  while (true) {
    // Get images
    std::unique_ptr<OgaImages> images;
    int num_images;
    std::tie(images, num_images) = GetUserImages(image_paths, interactive);

    // Get audios
    std::unique_ptr<OgaAudios> audios;
    int num_audios;
    std::tie(audios, num_audios) = GetUserAudios(audio_paths, interactive);

    // Get user prompt
    std::string text = GetUserPrompt(user_prompt, interactive);
    signal(SIGINT, TerminateGeneration);
    if (text == "quit()") {
      break;  // Exit the loop
    }

    // Construct user content based on inputs
    auto type = model->GetType();
    nlohmann::ordered_json user_content = GetUserContent(std::string(type), num_images, num_audios, text);

    // Add user message to list of messages
    nlohmann::ordered_json user_message = nlohmann::ordered_json{{"role", "user"}, {"content", user_content}};
    input_list.push_back(user_message);
    nlohmann::ordered_json j = input_list;
    std::string messages = j.dump();

    // Start timings
    bool is_first_token = true;
    Timing timing;
    timing.RecordStartTimestamp();

    // Initialize generator params
    auto params = OgaGeneratorParams::Create(*model);
    SetSearchOptions(*params, generator_params_args, verbose);

    // Initialize guidance info
    if (!guidance_args.response_format.empty()) {
      params->SetGuidance(guidance_type.c_str(), guidance_data.c_str());
      if (verbose) {
        std::cout << std::endl;
        std::cout << "Guidance type is: " << guidance_type << std::endl;
        std::cout << "Guidance data is: \n"
                  << guidance_data << std::endl;
        std::cout << std::endl;
      }
    }

    // Create generator
    auto generator = OgaGenerator::Create(*model, *params);
    g_generator = generator.get();  // Store the current generator for termination
    if (verbose) std::cout << "Generator created" << std::endl;

    // Apply chat template
    std::string prompt;
    try {
      bool add_generation_prompt = true;
      prompt = ApplyChatTemplate(model_path, *tokenizer, messages, add_generation_prompt, tools);
    } catch (...) {
      prompt = text;
    }
    if (verbose) std::cout << "Prompt: " << prompt << "\n"
                           << std::endl;

    // Encode combined system + user prompt and append inputs to model
    auto input_tensors = processor->ProcessImagesAndAudios(prompt.c_str(), images.get(), audios.get());
    generator->SetInputs(*input_tensors);
    const int prompt_tokens_length = generator->TokenCount();

    // Run generation loop
    if (verbose) std::cout << "Running generation loop..." << std::endl;
    std::cout << std::endl;
    std::cout << "Output: ";
    try {
      while (!generator->IsDone()) {
        generator->GenerateNextToken();

        if (is_first_token) {
          timing.RecordFirstTokenTimestamp();
          is_first_token = false;
        }

        const auto new_token = generator->GetNextTokens()[0];
        std::cout << stream->Decode(new_token) << std::flush;
      }
    } catch (const std::exception& e) {
      std::cout << "\n"
                << "Terminating generation: " << e.what() << std::endl;
    }
    timing.RecordEndTimestamp();

    // Clear the generator after use
    g_generator = nullptr;

    // Remove user message from list of messages
    input_list.pop_back();

    const int new_tokens_length = generator->TokenCount() - prompt_tokens_length;
    timing.Log(prompt_tokens_length, new_tokens_length);

    std::cout << "\n\n\n";
    if (!interactive) break;
  }
}

int main(int argc, char** argv) {
  // Get command-line args
  GeneratorParamsArgs generator_params_args;
  GuidanceArgs guidance_args;
  std::string model_path, ep = "follow_config", ep_path = "", system_prompt = "You are a helpful AI assistant.", user_prompt = "What color is the sky?";
  bool verbose = false, debug = false, interactive = true, rewind = true;
  std::vector<std::string> image_paths;
  std::vector<std::string> audio_paths;
  if (!ParseArgs(argc, argv, generator_params_args, guidance_args, model_path, ep, ep_path, system_prompt, user_prompt, verbose, debug, interactive, rewind, image_paths, audio_paths)) {
    return -1;
  }

  // Responsible for cleaning up the library during shutdown
  OgaHandle handle;

  std::cout << "--------------------------" << std::endl;
  std::cout << "Hello, ORT GenAI Model-MM!" << std::endl;
  std::cout << "--------------------------" << std::endl;

  std::cout << "Model path: " << model_path << std::endl;
  std::cout << "Execution provider: " << ep << std::endl;
  if (!ep_path.empty()) std::cout << "Execution provider path: " << ep_path << std::endl;
  std::cout << "System prompt: " << system_prompt << std::endl;
  if (!interactive) std::cout << "User prompt: " << user_prompt << std::endl;
  std::cout << "Verbose: " << verbose << std::endl;
  std::cout << "Interactive: " << interactive << std::endl;
  std::cout << "--------------------------" << std::endl;
  std::cout << std::endl;

  try {
    CXX_API(generator_params_args, guidance_args, model_path, ep, ep_path, image_paths, audio_paths, system_prompt, user_prompt, verbose, debug, interactive);
  } catch (const std::exception& e) {
    std::cerr << "Error: " << e.what() << std::endl;
    return -1;
  }

  return 0;
}