/******************************************************************************
 *
 * Copyright (c) 2019, the Perspective Authors.
 *
 * This file is part of the Perspective library, distributed under the terms of
 * the Apache License 2.0.  The full license can be found in the LICENSE file.
 *
 */
#ifdef PSP_ENABLE_PYTHON

#include <perspective/arrow_loader.h>
#include <perspective/base.h>
#include <perspective/binding.h>
#include <perspective/python/accessor.h>
#include <perspective/python/base.h>
#include <perspective/python/fill.h>
#include <perspective/python/numpy.h>
#include <perspective/python/table.h>
#include <perspective/python/utils.h>

namespace perspective {
namespace binding {
    using namespace perspective::apachearrow;

    /******************************************************************************
     *
     * Table API
     */

    std::shared_ptr<Table>
    make_table_py(t_val table, t_data_accessor accessor, std::uint32_t limit,
        std::string index, t_op op, bool is_update, bool is_arrow, bool is_csv,
        t_uindex port_id) {
        bool table_initialized = !table.is_none();
        std::shared_ptr<t_pool> pool;
        std::shared_ptr<Table> tbl;
        std::shared_ptr<t_gnode> gnode;
        std::uint32_t offset;
        void* ptr = nullptr;

        // If the Table has already been created, use it
        if (table_initialized) {
            tbl = table.cast<std::shared_ptr<Table>>();
            pool = tbl->get_pool();
            gnode = tbl->get_gnode();
            offset = tbl->get_offset();
            is_update = (is_update || gnode->mapping_size() > 0);
        } else {
            pool = std::make_shared<t_pool>();
        }

        std::vector<std::string> column_names;
        std::vector<t_dtype> data_types;
        ArrowLoader arrow_loader;
        numpy::NumpyLoader numpy_loader(accessor);

        // don't call `is_numpy` on an arrow binary
        bool is_numpy
            = !is_arrow && !is_csv && accessor.attr("_is_numpy").cast<bool>();

        // Determine metadata
        bool is_delete = op == OP_DELETE;
        if (is_arrow && !is_delete) {
            std::string csv_string;
            std::int32_t binary_size;

            if (is_csv) {
                // Load a string in CSV format
                csv_string = accessor.cast<std::string>();
            } else {
                // Load an arrow binary
                py::bytes bytes = accessor.cast<py::bytes>();
                binary_size = bytes.attr("__len__")().cast<std::int32_t>();
                ptr = malloc(binary_size);
                std::memcpy(
                    ptr, bytes.cast<std::string>().c_str(), binary_size);
            }

            {
                PerspectiveScopedGILRelease acquire(
                    pool->get_event_loop_thread_id());

                // With the GIL released, load the arrow
                if (is_csv) {
                    auto map = std::unordered_map<std::string,
                        std::shared_ptr<arrow::DataType>>();

                    if (is_update) {
                        auto gnode_output_schema = gnode->get_output_schema();
                        auto schema = gnode_output_schema.drop({"psp_okey"});
                        auto column_names = schema.columns();
                        auto data_types = schema.types();

                        for (auto idx = 0; idx < column_names.size(); ++idx) {
                            const std::string& name = column_names[idx];
                            const t_dtype& type = data_types[idx];
                            switch (type) {
                                case DTYPE_FLOAT32:
                                    map[name]
                                        = std::make_shared<arrow::FloatType>();
                                    break;
                                case DTYPE_FLOAT64:
                                    map[name]
                                        = std::make_shared<arrow::DoubleType>();
                                    break;
                                case DTYPE_STR:
                                    map[name]
                                        = std::make_shared<arrow::StringType>();
                                    break;
                                case DTYPE_BOOL:
                                    map[name] = std::make_shared<
                                        arrow::BooleanType>();
                                    break;
                                case DTYPE_UINT32:
                                    map[name]
                                        = std::make_shared<arrow::UInt32Type>();
                                    break;
                                case DTYPE_UINT64:
                                    map[name]
                                        = std::make_shared<arrow::UInt64Type>();
                                    break;
                                case DTYPE_INT32:
                                    map[name]
                                        = std::make_shared<arrow::Int32Type>();
                                    break;
                                case DTYPE_INT64:
                                    map[name]
                                        = std::make_shared<arrow::Int64Type>();
                                    break;
                                case DTYPE_TIME:
                                    map[name] = std::make_shared<
                                        arrow::TimestampType>();
                                    break;
                                case DTYPE_DATE:
                                    map[name]
                                        = std::make_shared<arrow::Date64Type>();
                                    break;
                                default:
                                    std::stringstream ss;
                                    ss << "Error loading arrow type "
                                       << dtype_to_str(type) << " for column "
                                       << name << std::endl;
                                    PSP_COMPLAIN_AND_ABORT(ss.str())
                                    break;
                            }
                        }
                    }

                    arrow_loader.init_csv(csv_string, is_update, map);
                } else {
                    arrow_loader.initialize((uintptr_t)ptr, binary_size);
                }

                // Always use the `Table` column names and data types on update.
                if (table_initialized && is_update) {
                    auto gnode_output_schema = gnode->get_output_schema();
                    auto schema = gnode_output_schema.drop({"psp_okey"});
                    column_names = schema.columns();
                    data_types = schema.types();

                    auto data_table = gnode->get_table();
                    if (data_table->size() == 0) {
                        /**
                         * If updating a table created from schema, a 32-bit
                         * int/float needs to be promoted to a 64-bit int/float
                         * if specified in the Arrow schema.
                         */
                        std::vector<t_dtype> arrow_dtypes
                            = arrow_loader.types();
                        for (auto idx = 0; idx < column_names.size(); ++idx) {
                            const std::string& name = column_names[idx];
                            bool can_retype = name != "psp_okey"
                                && name != "psp_pkey" && name != "psp_op";
                            bool is_32_bit = data_types[idx] == DTYPE_INT32
                                || data_types[idx] == DTYPE_FLOAT32;
                            if (can_retype && is_32_bit) {
                                t_dtype arrow_dtype = arrow_dtypes[idx];
                                switch (arrow_dtype) {
                                    case DTYPE_INT64:
                                    case DTYPE_FLOAT64: {
                                        std::cout
                                            << "Promoting column `"
                                            << column_names[idx]
                                            << "` to maintain consistency with "
                                               "Arrow type."
                                            << std::endl;
                                        gnode->promote_column(
                                            name, arrow_dtype);
                                    } break;
                                    default: {
                                        continue;
                                    }
                                }
                            }
                        }
                    }
                    // Make sure promoted types are used to construct data table
                    auto new_schema
                        = gnode->get_output_schema().drop({"psp_okey"});
                    data_types = new_schema.types();
                } else {
                    column_names = arrow_loader.names();
                    data_types = arrow_loader.types();
                }
            }
        } else if (is_update || is_delete) {
            /**
             * Use the names and types of the python accessor when
             * updating/deleting.
             *
             * This prevents the Table from looking up new columns present in an
             * update.
             *
             * Example: updating a Table with a DataFrame attempts to write the
             * "index" column, but if the table was not created from a
             * DataFrame, the "index" column would not exist.
             */
            if (is_numpy) {
                // `numpy_loader`s `m_names` and `m_types` variable contains
                // only the column names and data types present in the update
                // dataset, not the names/types of the entire `Table`.
                numpy_loader.init();
            }

            // `column_names` and `data_types` contain every single column in
            // the dataset, as well as `__INDEX__` if it exists.
            column_names
                = accessor.attr("names")().cast<std::vector<std::string>>();
            data_types = accessor.attr("types")().cast<std::vector<t_dtype>>();
        } else if (is_numpy) {
            /**
             * Numpy loading depends on both the `dtype` of the individual
             * arrays as well as the inferred type from Perspective. Using
             * `get_data_types` allows us to know the type of an array with
             * `dtype=object`.
             */
            numpy_loader.init();

            // This will contain every single column in the dataset, as the
            // first-time data load path does not mutate the `names` property of
            // `accessor`.
            column_names = numpy_loader.names();

            // Infer data type for each column, and then use a composite of
            // numpy dtype, inferred `t_dtype`, and stringified numpy dtype to
            // get the final, canonical data type mapping.
            std::vector<t_dtype> inferred_types
                = get_data_types(accessor.attr("data")(), 1, column_names,
                    accessor.attr("date_validator")().cast<t_val>());
            data_types = numpy_loader.reconcile_dtypes(inferred_types);
        } else {
            // Infer names and types
            t_val data = accessor.attr("data")();
            std::int32_t format
                = accessor.attr("format")().cast<std::int32_t>();
            column_names = get_column_names(data, format);
            data_types = get_data_types(data, format, column_names,
                accessor.attr("date_validator")().cast<t_val>());
        }

        if (!table_initialized) {
            tbl = std::make_shared<Table>(
                pool, column_names, data_types, limit, index);
            offset = 0;
        }

        // Create input schema - an input schema contains all columns to be
        // displayed AND index + operation columns
        t_schema input_schema(column_names, data_types);

        // strip implicit index, if present
        auto implicit_index_it
            = std::find(column_names.begin(), column_names.end(), "__INDEX__");
        if (implicit_index_it != column_names.end()) {
            auto idx = std::distance(column_names.begin(), implicit_index_it);
            // position of the column is at the same index in both vectors
            column_names.erase(column_names.begin() + idx);
            data_types.erase(data_types.begin() + idx);
        }

        // Create output schema - contains only columns to be displayed to the
        // user
        t_schema output_schema(column_names,
            data_types); // names + types might have been mutated at this point
                         // after implicit index removal
        t_data_table data_table(output_schema);
        data_table.init();
        std::uint32_t row_count;

        if (is_arrow) {
            PerspectiveScopedGILRelease acquire(
                pool->get_event_loop_thread_id());
            row_count = arrow_loader.row_count();
            data_table.extend(row_count);
            arrow_loader.fill_table(
                data_table, input_schema, index, offset, limit, is_update);
        } else if (is_numpy) {
            row_count = numpy_loader.row_count();
            data_table.extend(row_count);
            numpy_loader.fill_table(
                data_table, input_schema, index, offset, limit, is_update);
        } else {
            row_count = accessor.attr("row_count")().cast<std::int32_t>();
            data_table.extend(row_count);
            _fill_data(data_table, accessor, input_schema, index, offset, limit,
                is_update);
        }

        if (is_arrow && !is_csv) {
            free(ptr);
        }

        // calculate offset, limit, and set the gnode
        tbl->init(data_table, row_count, op, port_id);

        // pool->_process();
        return tbl;
    }

} // namespace binding
} // namespace perspective

#endif