// Licensed to the Apache Software Foundation (ASF) under one

// or more contributor license agreements.  See the NOTICE file

// distributed with this work for additional information

// regarding copyright ownership.  The ASF licenses this file

// to you under the Apache License, Version 2.0 (the

// "License"); you may not use this file except in compliance

// with the License.  You may obtain a copy of the License at

//

//   http://www.apache.org/licenses/LICENSE-2.0

//

// Unless required by applicable law or agreed to in writing,

// software distributed under the License is distributed on an

// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY

// KIND, either express or implied.  See the License for the

// specific language governing permissions and limitations

// under the License.

#include "core/column/column_array.h"

#include "exprs/function/ai/ai_classify.h"

#include "exprs/function/ai/ai_extract.h"

#include "exprs/function/ai/ai_filter.h"

#include "exprs/function/ai/ai_fix_grammar.h"

#include "exprs/function/ai/ai_generate.h"

#include "exprs/function/ai/ai_mask.h"

#include "exprs/function/ai/ai_sentiment.h"

#include "exprs/function/ai/ai_similarity.h"

#include "exprs/function/ai/ai_summarize.h"

#include "exprs/function/ai/ai_translate.h"

#include "exprs/function/ai/embed.h"

#include "exprs/function/simple_function_factory.h"

