/* -*- Mode: C++; tab-width: 4; c-basic-offset: 4; indent-tabs-mode: nil -*- */
/*
 *     Copyright 2016 Couchbase, Inc.
 *
 *   Licensed 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 "contrib/lcb-jsoncpp/lcb-jsoncpp.h"
#include "placeholders.h"
#include <algorithm>
#include <stdexcept>

namespace Pillowfight {

#define JSON_VALUE_SIZE 16

/**
 * Per-thread mutable state to generate the document. The populateIov()
 * method is called to populate IOVs suitable for passing to the storage
 * functions.
 */
class GeneratorState {
public:
    /**
     * Populate an IOV array
     * @param seq The sequence number (used as a 'selector')
     * @param[out] iov IOV containing pointers to buffers
     *
     * The buffers pointed to in the IOV array remain valid so long as
     * the populateIov() function is not called again on this state
     */
    virtual void populateIov(uint32_t seq, std::vector<lcb_IOV>& iov) = 0;
    virtual ~GeneratorState() {}
};

class SubdocGeneratorState {
public:
    /**
     * Populates subdocument command specifications
     * @param seq the sequence number of the current command
     * @param[in,out] specs container to hold the actual spec array.
     *  The spec array must have already been properly pre-sized.
     */
    virtual void populateLookup(uint32_t seq, std::vector<lcb_SDSPEC>& specs) = 0;
    virtual void populateMutate(uint32_t seq, std::vector<lcb_SDSPEC>& specs) = 0;
    virtual ~SubdocGeneratorState() {}
};

class DocGeneratorBase {
public:
    /**
     * Create the per-thread state for generating documents
     * @param total_gens Number of total generator threads
     * @param cur_gen The index of the current generator thread
     * @return An opaque state object. This should be deleted by the caller
     */
    virtual GeneratorState *createState(int total_gens, int cur_gen) const = 0;
    virtual SubdocGeneratorState *createSubdocState(int, int) const { return NULL; }
    virtual ~DocGeneratorBase() {}
};

/**
 * Generator class for raw objects. This contains a fixed buffer and will
 * simply vary in how 'long' the buffer is
 */
class RawDocGenerator : public DocGeneratorBase {
public:
    /**
     * Generate graded sizes based on a min,max specification. This allows us
     * to be more efficient by cutting corners on how 'random' the sizes
     * actually are. Rather than generating a 'random' size each time we need
     * a document, we split the range into a set of potential sizes (which are
     * also evenly distributed) and cycle between them.
     *
     * @param minsz Smallest desired size
     * @param maxsz Largest desired size
     * @param granularity How many grades to produce
     * @return A vector of sizes which fall between the range
     */
    static std::vector<size_t>
    gen_graded_sizes(uint32_t minsz, uint32_t maxsz, int grades=10) {
        std::vector<size_t> ret;

        size_t diff = maxsz - minsz;
        size_t factor = diff / grades;
        if (factor == 0 || minsz == maxsz) {
            ret.push_back(maxsz);
        } else {
            for (int ii = 0; ii < grades+1; ii++) {
                size_t size = minsz + (factor * ii);
                ret.push_back(size);
            }
        }
        return ret;
    }

    RawDocGenerator(uint32_t minsz, uint32_t maxsz)
    : m_sizes(gen_graded_sizes(minsz, maxsz)) {
        // Populate the buffer to its capacity
        m_buf.insert(0, maxsz, '#');
    }


    class MyState : public GeneratorState {
    public:
        MyState(const RawDocGenerator *parent) : m_parent(parent) {
        }

        const RawDocGenerator *m_parent;
        void populateIov(uint32_t seq, std::vector<lcb_IOV>& iov_out) {
            m_parent->populateIov(seq, iov_out);
        }
    };

    GeneratorState * createState(int, int) const {
        return new MyState(this);
    }

private:
    void populateIov(uint32_t seq, std::vector<lcb_IOV>& iov_out) const {
        iov_out.resize(1);
        size_t cursz = m_sizes[seq % m_sizes.size()];
        iov_out[0].iov_base = const_cast<char*>(m_buf.c_str());
        iov_out[0].iov_len = cursz;
    }
    std::string m_buf;
    std::vector<size_t> m_sizes;
};

/**
 * This 'generator' ignores sizes and generates documents as they are received
 * from premade buffers
 */
class PresetDocGenerator : public DocGeneratorBase {
public:
    /**
     * @param inputs List of fixed inputs to use
     */
    PresetDocGenerator(const std::vector<std::string>& inputs) : m_bufs(inputs) {
    }

    class MyState : public GeneratorState {
    public:
        MyState(const PresetDocGenerator *parent) : m_parent(parent) {
        }

        void populateIov(uint32_t seq, std::vector<lcb_IOV>& iov_out) {
            m_parent->populateIov(seq, iov_out);
        }

    private:
        const PresetDocGenerator *m_parent;
    };

