#include "CodeGenerator.h"
#include "InterfaceDescription.h"
#include "version.h"

#include <algorithm>
#include <exception>
#include <fstream>
#include <sstream>
#include <string>
#include <iomanip>
#include <iostream>

/**
 * Maps the names of lowercased IDL names to their size on the wire, in bytes. For the most part
 * this corresponds to the size of the types in C++ land.
 *
 * All binary types are represented as a pointer into the "payload" of the packet (the area after
 * all scalar types); each pointer takes up 8 bytes, a 4 byte offset and a 4 byte length.
 */
const std::unordered_map<std::string, size_t> CodeGenerator::gWireSizes{
    {"bool",   sizeof(bool)},
    {"int8",   sizeof(int8_t)},  {"int16",  sizeof(int16_t)}, 
    {"int32",  sizeof(int32_t)}, {"int64",  sizeof(int64_t)},
    {"uint8",  sizeof(uint8_t)}, {"uint16", sizeof(uint16_t)},
    {"uint32", sizeof(uint32_t)},{"uint64", sizeof(uint64_t)},
    {"float32",sizeof(float)},   {"float64",sizeof(double)},

    // pointer based types
    {"string", 2*sizeof(uint32_t)},  {"blob", 2*sizeof(uint32_t)},

    // void types aren't actually encoded so they take up no space
    {"void", 0},
};

/// Maps the given IDL type to whether it is blob based (and encoded by a pointer) or not
static const std::unordered_map<std::string, bool> gWireIsBlob{
    {"bool",   false},
    {"int8",   false}, {"int16",  false}, 
    {"int32",  false}, {"int64",  false},
    {"uint8",  false}, {"uint16", false},
    {"uint32", false}, {"uint64", false},
    {"float32",false}, {"float64",false},
    {"string", true},  {"blob",   true},
    {"void",   false},
};

/**
 * Convert string to lowercase
 */
static inline std::string LowerString(const std::string &in) {
    auto lower = in;
    std::transform(lower.begin(), lower.end(), lower.begin(), ::tolower);
    return lower;
}

/**
 * Is the given argument a blob type?
 */
static inline bool IsBlob(const InterfaceDescription::Argument &a) {
    return a.isBuiltinType() ? gWireIsBlob.at(LowerString(a.getTypeName())) : true;
}

/**
 * Returns the name of the given interface's identifier in the Type enum.
 *
 * @param fullyQualified When set, the namespace is prepended
 */
std::string CodeGenerator::SerGetMessageIdEnumName(const InterfaceDescription::Method &m,
        const bool fullyQualified) {
    std::string temp;

    if(fullyQualified) {
        temp = "internals::Type::";
    }

    auto off = temp.size();
    temp += m.getName();
    temp[off] = std::toupper(temp[off]);
    return temp;
}
/**
 * Returns the name of the given interface converted to the name of the request or response
 * structure.
 */
std::string CodeGenerator::SerGetMessageStructName(const InterfaceDescription::Method &m,
        const bool isResponse) {
    auto temp = m.getName();
    temp[0] = std::toupper(temp[0]);
    return temp + (isResponse ? "Response" : "Request");
}



/**
 * Generates the serialization header.
 */
void CodeGenerator::generateSerialization() {
    auto fileNameH = this->outDir / ("RpcHelpers_" + this->interface->getName() + ".hpp");
    this->serializationFile = fileNameH;

    std::cout << "    * Serialization: " << fileNameH.string() << ", " << std::endl;

    // generate the header first
    std::string includeGuardName("RPC_HELPERS_GENERATED_");
    includeGuardName.append(std::to_string(this->interface->getIdentifier()));

    std::ofstream header(fileNameH.string(), std::ofstream::trunc);
    this->serWriteInfoBlock(header);
    header << R"(#ifndef )" << includeGuardName << R"(
#define )" << includeGuardName << R"(

#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <cstdio>
#include <cstring>
#include <string>
#include <span>
#include <string_view>

#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Wunused-variable"
)" << std::endl;

    this->cppWriteIncludes(header);
    this->serWriteStructs(header);

    header << R"(
