// 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 "adbc_driver_manager.h"
#include <adbc.h>

#include <algorithm>
#include <cstring>
#include <string>
#include <unordered_map>
#include <utility>

#if defined(_WIN32)
#include <windows.h>  // Must come first

#include <libloaderapi.h>
#include <strsafe.h>
#else
#include <dlfcn.h>
#endif  // defined(_WIN32)

namespace {

// Platform-specific helpers

#if defined(_WIN32)
/// Append a description of the Windows error to the buffer.
void GetWinError(std::string* buffer) {
  DWORD rc = GetLastError();
  LPVOID message;

  FormatMessage(FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM |
                    FORMAT_MESSAGE_IGNORE_INSERTS,
                /*lpSource=*/nullptr, rc, MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
                reinterpret_cast<LPSTR>(&message), /*nSize=*/0, /*Arguments=*/nullptr);

  (*buffer) += '(';
  (*buffer) += std::to_string(rc);
  (*buffer) += ") ";
  (*buffer) += reinterpret_cast<char*>(message);
  LocalFree(message);
}

#endif  // defined(_WIN32)

// Error handling

void ReleaseError(struct AdbcError* error) {
  if (error) {
    if (error->message) delete[] error->message;
    error->message = nullptr;
    error->release = nullptr;
  }
}

void SetError(struct AdbcError* error, const std::string& message) {
  if (!error) return;
  if (error->message) {
    // Append
    std::string buffer = error->message;
    buffer.reserve(buffer.size() + message.size() + 1);
    buffer += '\n';
    buffer += message;
    error->release(error);

    error->message = new char[buffer.size() + 1];
    buffer.copy(error->message, buffer.size());
    error->message[buffer.size()] = '\0';
  } else {
    error->message = new char[message.size() + 1];
    message.copy(error->message, message.size());
    error->message[message.size()] = '\0';
  }
  error->release = ReleaseError;
}

// Driver state

/// Hold the driver DLL and the driver release callback in the driver struct.
struct ManagerDriverState {
  // The original release callback
  AdbcStatusCode (*driver_release)(struct AdbcDriver* driver, struct AdbcError* error);

#if defined(_WIN32)
  // The loaded DLL
  HMODULE handle;
#endif  // defined(_WIN32)
};

/// Unload the driver DLL.
static AdbcStatusCode ReleaseDriver(struct AdbcDriver* driver, struct AdbcError* error) {
  AdbcStatusCode status = ADBC_STATUS_OK;

  if (!driver->private_manager) return status;
  ManagerDriverState* state =
      reinterpret_cast<ManagerDriverState*>(driver->private_manager);

  if (state->driver_release) {
    status = state->driver_release(driver, error);
  }

#if defined(_WIN32)
  // TODO(apache/arrow-adbc#204): causes tests to segfault
  // if (!FreeLibrary(state->handle)) {
  //   std::string message = "FreeLibrary() failed: ";
  //   GetWinError(&message);
  //   SetError(error, message);
  // }
#endif  // defined(_WIN32)

  driver->private_manager = nullptr;
  delete state;
  return status;
}

// Default stubs

AdbcStatusCode DatabaseSetOption(struct AdbcDatabase* database, const char* key,
                                 const char* value, struct AdbcError* error) {
  return ADBC_STATUS_NOT_IMPLEMENTED;
}

AdbcStatusCode ConnectionCommit(struct AdbcConnection*, struct AdbcError* error) {
  return ADBC_STATUS_NOT_IMPLEMENTED;
}

AdbcStatusCode ConnectionGetInfo(struct AdbcConnection* connection, uint32_t* info_codes,
                                 size_t info_codes_length, struct ArrowArrayStream* out,
                                 struct AdbcError* error) {
  return ADBC_STATUS_NOT_IMPLEMENTED;
}

AdbcStatusCode ConnectionGetObjects(struct AdbcConnection*, int, const char*, const char*,
                                    const char*, const char**, const char*,
                                    struct ArrowArrayStream*, struct AdbcError* error) {
  return ADBC_STATUS_NOT_IMPLEMENTED;
}

AdbcStatusCode ConnectionGetTableSchema(struct AdbcConnection*, const char*, const char*,
                                        const char*, struct ArrowSchema*,
                                        struct AdbcError* error) {
  return ADBC_STATUS_NOT_IMPLEMENTED;
}

AdbcStatusCode ConnectionGetTableTypes(struct AdbcConnection*, struct ArrowArrayStream*,
                                       struct AdbcError* error) {
  return ADBC_STATUS_NOT_IMPLEMENTED;
}

AdbcStatusCode ConnectionReadPartition(struct AdbcConnection* connection,
                                       const uint8_t* serialized_partition,
                                       size_t serialized_length,
                                       struct ArrowArrayStream* out,
                                       struct AdbcError* error) {
  return ADBC_STATUS_NOT_IMPLEMENTED;
}

AdbcStatusCode ConnectionRollback(struct AdbcConnection*, struct AdbcError* error) {
  return ADBC_STATUS_NOT_IMPLEMENTED;
}

AdbcStatusCode ConnectionSetOption(struct AdbcConnection*, const char*, const char*,
                                   struct AdbcError* error) {
  return ADBC_STATUS_NOT_IMPLEMENTED;
}

AdbcStatusCode StatementBind(struct AdbcStatement*, struct ArrowArray*,
                             struct ArrowSchema*, struct AdbcError* error) {
  return ADBC_STATUS_NOT_IMPLEMENTED;
}

AdbcStatusCode StatementExecutePartitions(struct AdbcStatement* statement,
                                          struct ArrowSchema* schema,
                                          struct AdbcPartitions* partitions,
                                          int64_t* rows_affected,
                                          struct AdbcError* error) {
  return ADBC_STATUS_NOT_IMPLEMENTED;
}

AdbcStatusCode StatementGetParameterSchema(struct AdbcStatement* statement,
                                           struct ArrowSchema* schema,
                                           struct AdbcError* error) {
  return ADBC_STATUS_NOT_IMPLEMENTED;
}

AdbcStatusCode StatementPrepare(struct AdbcStatement*, struct AdbcError* error) {
  return ADBC_STATUS_NOT_IMPLEMENTED;
}

AdbcStatusCode StatementSetOption(struct AdbcStatement*, const char*, const char*,
                                  struct AdbcError* error) {
  return ADBC_STATUS_NOT_IMPLEMENTED;
}

AdbcStatusCode StatementSetSqlQuery(struct AdbcStatement*, const char*,
                                    struct AdbcError* error) {
  return ADBC_STATUS_NOT_IMPLEMENTED;
}

AdbcStatusCode StatementSetSubstraitPlan(struct AdbcStatement*, const uint8_t*, size_t,
                                         struct AdbcError* error) {
  return ADBC_STATUS_NOT_IMPLEMENTED;
}

/// Temporary state while the database is being configured.
struct TempDatabase {
  std::unordered_map<std::string, std::string> options;
  std::string driver;
  // Default name (see adbc.h)
  std::string entrypoint = "AdbcDriverInit";
  AdbcDriverInitFunc init_func = nullptr;
};

/// Temporary state while the database is being configured.
struct TempConnection {
  std::unordered_map<std::string, std::string> options;
};
}  // namespace