    GeneratorState *createState(int, int) const {
        return new MyState(this);
    }

protected:
    PresetDocGenerator() {}

    void populateIov(uint32_t seq, std::vector<lcb_IOV>& iov_out) const {
        iov_out.resize(1);
        const std::string& s = m_bufs[seq % m_bufs.size()];
        iov_out[0].iov_base = const_cast<char *>(s.c_str());
        iov_out[0].iov_len = s.size();
    }
    std::vector<std::string> m_bufs;
};

// This is the same as the normal document generator, except we generate
// the JSON first
class JsonDocGenerator : public PresetDocGenerator {
private:
    struct Doc {
        std::string m_doc;
        class Field {
        public:
            Field(const std::string& n, std::string& v) : m_name(n), m_value(v) {}
            const std::string& name() const { return m_name; }
            const std::string& value() const { return m_value; }
        private:
            std::string m_name;
            std::string m_value;
        };
        std::vector<Field> m_fields;
    };
    std::vector<Doc> m_docs;

public:
    /**
     * @param minsz Minimum JSON document size
     * @param maxsz Maximum JSON document size
     */
    JsonDocGenerator(uint32_t minsz, uint32_t maxsz)
    {
        genDocuments(minsz, maxsz, m_docs);
        for (size_t ii = 0; ii < m_docs.size(); ++ii) {
            m_bufs.push_back(m_docs[ii].m_doc);
        }
    }

    static void genDocuments(uint32_t minsz, uint32_t maxsz, std::vector<std::string>& out)
    {
        std::vector<Doc> docs;
        genDocuments(minsz, maxsz, docs);
        for (size_t ii = 0; ii < docs.size(); ++ii) {
            out.push_back(docs[ii].m_doc);
        }
    }

private:
    /**
     * Helper method to decrease the original size by a given amount.
     *
     * This also ensures the number never reaches below 0.
     *
     * @param[out] orig Pointer to original size
     * @param toDecr number by which to decrement
     */
    static void decrSize(int *orig, size_t toDecr)
    {
        *orig = std::max(0, *orig - static_cast<int>(toDecr));
    }

    static void
    genDocuments(uint32_t minsz, uint32_t maxsz, std::vector<Doc>& out)
    {
        std::vector<size_t> sizes = RawDocGenerator::gen_graded_sizes(minsz, maxsz);
        for (std::vector<size_t>::iterator ii = sizes.begin();
                ii != sizes.end(); ++ii) {
            out.push_back(generate(*ii));
        }
    }

    /**
     * Generates a "JSON" document of a given size. In order to remain
     * more or less in-tune with common document sizes, field names will be
     * "Field_$incr" and values will be evenly distributed as fixed 16 byte
     * strings. (See JSON_VALUE_SIZE)
     */
    static Doc generate(int docsize)
    {
        int counter = 0;
        char keybuf[128] = { 0 };
        Json::Value root(Json::objectValue);
        Json::FastWriter writer;
        Doc ret;

        while (docsize > 0) {
            decrSize(&docsize, sprintf(keybuf, "Field_%d", ++counter) + 3);
            size_t valsize = std::min(JSON_VALUE_SIZE, docsize);
            if (!valsize) {
                valsize = 1;
            }
            std::string value(valsize, '*');
            decrSize(&docsize, valsize + 3);
            root[keybuf] = value;
            value = '"' + value;
            value += '"';
            ret.m_fields.push_back(Doc::Field(keybuf, value));
        }
        ret.m_doc = writer.write(root);
        return ret;
    }

    class SDGenstate: public SubdocGeneratorState {
    public:
        SDGenstate(const std::vector<Doc>& docs) : m_pathix(0), m_docs(docs) {
        }

        void populateLookup(uint32_t seq, std::vector<lcb_SDSPEC>& specs) {
            populate(seq, specs, false);
        }
        void populateMutate(uint32_t seq, std::vector<lcb_SDSPEC>& specs) {
            populate(seq, specs, true);
        }

    private:
        void populate(uint32_t seq, std::vector<lcb_SDSPEC>& specs, bool mutate) {
            const Doc& d = doc(seq);
            specs.resize(std::min(d.m_fields.size(), specs.size()));
            for (size_t ii = 0; ii < d.m_fields.size() && ii < specs.size(); ++ii) {
                const Doc::Field& f = d.m_fields[m_pathix++ % d.m_fields.size()];
                lcb_SDSPEC& cur_spec = specs[ii];
                LCB_SDSPEC_SET_PATH(&cur_spec, f.name().c_str(), f.name().size());
                if (mutate) {
                    LCB_SDSPEC_SET_VALUE(&cur_spec, f.value().c_str(), f.value().size());
                    specs[ii].sdcmd = LCB_SDCMD_DICT_UPSERT;
                } else {
                    specs[ii].sdcmd = LCB_SDCMD_GET;
                }
            }
        }

