#include "nlohmann_reference_parser.hpp"
#include <simdutf.h>

NlohmannReferenceParser::ParseResult
NlohmannReferenceParser::parse(CommitRequest &request,
                               const std::string &json_str) {
  error_message_.clear();
  request.reset();

  try {
    nlohmann::json j = nlohmann::json::parse(json_str);

    // Parse required fields
    if (!j.contains("leader_id") || !j["leader_id"].is_string()) {
      error_message_ = "Missing or invalid leader_id";
      return ParseResult::ValidationError;
    }

    std::string leader_id = j["leader_id"];
    if (leader_id.empty()) {
      error_message_ = "Empty leader_id";
      return ParseResult::ValidationError;
    }
    request.set_leader_id(request.copy_to_arena(leader_id));

    if (!j.contains("read_version") ||
        !j["read_version"].is_number_unsigned()) {
      error_message_ = "Missing or invalid read_version";
      return ParseResult::ValidationError;
    }
    request.set_read_version(j["read_version"]);

    // Parse optional request_id
    if (j.contains("request_id")) {
      if (!j["request_id"].is_string()) {
        error_message_ = "Invalid request_id type";
        return ParseResult::ValidationError;
      }
      std::string request_id = j["request_id"];
      request.set_request_id(request.copy_to_arena(request_id));
    }

    // Parse optional preconditions
    if (j.contains("preconditions")) {
      if (!j["preconditions"].is_array()) {
        error_message_ = "Preconditions must be an array";
        return ParseResult::ValidationError;
      }
      if (!parse_preconditions(request, j["preconditions"])) {
        return ParseResult::ValidationError;
      }
    }

    // Parse optional operations
    if (j.contains("operations")) {
      if (!j["operations"].is_array()) {
        error_message_ = "Operations must be an array";
        return ParseResult::ValidationError;
      }
      if (!parse_operations(request, j["operations"])) {
        return ParseResult::ValidationError;
      }
    }

    request.finalize();
    return ParseResult::Success;

  } catch (const nlohmann::json::parse_error &e) {
    error_message_ = "JSON parse error: " + std::string(e.what());
    return ParseResult::ParseError;
  } catch (const nlohmann::json::exception &e) {
    error_message_ = "JSON error: " + std::string(e.what());
    return ParseResult::ParseError;
  }
}

bool NlohmannReferenceParser::parse_preconditions(
    CommitRequest &request, const nlohmann::json &preconditions_array) {
  for (const auto &precondition : preconditions_array) {
    if (!precondition.is_object()) {
      error_message_ = "Precondition must be an object";
      return false;
    }

    // Parse type
    if (!precondition.contains("type") || !precondition["type"].is_string()) {
      error_message_ = "Precondition missing type";
      return false;
    }

    std::string type_str = precondition["type"];
    Precondition::Type type = parse_precondition_type(type_str);
    if (type_str != "point_read" && type_str != "range_read") {
      error_message_ = "Invalid precondition type: " + type_str;
      return false;
    }

    // Parse version (optional, defaults to 0)
    uint64_t version = 0;
    if (precondition.contains("version")) {
      if (!precondition["version"].is_number_unsigned()) {
        error_message_ = "Invalid precondition version";
        return false;
      }
      version = precondition["version"];
    }

    if (type == Precondition::Type::PointRead) {
      // Point read requires key field
      if (!precondition.contains("key") || !precondition["key"].is_string()) {
        error_message_ = "Point read precondition missing key";
        return false;
      }

      std::string key_b64 = precondition["key"];
      std::string key_decoded = decode_base64(key_b64);

      request.add_precondition(type, version,
                               request.copy_to_arena(key_decoded));
    } else { // RangeRead
      // Range read requires begin and end fields
      if (!precondition.contains("begin") ||
          !precondition["begin"].is_string()) {
        error_message_ = "Range read precondition missing begin";
        return false;
      }
      if (!precondition.contains("end") || !precondition["end"].is_string()) {
        error_message_ = "Range read precondition missing end";
        return false;
      }

      std::string begin_b64 = precondition["begin"];
      std::string end_b64 = precondition["end"];
      std::string begin_decoded = decode_base64(begin_b64);
      std::string end_decoded = decode_base64(end_b64);

      request.add_precondition(type, version,
                               request.copy_to_arena(begin_decoded),
                               request.copy_to_arena(end_decoded));
    }
  }
  return true;
}

