package com.box.sdk;

import com.box.sdk.http.HttpMethod;
import com.eclipsesource.json.Json;
import com.eclipsesource.json.JsonArray;
import com.eclipsesource.json.JsonObject;
import java.net.URL;
import java.util.List;

public final class BoxAI {

  /** Ask AI url. */
  public static final URLTemplate SEND_AI_REQUEST_URL = new URLTemplate("ai/ask");
  /** Text gen AI url. */
  public static final URLTemplate SEND_AI_TEXT_GEN_REQUEST_URL = new URLTemplate("ai/text_gen");
  /** AI agent default config url. */
  public static final URLTemplate AI_AGENT_DEFAULT_CONFIG_URL = new URLTemplate("ai_agent_default");
  /** AI extract metadata freeform url. */
  public static final URLTemplate EXTRACT_METADATA_FREEFORM_URL = new URLTemplate("ai/extract");
  /** AI extract metadata structured url. */
  public static final URLTemplate EXTRACT_METADATA_STRUCTURED_URL =
      new URLTemplate("ai/extract_structured");

  private BoxAI() {}

  /**
   * Sends an AI request to supported LLMs and returns an answer specifically focused on the user's
   * question given the provided items.
   *
   * @param api the API connection to be used by the created user.
   * @param prompt The prompt provided by the client to be answered by the LLM.
   * @param items The items to be processed by the LLM, currently only files are supported.
   * @param mode The mode specifies if this request is for a single or multiple items.
   * @return The response from the AI.
   */
  public static BoxAIResponse sendAIRequest(
      BoxAPIConnection api, String prompt, List<BoxAIItem> items, Mode mode) {
    return sendAIRequest(api, prompt, items, mode, null, null, null);
  }

  /**
   * Sends an AI request to supported LLMs and returns an answer specifically focused on the user's
   * question given the provided items.
   *
   * @param api the API connection to be used by the created user.
   * @param prompt The prompt provided by the client to be answered by the LLM.
   * @param items The items to be processed by the LLM, currently only files are supported.
   * @param mode The mode specifies if this request is for a single or multiple items.
   * @param dialogueHistory The history of prompts and answers previously passed to the LLM. This
   *     provides additional context to the LLM in generating the response.
   * @param agent The AI agent configuration to be used for the request.
   * @param includeCitations Whether to include citations in the response.
   * @return The response from the AI.
   */
  public static BoxAIResponse sendAIRequest(
      BoxAPIConnection api,
      String prompt,
      List<BoxAIItem> items,
      Mode mode,
      List<BoxAIDialogueEntry> dialogueHistory,
      BoxAIAgentAsk agent,
      Boolean includeCitations) {
    URL url = SEND_AI_REQUEST_URL.build(api.getBaseURL());
    JsonObject requestJSON = new JsonObject();
    requestJSON.add("mode", mode.toString());
    requestJSON.add("prompt", prompt);

    JsonArray itemsJSON = new JsonArray();
    for (BoxAIItem item : items) {
      itemsJSON.add(item.getJSONObject());
    }
    requestJSON.add("items", itemsJSON);

    if (dialogueHistory != null) {
      JsonArray dialogueHistoryJSON = new JsonArray();
      for (BoxAIDialogueEntry dialogueEntry : dialogueHistory) {
        dialogueHistoryJSON.add(dialogueEntry.getJSONObject());
      }
      requestJSON.add("dialogue_history", dialogueHistoryJSON);
    }
    if (agent != null) {
      requestJSON.add("ai_agent", agent.getJSONObject());
    }
    if (includeCitations != null) {
      requestJSON.add("include_citations", includeCitations);
    }

    BoxJSONRequest req = new BoxJSONRequest(api, url, HttpMethod.POST);
    req.setBody(requestJSON.toString());

    try (BoxJSONResponse response = req.send()) {
      JsonObject responseJSON = Json.parse(response.getJSON()).asObject();
      return new BoxAIResponse(responseJSON);
    }
  }