        const Doc& doc(uint32_t seq) const { return m_docs[seq % m_docs.size()]; }
        size_t m_pathix;
        const std::vector<Doc>& m_docs;
    };
public:
    virtual SubdocGeneratorState *createSubdocState(int, int) const {
        return new SDGenstate(m_docs);
    }
};

struct TemplateSpec {
    std::string term;
    unsigned minval;
    unsigned maxval;
    bool sequential;
};

/**
 * Generate documents based on placeholder values. Each document (JSON or not)
 * may have one or more special replacement texts which can be substituted
 * with a random number
 */
class PlaceholderDocGenerator : public DocGeneratorBase {
public:
    /**
     * @param specs Placeholder specifications
     * @param inputs Documents to use for replacements
     */
    PlaceholderDocGenerator(
        const std::vector<std::string>& inputs,
        const std::vector<TemplateSpec>& specs) {

        initMatches(specs, inputs);
    }

protected:
    PlaceholderDocGenerator() {}

    /**
     * Really belongs in constructor, but is decoupled for subclasses
     */
    void initMatches(
        const std::vector<TemplateSpec>& specs,
        const std::vector<std::string>& inputs) {

        using namespace Placeholders;

        for (std::vector<TemplateSpec>::const_iterator ii = specs.begin();
                ii != specs.end(); ++ii) {
            Placeholders::Spec cur(ii->term, ii->minval, ii->maxval, ii->sequential);
            pl_specs.push_back(cur);
        }

        for (std::vector<std::string>::const_iterator ii = inputs.begin();
                ii != inputs.end(); ++ii) {
            matches.push_back(new DocumentMatches(*ii, pl_specs));
        }
    }

public:
    GeneratorState *createState(int total, int cur) const {
        return new MyState(this, total, cur);
    }

private:
    class MyState : public GeneratorState {
    public:
        MyState(const PlaceholderDocGenerator *parent, int total, int cur)
        : m_parent(parent) {
            using namespace Placeholders;
            for (size_t ii = 0; ii < parent->matches.size(); ii++) {
                m_substs.push_back(
                    Substitutions(parent->matches[ii], total, cur));
            }
            m_bufs.resize(m_substs.size());
        }

        void populateIov(uint32_t seq, std::vector<lcb_IOV>& iov_out) {
            using namespace Placeholders;
            size_t ix = seq % m_substs.size();
            Substitutions& subst = m_substs[ix];
            Substitutions::Backbuffer& buf = m_bufs[ix];
            subst.makeIovs(iov_out, buf);
        }

        std::vector<Placeholders::Substitutions::Backbuffer> m_bufs;
        std::vector<Placeholders::Substitutions> m_substs;
        const PlaceholderDocGenerator *m_parent;
    };

    // Match object for each doc
    std::vector<Placeholders::DocumentMatches*> matches;
    std::vector<Placeholders::Spec> pl_specs;
};

/**
 * Generate documents based on JSON fields. This adds on top of the normal
 * document generator in that JSON paths are used and explicit placeholders
 * are not required
 */
class PlaceholderJsonGenerator : public PlaceholderDocGenerator {
public:
    /**
     * @param specs List of specs. The term in the spec refers to the field
     * which is to be replaced, not the placeholder text
     * @param documents The documents to operate on
     */
    PlaceholderJsonGenerator(
        const std::vector<std::string>& documents,
        const std::vector<TemplateSpec>& specs) {

        initJsonPlaceholders(specs, documents);
    }

    /**
     * Generate the documents, and then generate specs for them
     * @param specs The specs to use
     * @param minsz Minimum document size
     * @param maxsz Maximum document size
     */
    PlaceholderJsonGenerator(uint32_t minsz, uint32_t maxsz,
        const std::vector<TemplateSpec>& specs) {

        std::vector<std::string> jsondocs;
        JsonDocGenerator::genDocuments(minsz, maxsz, jsondocs);
        initJsonPlaceholders(specs, jsondocs);
    }

private:
    void initJsonPlaceholders(
        const std::vector<TemplateSpec>& specs,
        const std::vector<std::string>& documents) {

        int serial = 0;
        Json::Reader reader;
        Json::FastWriter writer;
        std::vector<TemplateSpec> new_specs;
        std::vector<std::string> new_docs;

        for (size_t ii = 0; ii < documents.size(); ii++) {
            Json::Value root;
            if (!reader.parse(documents[ii], root)) {
                throw std::runtime_error("Couldn't parse one or more documents!");
            }

            for (size_t jj = 0; jj < specs.size(); ++jj) {
                char buf[64];
                const TemplateSpec& spec = specs[jj];
                sprintf(buf, "$__pillowfight_%d", serial++);
                root[spec.term] = buf;
                TemplateSpec new_spec = spec;

                std::string replace_term(buf);
                replace_term = '"' + replace_term + '"';
                new_spec.term = replace_term;
                new_specs.push_back(new_spec);
            }
            new_docs.push_back(writer.write(root));
        }
        initMatches(new_specs, new_docs);
    }
};

} // namespace