local ffi = require 'ffi'
ffi.cdef[[
    struct SFILE_CREATE_MPQ {
        unsigned long cbSize;         // Size of this structure, in bytes
        unsigned long dwMpqVersion;   // Version of the MPQ to be created
        void*         pvUserData;     // Reserved, must be NULL
        unsigned long cbUserData;     // Reserved, must be 0
        unsigned long dwStreamFlags;  // Stream flags for creating the MPQ
        unsigned long dwFileFlags1;   // File flags for (listfile). 0 = default
        unsigned long dwFileFlags2;   // File flags for (attributes). 0 = default
        unsigned long dwFileFlags3;   // File flags for (signature). 0 = default
        unsigned long dwAttrFlags;    // Flags for the (attributes) file. If 0, no attributes will be created
        unsigned long dwSectorSize;   // Sector size for compressed files
        unsigned long dwRawChunkSize; // Size of raw data chunk
        unsigned long dwMaxFileCount; // File limit for the MPQ
    };
    
    bool SFileCreateArchive2(const wchar_t* szMpqName, struct SFILE_CREATE_MPQ* pCreateInfo, uint32_t* phMpq);
    bool SFileOpenArchive(const wchar_t* szMpqName, unsigned long dwPriority, unsigned long dwFlags, uint32_t* phMpq);
    bool SFileCloseArchive(uint32_t hMpq);
    bool SFileAddFileEx(uint32_t hMpq, const wchar_t* szFileName, const char* szArchivedName, unsigned long dwFlags, unsigned long dwCompression, unsigned long dwCompressionNext);
    bool SFileExtractFile(uint32_t hMpq, const char* szToExtract, const wchar_t* szExtracted, unsigned long dwSearchScope);
    bool SFileHasFile(uint32_t hMpq, const char* szFileName);
    bool SFileSetMaxFileCount(uint32_t hMpq, unsigned long dwMaxFileCount);
    
    bool SFileCreateFile(uint32_t hMpq, const char* szArchivedName, unsigned long long FileTime, unsigned long dwFileSize, unsigned long lcLocale, unsigned long dwFlags, uint32_t* phFile);
    bool SFileWriteFile(uint32_t hFile, const void* pvData, unsigned long dwSize, unsigned long dwCompression);
    bool SFileFinishFile(uint32_t hFile);
    bool SFileOpenFileEx(uint32_t hMpq, const char* szFileName, unsigned long dwSearchScope, uint32_t* phFile);
    bool SFileReadFile(uint32_t hFile, void* lpBuffer, unsigned long dwToRead, unsigned long* pdwRead, void* lpOverlapped);
    unsigned long SFileGetFileSize(uint32_t hFile, unsigned long* pdwFileSizeHigh);
    bool SFileCloseFile(uint32_t hFile);
    bool SFileRemoveFile(uint32_t hMpq, const char* szFileName, unsigned long dwSearchScope);

    bool SFileGetFileInfo(uint32_t hMpqOrFile, int InfoClass, void * pvFileInfo, unsigned long cbFileInfo, unsigned long* pcbLengthNeeded);

    unsigned long SFileGetLocale();

    unsigned long GetLastError();
]]
ffi.cdef[[
    struct SYSTEMTIME {
        unsigned short wYear;
        unsigned short wMonth;
        unsigned short wDayOfWeek;
        unsigned short wDay;
        unsigned short wHour;
        unsigned short wMinute;
        unsigned short wSecond;
        unsigned short wMilliseconds;
    };
    struct FILETIME {
        unsigned long dwLowDateTime;
        unsigned long dwHighDateTime;
    };
    void GetSystemTime(struct SYSTEMTIME* lpSystemTime);
    int SystemTimeToFileTime(const struct SYSTEMTIME* lpSystemTime, struct FILETIME*lpFileTime);
]]

local SFileMpqNumberOfFiles = 36

require 'filesystem'
local uni = require 'ffi.unicode'
local stormlib = ffi.load('stormlib')

local function current_filetime()
    local systemtime = ffi.new('struct SYSTEMTIME')
    local filetime = ffi.new('struct FILETIME')
    ffi.C.GetSystemTime(systemtime)
    if not ffi.C.SystemTimeToFileTime(systemtime, filetime) then
        return 0
    end
    return filetime.dwLowDateTime | (filetime.dwHighDateTime << 32)
end


local wfile = {}
wfile.__index = wfile

function wfile:close()
    if self.handle == 0 then
        return
    end
    stormlib.SFileFinishFile(self.handle)
    self.handle = 0
end