  /**
   * Sends an AI request to supported LLMs and returns an answer specifically focused on the
   * creation of new text.
   *
   * @param api the API connection to be used by the created user.
   * @param prompt The prompt provided by the client to be answered by the LLM.
   * @param items The items to be processed by the LLM, currently only files are supported.
   * @return The response from the AI.
   */
  public static BoxAIResponse sendAITextGenRequest(
      BoxAPIConnection api, String prompt, List<BoxAIItem> items) {
    return sendAITextGenRequest(api, prompt, items, null);
  }

  /**
   * Sends an AI request to supported LLMs and returns an answer specifically focused on the
   * creation of new text.
   *
   * @param api the API connection to be used by the created user.
   * @param prompt The prompt provided by the client to be answered by the LLM.
   * @param items The items to be processed by the LLM, currently only files are supported.
   * @param dialogueHistory The history of prompts and answers previously passed to the LLM. This
   *     provides additional context to the LLM in generating the response.
   * @return The response from the AI.
   */
  public static BoxAIResponse sendAITextGenRequest(
      BoxAPIConnection api,
      String prompt,
      List<BoxAIItem> items,
      List<BoxAIDialogueEntry> dialogueHistory) {
    return sendAITextGenRequest(api, prompt, items, dialogueHistory, null);
  }

  /**
   * Sends an AI request to supported LLMs and returns an answer specifically focused on the
   * creation of new text.
   *
   * @param api the API connection to be used by the created user.
   * @param prompt The prompt provided by the client to be answered by the LLM.
   * @param items The items to be processed by the LLM, currently only files are supported.
   * @param dialogueHistory The history of prompts and answers previously passed to the LLM. This
   *     provides additional context to the LLM in generating the response.
   * @param agent The AI agent configuration to be used for the request.
   * @return The response from the AI.
   */
  public static BoxAIResponse sendAITextGenRequest(
      BoxAPIConnection api,
      String prompt,
      List<BoxAIItem> items,
      List<BoxAIDialogueEntry> dialogueHistory,
      BoxAIAgentTextGen agent) {
    URL url = SEND_AI_TEXT_GEN_REQUEST_URL.build(api.getBaseURL());
    JsonObject requestJSON = new JsonObject();
    requestJSON.add("prompt", prompt);

    JsonArray itemsJSON = new JsonArray();
    for (BoxAIItem item : items) {
      itemsJSON.add(item.getJSONObject());
    }
    requestJSON.add("items", itemsJSON);

    if (dialogueHistory != null) {
      JsonArray dialogueHistoryJSON = new JsonArray();
      for (BoxAIDialogueEntry dialogueEntry : dialogueHistory) {
        dialogueHistoryJSON.add(dialogueEntry.getJSONObject());
      }
      requestJSON.add("dialogue_history", dialogueHistoryJSON);
    }

    if (agent != null) {
      requestJSON.add("ai_agent", agent.getJSONObject());
    }

    BoxJSONRequest req = new BoxJSONRequest(api, url, HttpMethod.POST);
    req.setBody(requestJSON.toString());

    try (BoxJSONResponse response = req.send()) {
      JsonObject responseJSON = Json.parse(response.getJSON()).asObject();
      return new BoxAIResponse(responseJSON);
    }
  }

  /**
   * Get the default AI Agent use for the given mode.
   *
   * @param api The API connection to be used by the created user.
   * @param mode The mode to filter the agent config to return.
   * @return A successful response including the default agent configuration.
   */
  public static BoxAIAgent getAiAgentDefaultConfig(BoxAPIConnection api, BoxAIAgent.Mode mode) {
    return getAiAgentDefaultConfig(api, mode, null, null);
  }