// Direct implementations of API methods

AdbcStatusCode AdbcDatabaseNew(struct AdbcDatabase* database, struct AdbcError* error) {
  // Allocate a temporary structure to store options pre-Init
  database->private_data = new TempDatabase();
  database->private_driver = nullptr;
  return ADBC_STATUS_OK;
}

AdbcStatusCode AdbcDatabaseSetOption(struct AdbcDatabase* database, const char* key,
                                     const char* value, struct AdbcError* error) {
  if (database->private_driver) {
    return database->private_driver->DatabaseSetOption(database, key, value, error);
  }

  TempDatabase* args = reinterpret_cast<TempDatabase*>(database->private_data);
  if (std::strcmp(key, "driver") == 0) {
    args->driver = value;
  } else if (std::strcmp(key, "entrypoint") == 0) {
    args->entrypoint = value;
  } else {
    args->options[key] = value;
  }
  return ADBC_STATUS_OK;
}

AdbcStatusCode AdbcDriverManagerDatabaseSetInitFunc(struct AdbcDatabase* database,
                                                    AdbcDriverInitFunc init_func,
                                                    struct AdbcError* error) {
  if (database->private_driver) {
    return ADBC_STATUS_INVALID_STATE;
  }

  TempDatabase* args = reinterpret_cast<TempDatabase*>(database->private_data);
  args->init_func = init_func;
  return ADBC_STATUS_OK;
}

AdbcStatusCode AdbcDatabaseInit(struct AdbcDatabase* database, struct AdbcError* error) {
  if (!database->private_data) {
    SetError(error, "Must call AdbcDatabaseNew first");
    return ADBC_STATUS_INVALID_STATE;
  }
  TempDatabase* args = reinterpret_cast<TempDatabase*>(database->private_data);
  if (args->init_func) {
    // Do nothing
  } else if (args->driver.empty()) {
    SetError(error, "Must provide 'driver' parameter");
    return ADBC_STATUS_INVALID_ARGUMENT;
  }

  database->private_driver = new AdbcDriver;
  std::memset(database->private_driver, 0, sizeof(AdbcDriver));
  AdbcStatusCode status;
  // So we don't confuse a driver into thinking it's initialized already
  database->private_data = nullptr;
  if (args->init_func) {
    status = AdbcLoadDriverFromInitFunc(args->init_func, ADBC_VERSION_1_0_0,
                                        database->private_driver, error);
  } else {
    status = AdbcLoadDriver(args->driver.c_str(), args->entrypoint.c_str(),
                            ADBC_VERSION_1_0_0, database->private_driver, error);
  }
  if (status != ADBC_STATUS_OK) {
    // Restore private_data so it will be released by AdbcDatabaseRelease
    database->private_data = args;
    if (database->private_driver->release) {
      database->private_driver->release(database->private_driver, error);
    }
    delete database->private_driver;
    database->private_driver = nullptr;
    return status;
  }
  status = database->private_driver->DatabaseNew(database, error);
  if (status != ADBC_STATUS_OK) {
    if (database->private_driver->release) {
      database->private_driver->release(database->private_driver, error);
    }
    delete database->private_driver;
    database->private_driver = nullptr;
    return status;
  }
  for (const auto& option : args->options) {
    status = database->private_driver->DatabaseSetOption(database, option.first.c_str(),
                                                         option.second.c_str(), error);
    if (status != ADBC_STATUS_OK) {
      delete args;
      // Release the database
      std::ignore = database->private_driver->DatabaseRelease(database, error);
      if (database->private_driver->release) {
        database->private_driver->release(database->private_driver, error);
      }
      delete database->private_driver;
      database->private_driver = nullptr;
      // Should be redundant, but ensure that AdbcDatabaseRelease
      // below doesn't think that it contains a TempDatabase
      database->private_data = nullptr;
      return status;
    }
  }
  delete args;
  return database->private_driver->DatabaseInit(database, error);
}

AdbcStatusCode AdbcDatabaseRelease(struct AdbcDatabase* database,
                                   struct AdbcError* error) {
  if (!database->private_driver) {
    if (database->private_data) {
      TempDatabase* args = reinterpret_cast<TempDatabase*>(database->private_data);
      delete args;
      database->private_data = nullptr;
      return ADBC_STATUS_OK;