/*******************************************************************************
 * Copyright (c) 2014, Jan Koester jan.koester@gmx.net
 * All rights reserved.
 * 
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 * Redistributions of source code must retain the above copyright
 *      notice, this list of conditions and the following disclaimer.
 * Redistributions in binary form must reproduce the above copyright
 *      notice, this list of conditions and the following disclaimer in the
 *      documentation and/or other materials provided with the distribution.
 * Neither the name of the <organization> nor the
 *      names of its contributors may be used to endorse or promote products
 *      derived from this software without specific prior written permission.
 * 
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL <COPYRIGHT HOLDER> BE LIABLE FOR ANY
 * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
 * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *******************************************************************************/

#include <cstring>
#include <stdexcept>

#include "hpack.h"

namespace libhttppp::hpack {

// ============================================================
// HPACK Static Table (RFC 7541, Appendix A)
// ============================================================
const Decoder::StaticEntry Decoder::STATIC_TABLE[] = {
    { "",                         ""              },  // index 0 (unused)
    { ":authority",               ""              },  // 1
    { ":method",                  "GET"           },  // 2
    { ":method",                  "POST"          },  // 3
    { ":path",                    "/"             },  // 4
    { ":path",                    "/index.html"   },  // 5
    { ":scheme",                  "http"          },  // 6
    { ":scheme",                  "https"         },  // 7
    { ":status",                  "200"           },  // 8
    { ":status",                  "204"           },  // 9
    { ":status",                  "206"           },  // 10
    { ":status",                  "304"           },  // 11
    { ":status",                  "400"           },  // 12
    { ":status",                  "404"           },  // 13
    { ":status",                  "500"           },  // 14
    { "accept-charset",           ""              },  // 15
    { "accept-encoding",          "gzip, deflate" },  // 16
    { "accept-language",          ""              },  // 17
    { "accept-ranges",            ""              },  // 18
    { "accept",                   ""              },  // 19
    { "access-control-allow-origin", ""           },  // 20
    { "age",                      ""              },  // 21
    { "allow",                    ""              },  // 22
    { "authorization",            ""              },  // 23
    { "cache-control",            ""              },  // 24
    { "content-disposition",      ""              },  // 25
    { "content-encoding",         ""              },  // 26
    { "content-language",         ""              },  // 27
    { "content-length",           ""              },  // 28
    { "content-location",         ""              },  // 29
    { "content-range",            ""              },  // 30
    { "content-type",             ""              },  // 31
    { "cookie",                   ""              },  // 32
    { "date",                     ""              },  // 33
    { "etag",                     ""              },  // 34
    { "expect",                   ""              },  // 35
    { "expires",                  ""              },  // 36
    { "from",                     ""              },  // 37
    { "host",                     ""              },  // 38
    { "if-match",                 ""              },  // 39
    { "if-modified-since",        ""              },  // 40
    { "if-none-match",            ""              },  // 41
    { "if-range",                 ""              },  // 42
    { "if-unmodified-since",      ""              },  // 43
    { "last-modified",            ""              },  // 44
    { "link",                     ""              },  // 45
    { "location",                 ""              },  // 46
    { "max-forwards",             ""              },  // 47
    { "proxy-authenticate",       ""              },  // 48
    { "proxy-authorization",      ""              },  // 49
    { "range",                    ""              },  // 50
    { "referer",                  ""              },  // 51
    { "refresh",                  ""              },  // 52
    { "retry-after",              ""              },  // 53
    { "server",                   ""              },  // 54
    { "set-cookie",               ""              },  // 55
    { "strict-transport-security", ""             },  // 56
    { "transfer-encoding",        ""              },  // 57
    { "user-agent",               ""              },  // 58
    { "vary",                     ""              },  // 59
    { "via",                      ""              },  // 60
    { "www-authenticate",         ""              },  // 61
};
const size_t Decoder::STATIC_TABLE_SIZE = 62; // 0..61

// ============================================================
// HPACK Huffman table (RFC 7541, Appendix B)
// ============================================================
struct HuffEntry { uint32_t code; uint8_t bits; };
static const HuffEntry HUFF_TABLE[257] = {
    {0x1ff8,     13}, {0x7fffd8,   23}, {0xfffffe2,  28}, {0xfffffe3,  28},
    {0xfffffe4,  28}, {0xfffffe5,  28}, {0xfffffe6,  28}, {0xfffffe7,  28},
    {0xfffffe8,  28}, {0xffffea,   24}, {0x3ffffffc, 30}, {0xfffffe9,  28},
    {0xfffffea,  28}, {0x3ffffffd, 30}, {0xfffffeb,  28}, {0xfffffec,  28},
    {0xfffffed,  28}, {0xfffffee,  28}, {0xfffffef,  28}, {0xffffff0,  28},
    {0xffffff1,  28}, {0xffffff2,  28}, {0x3ffffffe, 30}, {0xffffff3,  28},
    {0xffffff4,  28}, {0xffffff5,  28}, {0xffffff6,  28}, {0xffffff7,  28},
    {0xffffff8,  28}, {0xffffff9,  28}, {0xffffffa,  28}, {0xffffffb,  28},
    {0x14,        6}, {0x3f8,      10}, {0x3f9,      10}, {0xffa,      12},
    {0x1ff9,     13}, {0x15,        6}, {0xf8,        8}, {0x7fa,      11},
    {0x3fa,      10}, {0x3fb,      10}, {0xf9,        8}, {0x7fb,      11},
    {0xfa,        8}, {0x16,        6}, {0x17,        6}, {0x18,        6},
    {0x0,         5}, {0x1,         5}, {0x2,         5}, {0x19,        6},
    {0x1a,        6}, {0x1b,        6}, {0x1c,        6}, {0x1d,        6},
    {0x1e,        6}, {0x1f,        6}, {0x5c,        7}, {0xfb,        8},
    {0x7ffc,     15}, {0x20,        6}, {0xffb,      12}, {0x3fc,      10},
    {0x1ffa,     13}, {0x21,        6}, {0x5d,        7}, {0x5e,        7},
    {0x5f,        7}, {0x60,        7}, {0x61,        7}, {0x62,        7},
    {0x63,        7}, {0x64,        7}, {0x65,        7}, {0x66,        7},
    {0x67,        7}, {0x68,        7}, {0x69,        7}, {0x6a,        7},
    {0x6b,        7}, {0x6c,        7}, {0x6d,        7}, {0x6e,        7},
    {0x6f,        7}, {0x70,        7}, {0x71,        7}, {0x72,        7},
    {0xfc,        8}, {0x73,        7}, {0xfd,        8}, {0x1ffb,     13},
    {0x7fff0,    19}, {0x1ffc,     13}, {0x3ffc,     14}, {0x22,        6},
    {0x7ffd,     15}, {0x3,         5}, {0x23,        6}, {0x4,         5},
    {0x24,        6}, {0x5,         5}, {0x25,        6}, {0x26,        6},
    {0x27,        6}, {0x6,         5}, {0x74,        7}, {0x75,        7},
    {0x28,        6}, {0x29,        6}, {0x2a,        6}, {0x7,         5},
    {0x2b,        6}, {0x76,        7}, {0x2c,        6}, {0x8,         5},
    {0x9,         5}, {0x2d,        6}, {0x77,        7}, {0x78,        7},
    {0x79,        7}, {0x7a,        7}, {0x7b,        7}, {0x7fffe,    19},
    {0x7fc,      11}, {0x3ffd,     14}, {0x1ffd,     13}, {0xffffffc,  28},
    {0xfffe6,    20}, {0x3fffd2,   22}, {0xfffe7,    20}, {0xfffe8,    20},
    {0x3fffd3,   22}, {0x3fffd4,   22}, {0x3fffd5,   22}, {0x7fffd9,   23},
    {0x3fffd6,   22}, {0x7fffda,   23}, {0x7fffdb,   23}, {0x7fffdc,   23},
    {0x7fffdd,   23}, {0x7fffde,   23}, {0xffffeb,   24}, {0x7fffdf,   23},
    {0xffffec,   24}, {0xffffed,   24}, {0x3fffd7,   22}, {0x7fffe0,   23},
    {0xffffee,   24}, {0x7fffe1,   23}, {0x7fffe2,   23}, {0x7fffe3,   23},
    {0x7fffe4,   23}, {0x1fffdc,   21}, {0x3fffd8,   22}, {0x7fffe5,   23},
    {0x3fffd9,   22}, {0x7fffe6,   23}, {0x7fffe7,   23}, {0xffffef,   24},
    {0x3fffda,   22}, {0x1fffdd,   21}, {0xfffe9,    20}, {0x3fffdb,   22},
    {0x3fffdc,   22}, {0x7fffe8,   23}, {0x7fffe9,   23}, {0x1fffde,   21},
    {0x7fffea,   23}, {0x3fffdd,   22}, {0x3fffde,   22}, {0xfffff0,   24},
    {0x1fffdf,   21}, {0x3fffdf,   22}, {0x7fffeb,   23}, {0x7fffec,   23},
    {0x1fffe0,   21}, {0x1fffe1,   21}, {0x3fffe0,   22}, {0x1fffe2,   21},
    {0x7fffed,   23}, {0x3fffe1,   22}, {0x7fffee,   23}, {0x7fffef,   23},
    {0xfffea,    20}, {0x3fffe2,   22}, {0x3fffe3,   22}, {0x3fffe4,   22},
    {0x7ffff0,   23}, {0x3fffe5,   22}, {0x3fffe6,   22}, {0x7ffff1,   23},
    {0x3ffffe0,  26}, {0x3ffffe1,  26}, {0xfffeb,    20}, {0x7fff1,    19},
    {0x3fffe7,   22}, {0x7ffff2,   23}, {0x3fffe8,   22}, {0x1ffffec,  25},
    {0x3ffffe2,  26}, {0x3ffffe3,  26}, {0x3ffffe4,  26}, {0x7ffffde,  27},
    {0x7ffffdf,  27}, {0x3ffffe5,  26}, {0xfffff1,   24}, {0x1ffffed,  25},
    {0x7fff2,    19}, {0x1fffe3,   21}, {0x3ffffe6,  26}, {0x7ffffe0,  27},
    {0x7ffffe1,  27}, {0x3ffffe7,  26}, {0x7ffffe2,  27}, {0xfffff2,   24},
    {0x1fffe4,   21}, {0x1fffe5,   21}, {0x3ffffe8,  26}, {0x3ffffe9,  26},
    {0xffffffd,  28}, {0x7ffffe3,  27}, {0x7ffffe4,  27}, {0x7ffffe5,  27},
    {0xfffec,    20}, {0xfffff3,   24}, {0xfffed,    20}, {0x1fffe6,   21},
    {0x3fffe9,   22}, {0x1fffe7,   21}, {0x1fffe8,   21}, {0x7ffff3,   23},
    {0x3fffea,   22}, {0x3fffeb,   22}, {0x1ffffee,  25}, {0x1ffffef,  25},
    {0xfffff4,   24}, {0xfffff5,   24}, {0x3ffffea,  26}, {0x7ffff4,   23},
    {0x3ffffeb,  26}, {0x7ffffe6,  27}, {0x3ffffec,  26}, {0x3ffffed,  26},
    {0x7ffffe7,  27}, {0x7ffffe8,  27}, {0x7ffffe9,  27}, {0x7ffffea,  27},
    {0x7ffffeb,  27}, {0xffffffe,  28}, {0x7ffffec,  27}, {0x7ffffed,  27},
    {0x7ffffee,  27}, {0x7ffffef,  27}, {0x7fffff0,  27}, {0x3ffffee,  26},
    {0x3fffffff, 30}, // EOS (256)
};

std::string Decoder::huffmanDecode(const uint8_t *data, size_t len) {
    std::string out;
    uint64_t bits = 0;
    int nbits = 0;

    for (size_t i = 0; i < len; ++i) {
        bits = (bits << 8) | data[i];
        nbits += 8;

        while (nbits >= 5) { // shortest Huffman code is 5 bits
            bool found = false;
            for (int sym = 0; sym < 256; ++sym) {
                uint8_t codelen = HUFF_TABLE[sym].bits;
                if (codelen > (uint8_t)nbits) continue;
                uint32_t code = HUFF_TABLE[sym].code;
                uint32_t candidate = static_cast<uint32_t>((bits >> (nbits - codelen)) & ((1ULL << codelen) - 1));
                if (candidate == code) {
                    out.push_back(static_cast<char>(sym));
                    nbits -= codelen;
                    bits &= (1ULL << nbits) - 1;
                    found = true;
                    break;
                }
            }
            if (!found) break; // remaining bits are padding
        }
    }
    return out;
}

uint32_t Decoder::decodeInt(const uint8_t *data, size_t len, uint8_t prefix_bits, size_t &bytes_read) {
    if (len == 0) { bytes_read = 0; return 0; }
    uint32_t max_prefix = (1u << prefix_bits) - 1u;
    uint32_t value = data[0] & max_prefix;
    bytes_read = 1;
    if (value < max_prefix) return value;
    uint32_t m = 0;
    while (bytes_read < len) {
        uint8_t b = data[bytes_read++];
        if (m < 32) {
            value += static_cast<uint32_t>(b & 0x7f) << m;
        }
        m += 7;
        if (!(b & 0x80)) return value;
        if (m >= 32) return value; // prevent overflow / runaway
    }
    return value;
}

std::string Decoder::decodeString(const uint8_t *data, size_t len, size_t &bytes_read) {
    if (len == 0) { bytes_read = 0; return ""; }
    bool huffman = (data[0] & 0x80) != 0;
    size_t int_bytes = 0;
    uint32_t slen = decodeInt(data, len, 7, int_bytes);
    bytes_read = int_bytes + slen;
    if (bytes_read > len) { bytes_read = len; return ""; }
    if (huffman) {
        return huffmanDecode(data + int_bytes, slen);
    }
    return std::string(reinterpret_cast<const char*>(data + int_bytes), slen);
}

std::vector<HeaderField> Decoder::decode(const std::string &data) {
    return decode(reinterpret_cast<const uint8_t*>(data.data()), data.size());
}

std::vector<HeaderField> Decoder::decodeStateless(const uint8_t *data, size_t len) {
    Decoder tmp;
    return tmp.decode(data, len);
}

// ----------------------------------------------------------------
// Dynamic table helpers (RFC 7541 §4)
// ----------------------------------------------------------------
// Entry size = name length + value length + 32 (RFC 7541 §4.1)
static size_t entrySize(const std::string &name, const std::string &value) {
    return name.size() + value.size() + 32;
}

void Decoder::evict(size_t required) {
    while (_dynTableSize + required > _dynTableMaxSize && !_dynTable.empty()) {
        auto &back = _dynTable.back();
        _dynTableSize -= entrySize(back.name, back.value);
        _dynTable.pop_back();
    }
}

void Decoder::addEntry(const std::string &name, const std::string &value) {
    size_t sz = entrySize(name, value);
    evict(sz);
    if (sz <= _dynTableMaxSize) {
        _dynTable.push_front({name, value});
        _dynTableSize += sz;
    }
}

bool Decoder::lookupIndex(uint32_t idx, std::string &name, std::string &value) const {
    if (idx == 0) return false;
    if (idx < STATIC_TABLE_SIZE) {
        name = STATIC_TABLE[idx].name;
        value = STATIC_TABLE[idx].value;
        return true;
    }
    // Dynamic table: index 62 = first entry (front)
    size_t dynIdx = idx - STATIC_TABLE_SIZE;
    if (dynIdx < _dynTable.size()) {
        name = _dynTable[dynIdx].name;
        value = _dynTable[dynIdx].value;
        return true;
    }
    return false;
}

bool Decoder::lookupIndexName(uint32_t idx, std::string &name) const {
    if (idx == 0) return false;
    if (idx < STATIC_TABLE_SIZE) {
        name = STATIC_TABLE[idx].name;
        return true;
    }
    size_t dynIdx = idx - STATIC_TABLE_SIZE;
    if (dynIdx < _dynTable.size()) {
        name = _dynTable[dynIdx].name;
        return true;
    }
    return false;
}

std::vector<HeaderField> Decoder::decode(const uint8_t *data, size_t len) {
    std::vector<HeaderField> headers;
    size_t off = 0;

    while (off < len) {
        size_t prev_off = off;   // progress guard
        uint8_t byte = data[off];

        if (byte & 0x80) {
            // Indexed header field (section 6.1): top bit set
            size_t int_bytes = 0;
            uint32_t idx = decodeInt(data + off, len - off, 7, int_bytes);
            off += int_bytes;
            std::string name, value;
            if (lookupIndex(idx, name, value)) {
                headers.push_back({name, value});
            }
        } else if ((byte & 0xC0) == 0x40) {
            // Literal with incremental indexing (section 6.2.1)
            size_t int_bytes = 0;
            uint32_t idx = decodeInt(data + off, len - off, 6, int_bytes);
            off += int_bytes;
            std::string name;
            if (idx > 0) {
                lookupIndexName(idx, name);
            } else {
                size_t str_bytes = 0;
                name = decodeString(data + off, len - off, str_bytes);
                off += str_bytes;
            }
            size_t str_bytes = 0;
            std::string value = decodeString(data + off, len - off, str_bytes);
            off += str_bytes;
            headers.push_back({name, value, true});
            // Add to dynamic table
            addEntry(name, value);
        } else if ((byte & 0xF0) == 0x00 || (byte & 0xF0) == 0x10) {
            // Literal without indexing (section 6.2.2) or never indexed (section 6.2.3)
            uint8_t prefix = 4;
            size_t int_bytes = 0;
            uint32_t idx = decodeInt(data + off, len - off, prefix, int_bytes);
            off += int_bytes;
            std::string name;
            if (idx > 0) {
                lookupIndexName(idx, name);
            } else {
                size_t str_bytes = 0;
                name = decodeString(data + off, len - off, str_bytes);
                off += str_bytes;
            }
            size_t str_bytes = 0;
            std::string value = decodeString(data + off, len - off, str_bytes);
            off += str_bytes;
            headers.push_back({name, value});
        } else if ((byte & 0xE0) == 0x20) {
            // Dynamic table size update (section 6.3)
            size_t int_bytes = 0;
            uint32_t newMax = decodeInt(data + off, len - off, 5, int_bytes);
            off += int_bytes;
            _dynTableMaxSize = newMax;
            evict(0); // evict entries that exceed new max
        } else {
            ++off; // skip unknown
        }

        // Safety: if no byte was consumed this iteration, advance by one
        // to avoid an infinite loop on malformed input.
        if (off == prev_off) {
            ++off;
        }
    }
    return headers;
}

static void appendInt(std::string &out, uint8_t prefix_bits, uint8_t prefix_mask, uint32_t value) {
    uint8_t max_prefix = static_cast<uint8_t>((1u << prefix_bits) - 1u);
    if (value < max_prefix) {
        out.push_back(static_cast<char>(prefix_mask | value));
        return;
    }
    out.push_back(static_cast<char>(prefix_mask | max_prefix));
    value -= max_prefix;
    while (value >= 128) {
        out.push_back(static_cast<char>((value % 128) + 128));
        value /= 128;
    }
    out.push_back(static_cast<char>(value));
}

void Encoder::appendLiteral(std::string &out,
                            const std::string &name,
                            const std::string &value,
                            bool incremental_index) {
    uint8_t prefix_mask = incremental_index ? 0x40 : 0x00;
    appendInt(out, 6, prefix_mask, 0);
    appendInt(out, 7, 0x00, static_cast<uint32_t>(name.size()));
    out.append(name);
    appendInt(out, 7, 0x00, static_cast<uint32_t>(value.size()));
    out.append(value);
}

void Encoder::appendLiteralWithNameIndex(std::string &out,
                                         uint8_t name_index,
                                         const std::string &value,
                                         bool incremental_index) {
    uint8_t prefix_mask = incremental_index ? 0x40 : 0x00;
    appendInt(out, 6, prefix_mask, name_index);
    appendInt(out, 7, 0x00, static_cast<uint32_t>(value.size()));
    out.append(value);
}

std::string Encoder::encodeResponseHeaders(uint16_t status_code,
                                           const std::string &content_type,
                                           size_t content_length,
                                           const std::vector<HeaderField> &extra) {
    std::string out;

    if (status_code == 200) {
        out.push_back(static_cast<char>(0x88));  // index 8
    } else if (status_code == 204) {
        out.push_back(static_cast<char>(0x89));  // index 9
    } else if (status_code == 206) {
        out.push_back(static_cast<char>(0x8a));  // index 10
    } else if (status_code == 304) {
        out.push_back(static_cast<char>(0x8b));  // index 11
    } else {
        std::string status = std::to_string(status_code);
        appendLiteralWithNameIndex(out, 0x08, status, false);
    }

    if (!content_type.empty()) {
        appendLiteralWithNameIndex(out, 0x1f, content_type, true);
    }

    if (content_length > 0) {
        std::string len = std::to_string(content_length);
        appendLiteral(out, "content-length", len, false);
    }

    for (const auto &h : extra) {
        if (h.name.empty()) {
            continue;
        }
        appendLiteral(out, h.name, h.value, h.incremental_index);
    }

    return out;
}

std::string Encoder::encodeRequestHeaders(const std::string &method,
                                           const std::string &path,
                                           const std::string &scheme,
                                           const std::string &authority,
                                           const std::vector<HeaderField> &extra) {
    std::string out;

    // :method — static table index 2 = GET, 3 = POST
    if (method == "GET") {
        out.push_back(static_cast<char>(0x82)); // indexed: 2
    } else if (method == "POST") {
        out.push_back(static_cast<char>(0x83)); // indexed: 3
    } else {
        // :method name is at index 2, literal with name index
        appendLiteralWithNameIndex(out, 0x02, method, false);
    }

    // :path — static table index 4 = "/", 5 = "/index.html"
    if (path == "/") {
        out.push_back(static_cast<char>(0x84)); // indexed: 4
    } else if (path == "/index.html") {
        out.push_back(static_cast<char>(0x85)); // indexed: 5
    } else {
        // :path name is at index 4, literal with name index
        appendLiteralWithNameIndex(out, 0x04, path, false);
    }

    // :scheme — static table index 6 = "http", 7 = "https"
    if (scheme == "http") {
        out.push_back(static_cast<char>(0x86)); // indexed: 6
    } else if (scheme == "https") {
        out.push_back(static_cast<char>(0x87)); // indexed: 7
    } else {
        appendLiteralWithNameIndex(out, 0x06, scheme, false);
    }

    // :authority — static table index 1 (name only, no value)
    if (!authority.empty()) {
        appendLiteralWithNameIndex(out, 0x01, authority, false);
    }

    // Extra headers
    for (const auto &h : extra) {
        if (h.name.empty()) continue;
        appendLiteral(out, h.name, h.value, h.incremental_index);
    }

    return out;
}

} // namespace libhttppp::hpack