  /**
   * Get the default AI Agent use for the given mode.
   *
   * @param api The API connection to be used by the created user.
   * @param mode The mode to filter the agent config to return.
   * @param language The language to filter the agent config to return.
   * @param model The model to filter the agent config to return.
   * @return A successful response including the default agent configuration.
   */
  public static BoxAIAgent getAiAgentDefaultConfig(
      BoxAPIConnection api, BoxAIAgent.Mode mode, String language, String model) {
    QueryStringBuilder builder = new QueryStringBuilder();
    builder.appendParam("mode", mode.toString());
    if (language != null) {
      builder.appendParam("language", language);
    }
    if (model != null) {
      builder.appendParam("model", model);
    }
    URL url = AI_AGENT_DEFAULT_CONFIG_URL.buildWithQuery(api.getBaseURL(), builder.toString());
    BoxAPIRequest req = new BoxAPIRequest(api, url, HttpMethod.GET);
    try (BoxJSONResponse response = (BoxJSONResponse) req.send()) {
      JsonObject responseJSON = Json.parse(response.getJSON()).asObject();
      return BoxAIAgent.parse(responseJSON);
    }
  }

  public enum Mode {
    /** Multiple items */
    MULTIPLE_ITEM_QA("multiple_item_qa"),

    /** Single item */
    SINGLE_ITEM_QA("single_item_qa");

    private final String mode;

    Mode(String mode) {
      this.mode = mode;
    }

    static BoxAI.Mode fromJSONValue(String jsonValue) {
      if (jsonValue.equals("multiple_item_qa")) {
        return BoxAI.Mode.MULTIPLE_ITEM_QA;
      } else if (jsonValue.equals("single_item_qa")) {
        return BoxAI.Mode.SINGLE_ITEM_QA;
      } else {
        System.out.print("Invalid AI mode.");
        return null;
      }
    }

    String toJSONValue() {
      return this.mode;
    }

    public String toString() {
      return this.mode;
    }
  }

  /**
   * Sends an AI request to supported Large Language Models (LLMs) and extracts metadata in form of
   * key-value pairs. Freeform metadata extraction does not require any metadata template setup
   * before sending the request.
   *
   * @param api the API connection to be used by the created user.
   * @param prompt The prompt provided by the client to be answered by the LLM.
   * @param items The items to be processed by the LLM, currently only files are supported.
   * @return The response from the AI.
   */
  public static BoxAIResponse extractMetadataFreeform(
      BoxAPIConnection api, String prompt, List<BoxAIItem> items) {
    return extractMetadataFreeform(api, prompt, items, null);
  }

  /**
   * Sends an AI request to supported Large Language Models (LLMs) and extracts metadata in form of
   * key-value pairs. Freeform metadata extraction does not require any metadata template setup
   * before sending the request.
   *
   * @param api the API connection to be used by the created user.
   * @param prompt The prompt provided by the client to be answered by the LLM.
   * @param items The items to be processed by the LLM, currently only files are supported.
   * @param agent The AI agent configuration to be used for the request.
   * @return The response from the AI.
   */
  public static BoxAIResponse extractMetadataFreeform(
      BoxAPIConnection api, String prompt, List<BoxAIItem> items, BoxAIAgentExtract agent) {
    URL url = EXTRACT_METADATA_FREEFORM_URL.build(api.getBaseURL());

    JsonObject requestJSON = new JsonObject();
    JsonArray itemsJSON = new JsonArray();
    for (BoxAIItem item : items) {
      itemsJSON.add(item.getJSONObject());
    }
    requestJSON.add("items", itemsJSON);
    requestJSON.add("prompt", prompt);
    if (agent != null) {
      requestJSON.add("ai_agent", agent.getJSONObject());
    }

    BoxJSONRequest req = new BoxJSONRequest(api, url, HttpMethod.POST);
    req.setBody(requestJSON.toString());

    try (BoxJSONResponse response = req.send()) {
      JsonObject responseJSON = Json.parse(response.getJSON()).asObject();
      return new BoxAIResponse(responseJSON);
    }
  }

  /**
   * Sends an AI request to supported Large Language Models (LLMs) and returns extracted metadata as
   * a set of key-value pairs. For this request, you need to use an already defined metadata
   * template or a define a schema yourself.
   *
   * @param api The API connection to be used by the created user.
   * @param items The items to be processed by the LLM, currently only files are supported.
   * @param template The metadata template to be used for the request.
   * @return The response from the AI.
   */
  public static BoxAIExtractStructuredResponse extractMetadataStructured(
      BoxAPIConnection api, List<BoxAIItem> items, BoxAIExtractMetadataTemplate template) {
    return extractMetadataStructured(api, items, template, null, null);
  }