#pragma clang diagnostic push

#endif // defined()" << includeGuardName << R"()
)";
}

/**
 * Write the generic info block comment for the client stub.
 */
void CodeGenerator::serWriteInfoBlock(std::ofstream &os) {
    os << R"(/*
 * This RPC serialization code was autogenerated by idlc (version )"
    << std::string(gVERSION_HASH).substr(0, 8) << "). DO NOT EDIT!" << std::endl
    << " * Generated from " << this->interface->getSourceFilename() << " for interface "
       << this->interface->getName() << " at " << this->creationTimestamp << std::endl
       << " *" << R"(
 * The structs and methods within are used by the RPC system to serialize and deserialize the
 * arguments and return values on method calls. They work internally in the same way that encoding
 * custom types in RPC messages works.
 *
 * See the full RPC documentation for more details.
)"
       << " */" << std::endl;
}


/**
 * Writes out definitions for each of the messages' request and reply structs.
 *
 * This includes the sizes and offsets of each of the fields within it in the wire format.
 */
void CodeGenerator::serWriteStructs(std::ofstream &os) {
    // output info for each method
    os << R"(
namespace rpc {
/**
 * Method to handle a failure to deserialize a field; this will log the failure if built in debug
 * mode.
 */
static inline void HandleDecodeError(const char *typeName, const char *fieldName,
    const uintptr_t offset) {
    fprintf(stderr, "[RPC] Decode error for type %s, field %s at offset $%x\n", typeName, fieldName,
        offset);
}
static inline void HandleDecodeError(const char *typeName, const char *fieldName,
    const uintptr_t offset, const uint32_t blobDataOffset, const uint32_t blobSz) {
    fprintf(stderr, "[RPC] Decode error for type %s, field %s at offset $%x "
        "(blob offset $%x, $%x bytes)\n", typeName, fieldName, offset, blobDataOffset, blobSz);
}

/*
 * Serialization functions for various C++ built in types
 */
inline size_t bytesFor(const std::string &s) {
    return s.length();
}
inline bool serialize(std::span<std::byte> &out, const std::string &str) {
    if(str.empty()) return true;
    else if(out.size() < str.length()) return false;
    memcpy(out.data(), str.c_str(), str.length());
    return true;
}
inline bool deserialize(const std::span<std::byte> &in, std::string &outStr) {
    if(in.empty()) {
        outStr = "";
    } else {
        outStr = std::string(reinterpret_cast<const char *>(in.data()), in.size());
    }
    return true;
}

// XXX: this only works for POD types!
template<typename T>
inline size_t bytesFor(const std::vector<T> &s) {
    return s.size() * sizeof(T);
}
template<typename T>
inline bool serialize(std::span<std::byte> &out, const std::vector<T> &vec) {
    const size_t numBytes = vec.size() * sizeof(T);
    if(out.size() < numBytes) return false;
    memcpy(out.data(), vec.data(), out.size());
    return true;
}
template <typename T>
inline bool deserialize(const std::span<std::byte> &in, std::vector<T> &outVec) {
    const size_t elements = in.size() / sizeof(T);
    outVec.resize(elements);
    memcpy(outVec.data(), in.data(), in.size());
    return true;
}


/*
 * Definitions of serialization structures for messages and message replies. These use the
 * autogenerated stubs to convert to/from the wire format.
 */
namespace internals {
)";

    // write message types
    os << R"(/// Message ids for each of the RPC messages
enum class Type: uint64_t {
)";
    for(const auto &m : this->interface->getMethods()) {
        os << "    " << std::setw(48) << SerGetMessageIdEnumName(m, false) << " = 0x" << std::hex
           << m.getIdentifier() << std::dec << "ULL,\n";
    }
    os << "};" << std::endl;

    // then the structures
    for(const auto &method : this->interface->getMethods()) {
        this->serWriteMethod(os, method);
    }

os << R"(} // namespace rpc::internals

)";

    // write the associated serialization methods
    for(const auto &method : this->interface->getMethods()) {
        const auto requestName = SerGetMessageStructName(method, false);
        this->serWriteSerializers(os, method.getParameters(), requestName);

        if(!method.isAsync()) {
            const auto responseName = SerGetMessageStructName(method, true);
            this->serWriteSerializers(os, method.getReturns(), responseName);
        }
    }

    os << std::endl << "}; // namespace rpc" << std::endl;
}