namespace doris {

Status FunctionAIClassify::build_prompt(const Block& block, const ColumnNumbers& arguments,

                                        size_t row_num, std::string& prompt) const {

    // Get the text column

    const ColumnWithTypeAndName& text_column = block.get_by_position(arguments[1]);

    StringRef text = text_column.column->get_data_at(row_num);

    std::string text_str = std::string(text.data, text.size);

    // Get the labels array column

    const ColumnWithTypeAndName& labels_column = block.get_by_position(arguments[2]);

    const auto& [array_column, array_row_num] =

            check_column_const_set_readability(*labels_column.column, row_num);

    const auto* col_array = check_and_get_column<ColumnArray>(*array_column);

    if (col_array == nullptr) {

        return Status::InternalError(

                "labels argument for {} must be Array(String) or Array(Varchar)", name);

    }

    std::vector<std::string> label_values;

    const auto& data = col_array->get_data();

    const auto& offsets = col_array->get_offsets();

    size_t start = array_row_num > 0 ? offsets[array_row_num - 1] : 0;

    size_t end = offsets[array_row_num];

    for (size_t i = start; i < end; ++i) {

        Field field;

        data.get(i, field);

        label_values.emplace_back(field.template get<TYPE_STRING>());

    }

    std::string labels_str = "[";

    for (size_t i = 0; i < label_values.size(); ++i) {

        if (i > 0) {

            labels_str += ", ";

        }

        labels_str += "\"" + label_values[i] + "\"";

    }

    labels_str += "]";

    prompt = "Labels: " + labels_str + "\nText: " + text_str;

    return Status::OK();

}

Status FunctionAIExtract::build_prompt(const Block& block, const ColumnNumbers& arguments,

                                       size_t row_num, std::string& prompt) const {

    // Get the text column

    const ColumnWithTypeAndName& text_column = block.get_by_position(arguments[1]);

    StringRef text = text_column.column->get_data_at(row_num);

    std::string text_str = std::string(text.data, text.size);

    // Get the labels array column

    const ColumnWithTypeAndName& labels_column = block.get_by_position(arguments[2]);

    const auto& [array_column, array_row_num] =

            check_column_const_set_readability(*labels_column.column, row_num);

    const auto* col_array = check_and_get_column<ColumnArray>(*array_column);

    if (col_array == nullptr) {

        return Status::InternalError(

                "labels argument for {} must be Array(String) or Array(Varchar)", name);

    }

    std::vector<std::string> label_values;

    const auto& offsets = col_array->get_offsets();

    const auto& data = col_array->get_data();

    size_t start = array_row_num > 0 ? offsets[array_row_num - 1] : 0;

    size_t end = offsets[array_row_num];

    for (size_t i = start; i < end; ++i) {

        Field field;

        data.get(i, field);

        label_values.emplace_back(field.template get<TYPE_STRING>());

    }

    std::string labels_str = "[";

    for (size_t i = 0; i < label_values.size(); ++i) {

        if (i > 0) {

            labels_str += ", ";

        }

        labels_str += "\"" + label_values[i] + "\"";

    }

    labels_str += "]";

    prompt = "Labels: " + labels_str + "\nText: " + text_str;

    return Status::OK();

}

Status FunctionAIGenerate::build_prompt(const Block& block, const ColumnNumbers& arguments,

                                        size_t row_num, std::string& prompt) const {

    const ColumnWithTypeAndName& text_column = block.get_by_position(arguments[1]);

    StringRef text_ref = text_column.column->get_data_at(row_num);

    prompt = std::string(text_ref.data, text_ref.size);

    return Status::OK();

}

Status FunctionAIMask::build_prompt(const Block& block, const ColumnNumbers& arguments,

                                    size_t row_num, std::string& prompt) const {

    // Get the text column

    const ColumnWithTypeAndName& text_column = block.get_by_position(arguments[1]);

    StringRef text = text_column.column->get_data_at(row_num);

    std::string text_str = std::string(text.data, text.size);

    // Get the labels array column

    const ColumnWithTypeAndName& labels_column = block.get_by_position(arguments[2]);

    const auto& [array_column, array_row_num] =

            check_column_const_set_readability(*labels_column.column, row_num);

    const auto* col_array = check_and_get_column<ColumnArray>(*array_column);

    if (col_array == nullptr) {

        return Status::InternalError(

                "labels argument for {} must be Array(String) or Array(Varchar)", name);

    }

    std::vector<std::string> label_values;

    const auto& offsets = col_array->get_offsets();

    const auto& data = col_array->get_data();

    size_t start = array_row_num > 0 ? offsets[array_row_num - 1] : 0;

    size_t end = offsets[array_row_num];

    for (size_t i = start; i < end; ++i) {

        Field field;

        data.get(i, field);

        label_values.emplace_back(field.template get<TYPE_STRING>());

    }

    std::string labels_str = "[";

    for (size_t i = 0; i < label_values.size(); ++i) {

        if (i > 0) {

            labels_str += ", ";

        }

        labels_str += "\"" + label_values[i] + "\"";

    }

    labels_str += "]";

    prompt = "Labels: " + labels_str + "\nText: " + text_str;

    return Status::OK();

}

Status FunctionAISimilarity::build_prompt(const Block& block, const ColumnNumbers& arguments,

                                          size_t row_num, std::string& prompt) const {

    // text1

    const ColumnWithTypeAndName& text_column_1 = block.get_by_position(arguments[1]);

    StringRef text_1 = text_column_1.column.get()->get_data_at(row_num);

    std::string text_str_1 = std::string(text_1.data, text_1.size);

    // text2

    const ColumnWithTypeAndName& text_column_2 = block.get_by_position(arguments[2]);

    StringRef text_2 = text_column_2.column.get()->get_data_at(row_num);

    std::string text_str_2 = std::string(text_2.data, text_2.size);

    prompt = "Text 1: " + text_str_1 + "\nText 2: " + text_str_2;

    return Status::OK();

}

Status FunctionAITranslate::build_prompt(const Block& block, const ColumnNumbers& arguments,

                                         size_t row_num, std::string& prompt) const {

    // text

    const ColumnWithTypeAndName& text_column = block.get_by_position(arguments[1]);

    StringRef text = text_column.column.get()->get_data_at(row_num);

    std::string text_str = std::string(text.data, text.size);

    // target language

    const ColumnWithTypeAndName& lang_column = block.get_by_position(arguments[2]);

    StringRef lang = lang_column.column.get()->get_data_at(row_num);

    std::string target_lang = std::string(lang.data, lang.size);

    prompt = "Translate the following text to " + target_lang + ".\nText: " + text_str;

    return Status::OK();

}

void register_function_ai(SimpleFunctionFactory& factory) {

    factory.register_function<FunctionEmbed>();

    factory.register_function<FunctionAIClassify>();

    factory.register_function<FunctionAIExtract>();

    factory.register_function<FunctionAIFilter>();

    factory.register_function<FunctionAIFixGrammar>();

    factory.register_function<FunctionAIGenerate>();

    factory.register_function<FunctionAIMask>();

    factory.register_function<FunctionAISentiment>();

    factory.register_function<FunctionAISimilarity>();

    factory.register_function<FunctionAISummarize>();

    factory.register_function<FunctionAITranslate>();

}

} // namespace doris