function wfile:write(buf)
    if self.handle == 0 then
        return false
    end
    return stormlib.SFileWriteFile(self.handle, buf, #buf, 0x02)
end

local rfile = {}
rfile.__index = rfile

function rfile:close()
    if self.handle == 0 then
        return
    end
    stormlib.SFileCloseFile(self.handle)
    self.handle = 0
end

function rfile:size()
    if self.handle == 0 then
        return 0
    end
    local size_hi = ffi.new('unsigned long[1]', 0)
    local size_lo = stormlib.SFileGetFileSize(self.handle, size_hi)
    return size_lo | (size_hi[0] << 32)
end

function rfile:read(n)
    if self.handle == 0 then
        return nil
    end
    if not n then
        n = self:size()
    end
    local buf = ffi.new('char[?]', n)
    local pread = ffi.new('unsigned long[1]', 0)
    if not stormlib.SFileReadFile(self.handle, buf, n, pread, nil) then
        return nil
    end
    return ffi.string(buf, pread[0])
end

local archive = {}
archive.__index = archive

function archive:close()
    if self.handle == 0 then
        return
    end
    stormlib.SFileCloseArchive(self.handle)
    self.handle = 0
end

function archive:add_file(name, path)
    if self.handle == 0 then
        return false
    end
    local wpath = uni.u2w(path:string())
    return stormlib.SFileAddFileEx(self.handle, wpath, name,
            0x00000200 | 0x80000000, -- MPQ_FILE_COMPRESS | MPQ_FILE_REPLACEEXISTING,
            0x02, -- MPQ_COMPRESSION_ZLIB,
            0x02 --MPQ_COMPRESSION_ZLIB
            )
end

function archive:extract(name, path)
    if self.handle == 0 then
        return false
    end
    local dir = path:parent_path()
    if not fs.exists(dir) then
        fs.create_directories(dir)
    end
    local wpath = uni.u2w(path:string())
    return stormlib.SFileExtractFile(self.handle, name, wpath,
            0 --SFILE_OPEN_FROM_MPQ
            )
end

function archive:has_file(name)
    if self.handle == 0 then
        return false
    end
    return stormlib.SFileHasFile(self.handle, name)
end

function archive:remove_file(name)
    if self.handle == 0 then
        return false
    end
    return stormlib.SFileRemoveFile(self.handle, name, 0)
end

function archive:open_file(name)
    if self.handle == 0 then
        return nil
    end
    local phandle = ffi.new('uint32_t[1]', 0)
    if not stormlib.SFileOpenFileEx(self.handle, name, 0, phandle) then
        return nil
    end
    return setmetatable({ handle = phandle[0] }, rfile)
end

function archive:create_file(name, size, filetime)
    if self.handle == 0 then
        return nil
    end
    if not filetime then
        filetime = current_filetime()
    end
    local phandle = ffi.new('uint32_t[1]', 0)
    if not stormlib.SFileCreateFile(self.handle, name, filetime, size, stormlib.SFileGetLocale(), 0x00000200 | 0x80000000, phandle) then
        return nil
    end
    return setmetatable({ handle = phandle[0] }, wfile)
end

function archive:load_file(name)
    if self.handle == 0 then
        return nil
    end
    local file = self:open_file(name)
    if not file then
        return nil
    end
    local content = file:read()
    file:close()
    return content
end

function archive:save_file(name, buf, filetime)
    if self.handle == 0 then
        return false
    end
    local file = self:create_file(name, #buf, filetime)
    if not file then
        return false
    end
    file:write(buf)
    file:close()
    return true
end

function archive:number_of_files()
    if self.handle == 0 then
        return 0
    end
    local pinfo = ffi.new('uint32_t[1]', 0)
    if not stormlib.SFileGetFileInfo(self.handle, SFileMpqNumberOfFiles, pinfo, 4, nil) then
        return 0
    end
    return pinfo[0]
end

local m = {}
function m.open(path, readonly, filecount)
    local wpath = uni.u2w(path:string())
    local phandle = ffi.new('uint32_t[1]', 0)
    local flag = 0
    if readonly then
        flag = 0x100
    end
    if not stormlib.SFileOpenArchive(wpath, 0, flag, phandle) then
        return nil
    end
    if filecount then
        stormlib.SFileSetMaxFileCount(phandle[0], filecount)
    end
    return setmetatable({ handle = phandle[0] }, archive)
end
function m.create(path, filecount, encrypt)
    local wpath = uni.u2w(path:string())
    local phandle = ffi.new('uint32_t[1]', 0)
    local info = ffi.new('struct SFILE_CREATE_MPQ')
    info.cbSize = ffi.sizeof('struct SFILE_CREATE_MPQ')
    info.dwMpqVersion   = 0 --MPQ_FORMAT_VERSION_1
    info.pvUserData     = nil
    info.cbUserData     = 0
    info.dwStreamFlags  = 0 --STREAM_PROVIDER_FLAT | BASE_PROVIDER_FILE
    if encrypt then
        info.dwFileFlags1   = 0
        info.dwFileFlags2   = 0
        info.dwFileFlags3   = 0
    else
        info.dwFileFlags1   = 0x80000000 --MPQ_FILE_EXISTS
        info.dwFileFlags2   = 0x80000000 --MPQ_FILE_EXISTS
        info.dwFileFlags3   = 0x80000000 --MPQ_FILE_EXISTS
    end
    info.dwAttrFlags    = 7 --MPQ_ATTRIBUTE_CRC32 | MPQ_ATTRIBUTE_FILETIME | MPQ_ATTRIBUTE_MD5
    info.dwSectorSize   = 0x10000
    info.dwRawChunkSize = 0
    info.dwMaxFileCount = filecount
    if not stormlib.SFileCreateArchive2(wpath, info, phandle) then
        return nil
    end
    return setmetatable({ handle = phandle[0] }, archive)
end
function m.attach(handle)
    return setmetatable({ handle = handle }, archive)
end
return m