/**
 * Writes out the request and optional reply structures for the given method.
 *
 * Note that we still generate structs even if the request/reply has no arguments; these will be
 * empty and it's the compiler's job to optimize calls to them away, since the serializers end up
 * as no-ops.
 */
void CodeGenerator::serWriteMethod(std::ofstream &os, const Method &m) {
    // start with the request structure
    os << R"(/**
 * Request structure for method ')" << m.getName() << R"('
 */
struct )" << SerGetMessageStructName(m, false) << " {" << std::endl;
    this->serWriteArgs(os, m.getParameters());
    os << "};" << std::endl;

    // if method is not async, write its reply structure
    if(!m.isAsync()) {
        const auto name = SerGetMessageStructName(m, true);

        os << R"(/**
 * Reply structure for method ')" << m.getName() << R"('
 */
struct )" << name << " {" << std::endl;
        this->serWriteArgs(os, m.getReturns());
        os << "};" << std::endl;
    }

    os << std::endl;
}

/**
 * Writes the arguments for the method out into the struct's argument form.
 *
 * Note that arguments are serialized in the order that they're specified in this arguments vector,
 * which is the declaration order.
 */
void CodeGenerator::serWriteArgs(std::ofstream &os, const std::vector<Argument> &args) {
    // write the actual storage
    for(const auto &a : args) {
        os << "    " << CppTypenameForArg(a, false) << " " << a.getName() << ";" << std::endl;
    }
    os << std::endl;

    /*
     * Write an array that contains the number of bytes that each element takes up on the wire.
     * This is followed by one that contains the offsets of each element from the start of the
     * message.
     */
    std::stringstream sizes, offsets;
    size_t scalarBytes{0};

    for(size_t i = 0; i < args.size(); i++) {
        const auto &a = args[i];
        size_t bytes;

        if(a.isBuiltinType()) {
            bytes = gWireSizes.at(LowerString(a.getTypeName()));
        } else {
            // non-builtin types are encoded as blobs
            bytes = 2*sizeof(uint32_t);
        }

        sizes << std::setw(2) << bytes;
        offsets << std::setw(2) << scalarBytes;

        scalarBytes += bytes;

        if(i != (args.size() - 1)) {
            sizes << ", ";
            offsets << ", ";
        }
    }

    os << "    constexpr static const size_t kElementSizes[" << args.size() << "] {" << std::endl
       << "    " << sizes.str() << std::endl << "    };" << std::endl
       << "    constexpr static const size_t kElementOffsets[" << args.size() << "] {" << std::endl
       << "    " << offsets.str() << std::endl << "    };" << std::endl;

    os << "    constexpr static const size_t kScalarBytes{" << scalarBytes << "};" << std::endl;

    auto blobStart = ((scalarBytes + 7) / 8) * 8;
    os << "    constexpr static const size_t kBlobStartOffset{" << blobStart << "};" << std::endl;
}

/**
 * Writes the serialization methods for the message structure with the given arguments and type
 * name.
 */
void CodeGenerator::serWriteSerializers(std::ofstream &os, const std::vector<Argument> &args,
        const std::string &structName) {
    os << std::endl;

    /**
     * Get the total number of bytes required for the message. This consists of the scalar bytes,
     * rounded up to the nearest 8 byte boundary. Then, it is followed by each of the blob datas
     * packed immediately after one another.
     *
     * The basic takeaway is that nothing is ever guaranteed to be aligned :D
     */
os << R"(inline size_t bytesFor(const internals::)" << structName << R"( &x) {
    using namespace internals;
    size_t len = )" << structName << R"(::kBlobStartOffset;
)";

    for(const auto &a : args) {
        // figure out if this is a blob
        if(!IsBlob(a)) continue;
        os << "    len += bytesFor(x." << a.getName() << ");" << std::endl;
    }

