#include "pch.h"
#include "binary_stream.h"
#include "asset.h"

BinaryStreamReader::~BinaryStreamReader()
{
    m_ptr = m_cur = nullptr;
    m_size = 0;
}

void BinaryStreamReader::init(uint8_t* data_ptr, size_t size, ByteOrder byte_order /*= ByteOrder::Host*/)
{
    m_ptr = m_cur = data_ptr;
    m_size = size;
    m_byte_order = byte_order;
    m_swap = byte_order == ByteOrder::Host ? false : byte_order != sys_order();
}

bool BinaryStreamReader::load(const std::string& path, ByteOrder byte_order /*= ByteOrder::Host*/)
{
    if (m_asset.open(path.c_str()))
    {
        m_asset.read_all();
        init(m_asset.m_data, m_asset.m_len, byte_order);
        return true;
    }
    return false;
}

size_t BinaryStreamReader::pos()
{
    return std::distance(m_ptr, m_cur);
}

void BinaryStreamReader::skip(size_t bytes)
{
    m_cur += bytes;
}

void BinaryStreamReader::snap()
{
    if ((size_t)m_cur % 4 != 0) m_cur += 4 - (size_t)m_cur % 4;
}

bool BinaryStreamReader::eof()
{
    return std::distance(m_ptr, m_cur) >= m_size;
}

bool BinaryStreamReader::has_data(size_t sz)
{
    return std::distance(m_ptr, m_cur + sz) <= m_size;
}

uint8_t BinaryStreamReader::ru8()
{
    return read<uint8_t>();
}

uint16_t BinaryStreamReader::ru16()
{
    return m_swap ? swap(read<uint16_t>()) : read<uint16_t>();
}

uint32_t BinaryStreamReader::ru32()
{
    return m_swap ? swap(read<uint32_t>()) : read<uint32_t>();
}

uint64_t BinaryStreamReader::ru64()
{
    return m_swap ? swap(read<uint64_t>()) : read<uint64_t>();
}

int8_t BinaryStreamReader::ri8()
{
    return read<int8_t>();
}

int16_t BinaryStreamReader::ri16()
{
    return m_swap ? swap(read<int16_t>()) : read<int16_t>();
}

int32_t BinaryStreamReader::ri32()
{
    return m_swap ? swap(read<int32_t>()) : read<int32_t>();
}

int64_t BinaryStreamReader::ri64()
{
    return m_swap ? swap(read<int64_t>()) : read<int64_t>();
}

float BinaryStreamReader::rflt()
{
    return m_swap ? swap(read<float>()) : read<float>();
}

double BinaryStreamReader::rdbl()
{
    return m_swap ? swap(read<double>()) : read<double>();
}

std::string BinaryStreamReader::pick(size_t chars)
{
    return { (char*)m_cur, chars };
}

std::string BinaryStreamReader::rstring(size_t len)
{
    return { advance<char>(len), len };
}

std::string BinaryStreamReader::rstring()
{
    auto len = ru32();
    return { advance<char>(len), len };
}

std::wstring BinaryStreamReader::rwstring(size_t len)
{
    auto ptr = advance<char>(len * 2);
    //for (int i = 0; i < len; i++)
    //    std::swap(ptr[i * 2], ptr[i * 2 + 1]);

    // right trim trailing zeroes
    if (len > 0)
    {
        auto wptr = (uint16_t*)ptr;
        for (int i = (int)len - 1; i >= 0; i--)
            if (wptr[i] == 0) len--;
    }
    
    // wide to UTF-16le
    std::wstring_convert<std::codecvt_utf16<wchar_t, 0x10ffff>> converter;
    return converter.from_bytes(ptr, ptr + len * 2);
}

std::wstring BinaryStreamReader::rwstring()
{
    auto len = ru32();
    return rwstring(len);
}

std::string BinaryStreamReader::rpascal()
{
    auto len = ru8();
    return rstring(len);
}

std::vector<glm::uint8_t> BinaryStreamReader::rraw(size_t bytes)
{
    std::vector<uint8_t> ret(bytes);
    std::copy_n(m_cur, bytes, ret.data());
    skip(bytes);
    return ret;
}

std::vector<glm::uint8_t> BinaryStreamReader::rraw()
{
    auto size = ru32();
    return rraw(size);
}

std::vector<glm::uint8_t> BinaryStreamReader::rrle(size_t encoded_bytes)
{
    std::vector<uint8_t> data;
    int32_t n;
    for (int j = 0; j < encoded_bytes;)
    {
        n = ri8();
        j++;
        // copy the following char -n + 1 times
        if (n < 0)
        {
            // NOP
            if (n == -128)
                continue;
            n = -n + 1;

            j++;
            auto ch = ru8();
            for (int c = 0; c < n; c++)
                data.push_back(ch);
        }
        else
        {
            // read the following n + 1 chars (no compr)
            for (int c = 0; c < n + 1; c++, j++)
                data.push_back(ru8());
        }
    }
    return data;
}

