#include "StringTable.h"

StringTable::StringTable(void) {}

// Load string table from a binary blob, filling out with the current localisation data only
StringTable::StringTable(PBYTE pbData, DWORD dwSize)
{
	src = byteArray(pbData, dwSize);

	ProcessStringTableData();
}

StringTable::~StringTable(void)
{
	Reset();
}

void StringTable::Reset(void)
{
	m_stringsMap.clear();
    m_stringsVec.clear();
    m_langSizeMap.clear();
    currentLanguage.clear();
}

void StringTable::ReloadStringTable()
{
    Reset();
    ProcessStringTableData();
}

void StringTable::ProcessStringTableData(void)
{
	if (src.length > 0 && src.data != nullptr) 
    {
        ByteArrayInputStream bais(src);
        DataInputStream dis(&bais);
        
        ReadAllFromStream(&dis);
        
        dis.close();
        bais.reset();
    }
}

void StringTable::ReadAllFromStream(DataInputStream* dis)
{
    int versionNumber = dis->readInt();
    int languagesCount = dis->readInt();

    m_langSizeMap.clear();

    for(int i = 0; i < languagesCount; ++i)
    {
        wstring langId = dis->readUTF();
        int langSize = dis->readInt();
        m_langSizeMap.push_back(make_pair(langId, langSize));
    }

    int64_t bytesToSkip = 0;
    int dataSize = 0;

    if(FindLanguage(currentLanguage, bytesToSkip, dataSize))
    {
        dis->skip(bytesToSkip);
        ReadLanguageFromStream(dis, dataSize);
    }
}

bool StringTable::FindLanguage(wstring& outLangId, int64_t& outBytesToSkip, int& outDataSize)
{
    vector<wstring> locales;
    app.getLocale(locales);

    outBytesToSkip = 0;

    for (auto it_locales = locales.begin(); it_locales != locales.end(); ++it_locales)
    {
        outBytesToSkip = 0;

        for(auto& it : m_langSizeMap)
        {
            if(it.first.compare(*it_locales) == 0)
            {
                app.DebugPrintf("StringTable:: Found language '%ls'.\n", it_locales->c_str());
                outDataSize = it.second;
                outLangId = *it_locales;
                return true;
            }
            outBytesToSkip += it.second;
        }

        app.DebugPrintf("StringTable:: Can't find language '%ls'.\n", it_locales->c_str());
    }

    app.DebugPrintf("Failed to get language\n");
    return false;
}

void StringTable::ReadLanguageFromStream(DataInputStream* dis, int dataSize)
{
    // Read the language file for the selected language
    int langVersion = dis->readInt();
    wstring langId = dis->readUTF();
    int totalStrings = dis->readInt();

    app.DebugPrintf("StringTable:: Hybrid Mode. totalStrings = %d\n", totalStrings);

    m_stringsVec.reserve(totalStrings);
    m_stringsMap.reserve(totalStrings);

    for(int i = 0; i < totalStrings; ++i)
    {
        wstring stringId = dis->readUTF();
        wstring stringValue = dis->readUTF();

        m_stringsVec.push_back(std::move(stringValue));

        size_t indexInVector = m_stringsVec.size() - 1;
        m_stringsMap.emplace(std::move(stringId), indexInVector);
    }
}

void StringTable::write(DataOutputStream* dos)
{    
    dos->writeInt(1);
    dos->writeInt(1);

    int langVersion = 1;

    dos->writeInt(langVersion);
    dos->writeUTF(currentLanguage.empty() ? L"en_US" : currentLanguage);

    dos->writeInt(m_stringsMap.size());
    
    for (auto const& [key, index] : m_stringsMap)
    {
        dos->writeUTF(key);
        dos->writeUTF(m_stringsVec[index]);
    }
}

void StringTable::getData(PBYTE *ppData, UINT *pSize)
{
	*ppData = src.data;
	*pSize = src.length;
}

LPCWSTR StringTable::getString(const wstring &id)
{
    auto it = m_stringsMap.find(id);

    if(it != m_stringsMap.end())
		return m_stringsVec[it->second].c_str();
	else
		return L"";
}

LPCWSTR StringTable::getString(int id)
{
	if (id >= 0 && id < m_stringsVec.size())
	{
		LPCWSTR pwchString=m_stringsVec[id].c_str();
		return pwchString;
	}
	else
		return L"";
}