os << R"(
    return len;
}
)";

    /**
     * Serializes a message; this generates code for each of the arguments in order, using the
     * indices provided in the message's arrays.
     */
    os << R"(inline bool serialize(std::span<std::byte> &out, const internals::)" << structName << R"( &x) {
    using namespace internals;
    uint32_t blobOff = )" << structName << R"(::kBlobStartOffset;
)";

    for(size_t i = 0; i < args.size(); i++) {
        const auto &a = args[i];
        os << "    {" << std::endl;

        // get offset and number of allocated bytes
        os << "        const auto off = " << structName << "::kElementOffsets[" << i << "];\n"
           << "        const auto size = " << structName << "::kElementSizes[" << i << "];\n"
           << "        auto range = out.subspan(off, size);\n";

        // allocate blob storage and write pointer if needed
        if(IsBlob(a)) {
            // serialize blob data
            os << "        const uint32_t blobSz = bytesFor(x." << a.getName() << ");\n"
               << "        const uint32_t blobDataOffset = blobOff;\n"
               << "        auto blobRange = out.subspan(blobDataOffset, blobSz);\n"
               << "        if(!serialize(blobRange, x." << a.getName() << ")) return false;\n"
               << "        blobOff += blobSz;\n";

            // write blob pointers
            os << "        memcpy(range.data(), &blobDataOffset, sizeof(blobDataOffset));\n"
               << "        memcpy(range.data()+sizeof(blobDataOffset), &blobSz, sizeof(blobSz));\n";
        }
        // otherwise, it's a primitive type we can write directly
        else {
            os << "        memcpy(range.data(), &x." << a.getName() << ", range.size());\n";
        }

        os << "    }" << std::endl;
    }

    os << R"(
    return true;
})" << std::endl;

    /**
     * Deserializes a message. Arguments are unpacked in order and those with extra data are
     * decoded from the blob region.
     */
    os << R"(inline bool deserialize(const std::span<std::byte> &in, internals::)" << structName << R"( &x) {
    using namespace internals;
    if(in.size() < )" << structName << R"(::kScalarBytes) return false;
    const auto blobRegion = in.subspan()" << structName << R"(::kBlobStartOffset);
)";

    for(size_t i = 0; i < args.size(); i++) {
        const auto &a = args[i];
        os << "    {" << std::endl;

        // get the offset of the fields and its data in the input buffer
        os << "        const auto off = " << structName << "::kElementOffsets[" << i << "];\n"
           << "        const auto size = " << structName << "::kElementSizes[" << i << "];\n"
           << "        auto range = in.subspan(off, size);\n"
           << "        if(range.empty() || range.size() != size) return false;\n";

        // allocate blob storage and write pointer if needed
        if(IsBlob(a)) {
            // read out blob pointers
            os << "        uint32_t blobSz{0}, blobDataOffset{0};\n"
               << "        memcpy(&blobDataOffset, range.data(), sizeof(blobDataOffset));\n"
               << "        memcpy(&blobSz, range.data()+sizeof(blobDataOffset), sizeof(blobSz));\n";

            // get blob range and deserialize
            os << "        auto blobRange = in.subspan(blobDataOffset, blobSz);\n"
                << "       if(!deserialize(blobRange, x." << a.getName() << ")) {\n"
                << "            HandleDecodeError(\"" << structName << "\", \"" << a.getName()
                << "\", off, blobDataOffset, blobSz);\n"
                << "            return false;\n"
                << "        }\n";
        }
        // otherwise, it's a primitive type we can write directly
        else {
            os << "        memcpy(&x." << a.getName() << ", range.data(), range.size());\n";
        }

        os << "    }" << std::endl;
    }
    os << R"(
    return true;
})" << std::endl;
}