bool NlohmannReferenceParser::parse_operations(
    CommitRequest &request, const nlohmann::json &operations_array) {
  for (const auto &operation : operations_array) {
    if (!operation.is_object()) {
      error_message_ = "Operation must be an object";
      return false;
    }

    // Parse type
    if (!operation.contains("type") || !operation["type"].is_string()) {
      error_message_ = "Operation missing type";
      return false;
    }

    std::string type_str = operation["type"];
    Operation::Type type = parse_operation_type(type_str);
    if (type_str != "write" && type_str != "delete" &&
        type_str != "range_delete") {
      error_message_ = "Invalid operation type: " + type_str;
      return false;
    }

    if (type == Operation::Type::Write) {
      // Write requires key and value
      if (!operation.contains("key") || !operation["key"].is_string()) {
        error_message_ = "Write operation missing key";
        return false;
      }
      if (!operation.contains("value") || !operation["value"].is_string()) {
        error_message_ = "Write operation missing value";
        return false;
      }

      std::string key_b64 = operation["key"];
      std::string value_b64 = operation["value"];
      std::string key_decoded = decode_base64(key_b64);
      std::string value_decoded = decode_base64(value_b64);

      request.add_operation(type, request.copy_to_arena(key_decoded),
                            request.copy_to_arena(value_decoded));
    } else if (type == Operation::Type::Delete) {
      // Delete requires key
      if (!operation.contains("key") || !operation["key"].is_string()) {
        error_message_ = "Delete operation missing key";
        return false;
      }

      std::string key_b64 = operation["key"];
      std::string key_decoded = decode_base64(key_b64);

      request.add_operation(type, request.copy_to_arena(key_decoded));
    } else { // RangeDelete
      // Range delete requires begin and end
      if (!operation.contains("begin") || !operation["begin"].is_string()) {
        error_message_ = "Range delete operation missing begin";
        return false;
      }
      if (!operation.contains("end") || !operation["end"].is_string()) {
        error_message_ = "Range delete operation missing end";
        return false;
      }

      std::string begin_b64 = operation["begin"];
      std::string end_b64 = operation["end"];
      std::string begin_decoded = decode_base64(begin_b64);
      std::string end_decoded = decode_base64(end_b64);

      request.add_operation(type, request.copy_to_arena(begin_decoded),
                            request.copy_to_arena(end_decoded));
    }
  }
  return true;
}

std::string
NlohmannReferenceParser::decode_base64(const std::string &base64_str) {
  if (base64_str.empty()) {
    return "";
  }

  // Calculate required output size for worst case
  size_t max_output_size = simdutf::maximal_binary_length_from_base64(
      base64_str.data(), base64_str.size());

  // Allocate output buffer
  std::vector<char> output(max_output_size);

  // Decode using simdutf
  simdutf::result result =
      simdutf::base64_to_binary(base64_str.data(), base64_str.size(),
                                output.data(), simdutf::base64_default);

  if (result.error == simdutf::error_code::SUCCESS) {
    return std::string(output.data(), result.count);
  } else {
    // Return original string if decode fails (for non-base64 strings)
    return base64_str;
  }
}

Precondition::Type
NlohmannReferenceParser::parse_precondition_type(const std::string &type_str) {
  if (type_str == "point_read") {
    return Precondition::Type::PointRead;
  } else if (type_str == "range_read") {
    return Precondition::Type::RangeRead;
  }
  // Default fallback (should not happen if validation is correct)
  return Precondition::Type::PointRead;
}

Operation::Type
NlohmannReferenceParser::parse_operation_type(const std::string &type_str) {
  if (type_str == "write") {
    return Operation::Type::Write;
  } else if (type_str == "delete") {
    return Operation::Type::Delete;
  } else if (type_str == "range_delete") {
    return Operation::Type::RangeDelete;
  }
  // Default fallback (should not happen if validation is correct)
  return Operation::Type::Write;
}