  /**
   * Sends an AI request to supported Large Language Models (LLMs) and returns extracted metadata as
   * a set of key-value pairs. For this request, you need to use an already defined metadata
   * template or a define a schema yourself.
   *
   * @param api The API connection to be used by the created user.
   * @param items The items to be processed by the LLM, currently only files are supported.
   * @param template The metadata template to be used for the request.
   * @param agent The AI agent configuration to be used for the request.
   * @return The response from the AI.
   */
  public static BoxAIExtractStructuredResponse extractMetadataStructured(
      BoxAPIConnection api,
      List<BoxAIItem> items,
      BoxAIExtractMetadataTemplate template,
      BoxAIAgentExtractStructured agent) {
    return extractMetadataStructured(api, items, template, null, agent);
  }

  /**
   * Sends an AI request to supported Large Language Models (LLMs) and returns extracted metadata as
   * a set of key-value pairs. For this request, you need to use an already defined metadata
   * template or a define a schema yourself.
   *
   * @param api The API connection to be used by the created user.
   * @param items The items to be processed by the LLM, currently only files are supported.
   * @param fields The fields to be extracted from the items.
   * @return The response from the AI.
   */
  public static BoxAIExtractStructuredResponse extractMetadataStructured(
      BoxAPIConnection api, List<BoxAIItem> items, List<BoxAIExtractField> fields) {
    return extractMetadataStructured(api, items, null, fields, null);
  }

  /**
   * Sends an AI request to supported Large Language Models (LLMs) and returns extracted metadata as
   * a set of key-value pairs. For this request, you need to use an already defined metadata
   * template or a define a schema yourself.
   *
   * @param api The API connection to be used by the created user.
   * @param items The items to be processed by the LLM, currently only files are supported.
   * @param fields The fields to be extracted from the items.
   * @param agent The AI agent configuration to be used for the request.
   * @return The response from the AI.
   */
  public static BoxAIExtractStructuredResponse extractMetadataStructured(
      BoxAPIConnection api,
      List<BoxAIItem> items,
      List<BoxAIExtractField> fields,
      BoxAIAgentExtractStructured agent) {
    return extractMetadataStructured(api, items, null, fields, agent);
  }

  /**
   * Sends an AI request to supported Large Language Models (LLMs) and returns extracted metadata as
   * a set of key-value pairs. For this request, you need to use an already defined metadata
   * template or a define a schema yourself.
   *
   * @param api The API connection to be used by the created user.
   * @param items The items to be processed by the LLM, currently only files are supported.
   * @param template The metadata template to be used for the request.
   * @param fields The fields to be extracted from the items.
   * @param agent The AI agent configuration to be used for the request.
   * @return The response from the AI.
   */
  private static BoxAIExtractStructuredResponse extractMetadataStructured(
      BoxAPIConnection api,
      List<BoxAIItem> items,
      BoxAIExtractMetadataTemplate template,
      List<BoxAIExtractField> fields,
      BoxAIAgentExtractStructured agent) {
    URL url = EXTRACT_METADATA_STRUCTURED_URL.build(api.getBaseURL());

    JsonObject requestJSON = new JsonObject();
    JsonArray itemsJSON = new JsonArray();
    for (BoxAIItem item : items) {
      itemsJSON.add(item.getJSONObject());
    }
    requestJSON.add("items", itemsJSON);

    if (template != null) {
      requestJSON.add("metadata_template", template.getJSONObject());
    }

    if (fields != null) {
      JsonArray fieldsJSON = new JsonArray();
      for (BoxAIExtractField field : fields) {
        fieldsJSON.add(field.getJSONObject());
      }
      requestJSON.add("fields", fieldsJSON);
    }

    if (agent != null) {
      requestJSON.add("ai_agent", agent.getJSONObject());
    }

    BoxJSONRequest req = new BoxJSONRequest(api, url, HttpMethod.POST);
    req.setBody(requestJSON.toString());

    try (BoxJSONResponse response = req.send()) {
      return new BoxAIExtractStructuredResponse(response.getJSON());
    }
  }
}