std::string BinaryStreamReader::rkey_or_string()
{
    auto len = ru32();
    if (len == 0)
        len = 4;
    return rstring(len);
}

void BinaryStreamWriter::init(ByteOrder byte_order /*= ByteOrder::Host*/)
{
    m_byte_order = byte_order;
    m_swap = byte_order == ByteOrder::Host ? false : byte_order != sys_order();
}

bool BinaryStreamWriter::save(const std::string& path) const
{
    std::ofstream f(path, std::ios::binary);
    if (f.good())
    {
        f.write((char*)m_data.data(), m_data.size());
        return true;
    }
    return false;
}

void BinaryStreamWriter::skip(size_t bytes, uint8_t fill /*= 0*/)
{
    m_data.resize(m_data.size() + bytes);
}

void BinaryStreamWriter::snap()
{
    if ((size_t)m_data.size() % 4 != 0) m_data.resize(m_data.size() + 4 - (size_t)m_data.size() % 4);
}

void BinaryStreamWriter::wu8(uint8_t v)
{
    return write<uint8_t>(v);
}

void BinaryStreamWriter::wu16(uint16_t v)
{
    return m_swap ? write(swap(v)) : write(v);
}

void BinaryStreamWriter::wu32(uint32_t v)
{
    return m_swap ? write(swap(v)) : write(v);
}

void BinaryStreamWriter::wu64(uint64_t v)
{
    return m_swap ? write(swap(v)) : write(v);
}

void BinaryStreamWriter::wi8(int8_t v)
{
    return write<int8_t>(v);
}

void BinaryStreamWriter::wi16(int16_t v)
{
    return m_swap ? write(swap(v)) : write(v);
}

void BinaryStreamWriter::wi32(int32_t v)
{
    return m_swap ? write(swap(v)) : write(v);
}

void BinaryStreamWriter::wi64(int64_t v)
{
    return m_swap ? write(swap(v)) : write(v);
}

void BinaryStreamWriter::wflt(float v)
{
    return m_swap ? write(swap(v)) : write(v);
}

void BinaryStreamWriter::wdbl(double v)
{
    return m_swap ? write(swap(v)) : write(v);
}

void BinaryStreamWriter::wstring(std::string s)
{
    wu32((uint32_t)s.size());
    write(s.data(), s.size());
}

void BinaryStreamWriter::wstring_raw(std::string s)
{
    write(s.data(), s.size());
}

void BinaryStreamWriter::wwstring(std::wstring s)
{
    wu32((uint32_t)s.size());
    std::wstring_convert<std::codecvt_utf16<wchar_t, 0x10ffff>> converter;
    wstring_raw(converter.to_bytes(s));
}

void BinaryStreamWriter::wwstring_raw(std::wstring s)
{
    std::wstring_convert<std::codecvt_utf16<wchar_t, 0x10ffff>> converter;
    wstring(converter.to_bytes(s));
}

void BinaryStreamWriter::wpascal(std::string s)
{
    wu8((uint8_t)s.size());
    write(s.data(), s.size());
}

void BinaryStreamWriter::wraw(std::vector<uint8_t> raw)
{
    wu32((uint32_t)raw.size());
    write(raw.data(), raw.size());
}

void BinaryStreamWriter::wrle(std::vector<uint8_t> data)
{
    size_t idx = 0;
    auto ptr = data.data();

    // find the next sequence and return the position
    auto next_seq = [&]() -> std::pair<size_t, size_t> {
        size_t i = idx;
        while (i < data.size() - 1)
        {
            size_t count = 0;
            for (size_t j = i + 1; j < data.size() && data[i] == data[j]; j++)
                count++;
            if (count > 0)
                return { i, count + 1 };
            i++;
        }
        return { i + 1, 0 };
    };

    while (idx < data.size())
    {
        auto seq = next_seq();
        // non sequence chars
        if (seq.first - idx > 0)
        {
            wi8(seq.first - idx - 1);
            write(ptr + idx, seq.first - idx);
        }
        if (seq.second > 0)
        {
            wi8(-(int8_t)(seq.second - 1));
            write(ptr[seq.first]);
        }
        idx = seq.first + seq.second;
    }
}

void BinaryStreamWriter::wkey_or_string(std::string s)
{
    assert(s.size() > 0);
    wstring(s);
}