#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>
#include <zlib.h>

#include "tldat.h"

static const uint64_t KEYS[] =
{
    0x551b8a2ecd6197efLL,
    0x186a57f394700e34LL,
    0x3e9f2e302712b938LL,
    0xe14c2303ccc551f2LL,
    0xccf38ca1f5c17133LL,
    0x2353622f23b1c9dbLL,
    0x34afadac84ae7417LL,
    0xa5dcaca1d9365ebLL,
    0xf262becf99cd3c0fLL,
    0x3125b4b2f481962LL,
    0xcd5ec4039782a7aaLL,
    0x7e33b2fc317e77f3LL,
    0xbefb5409bb40d4faLL,
    0x3368c49e410e24efLL,
    0x1e9693617e3e6bbfLL,
    0x9c35278e358912b1LL
};

static void tldecrypt(uint64_t key, uint8_t *data, uint64_t size)
{
    uint64_t tmp, t;
    unsigned i;

    for (i = 0; i < ((size / 16) * 16); i += 8)
    {
        *(uint64_t *)(data + i) ^= key;
        tmp = key ^ 0x4e3362bf7a4c7c26LL;
        tmp ^= tmp << 13;
        tmp ^= tmp >> 7;
        t = ((int)(data[i + 7]) << ((((i - 1) & 7) << 3)) % 32) |
            ((int)(data[i + 4]) << ((((i - 4) & 7) << 3)) % 32) |
            ((int)(data[i + 1]) << ((((i + 1) & 7) << 3)) % 32) |
            ((int)(data[i + 5]) << ((((i - 3) & 7) << 3)) % 32) |
            ((int)(data[i]) << (((i & 7) << 3)) % 32) |
            ((int)(data[i + 2]) << ((((i + 2) & 7) << 3)) % 32) |
            ((int)(data[i + 6]) << ((((i - 2) & 7) << 3)) % 32) |
            ((int)(data[i + 3]) << ((((i + 3) & 7) << 3)) % 32);
        key = tmp ^ (tmp << 17) ^ t;
    }

    for (; i < size; i++)
        data[i] ^= key >> ((i & 7) << 3);
}

tldat_t * tldat_init(const char *path)
{
    struct stat st;
    FILE *hda = NULL, *hdb = NULL;
    tldat_t *tldat = NULL;

    if (tldat = (tldat_t *)malloc(sizeof(tldat_t)))
    {
        memset(tldat, 0, sizeof(tldat_t));
        tldat->path = (char *)calloc(strlen(path)+1, sizeof(char));
        strcpy(tldat->path, path);
    }
    else
        return NULL;

    char *hdapath = NULL, *hdbpath = NULL;
    if (stat(path, &st) == 0)
    {
        if (!S_ISDIR(st.st_mode))
        {
            tldat->err = ERR_INV_PATH;
            return tldat;
        }

        hdapath = (char *)calloc(strlen(path)+24, sizeof(char));
        strcpy(hdapath, path);
        strcat(hdapath, "/FILEHEADER.TOFHDA");
        hdbpath = (char *)calloc(strlen(path)+24, sizeof(char));
        strcpy(hdbpath, path);
        strcat(hdbpath, "/FILEHEADER.TOFHDB");
    }
    else
    {
        tldat->err = ERR_INV_PATH;
        return tldat;
    }

    uint64_t *hda_keys = NULL;
    if (hda = fopen(hdapath, "rb"))
    {
        stat(hdapath, &st);
        hda_keys = (uint64_t *)calloc(st.st_size/8, sizeof(uint64_t));
        fread(hda_keys, 8, st.st_size/8, hda);
        fclose(hda);
        free(hdapath);
    }
    else
    {
        tldat->err = ERR_FNF_HDA;
        free(hdapath);
        return tldat;
    }

    uint64_t key1 = hda_keys[12] & 15;
    uint64_t key2 = (hda_keys[4] >> 4) & 15;
    free(hda_keys);
    if (key1 == key2) key2 = (key2 + 1) & 15;
    key1 = KEYS[key1];
    key2 = KEYS[key2];
    key1 ^= key2;

    uint8_t * hdb_buf = NULL;
    if (hdb = fopen(hdbpath, "rb"))
    {
        stat(hdbpath, &st);
        hdb_buf = (uint8_t *)malloc(st.st_size);
        fread(hdb_buf, 1, st.st_size, hdb);
        tldecrypt(key1, hdb_buf, st.st_size);
        fclose(hdb);
        free(hdbpath);
    }
    else
    {
        tldat->err = ERR_FNF_HDB;
        free(hdbpath);
        return tldat;
    }

    memmove(&tldat->entries_offs, &hdb_buf[8], 8); // first 8 bytes are padding
    memmove(&tldat->nentries, &hdb_buf[16], 8);
    memmove(&tldat->info_offs, &hdb_buf[40], 8); // 16 bytes of padding
    memmove(&tldat->nfiles, &hdb_buf[48], 8);

    uint8_t *pos = &hdb_buf[tldat->entries_offs];
    tldat->crcs = (uint32_t *)calloc(tldat->nfiles, sizeof(uint32_t));
    tldat->keys = (uint32_t *)calloc(tldat->nfiles, sizeof(uint32_t));
    tldat->recs = (tldat_rec_t *)calloc(tldat->nfiles, sizeof(tldat_rec_t));
    for (uint64_t i = 0; i < tldat->nfiles; i++)
    {
        tldat->crcs[i] = *(uint32_t *)pos;
        pos += 4;
        tldat->keys[i] = *(uint32_t *)pos;
        pos += 4;
    }

    for (uint64_t i = 0; i < tldat->nfiles; i++)
    {
        memmove(&tldat->recs[i], &hdb_buf[tldat->info_offs], sizeof(tldat_rec_t) - sizeof(uint64_t));

        for (uint64_t x = 0; x < tldat->nentries; x++)
            if (tldat->crcs[x] == tldat->recs[i].crc)
            {
                tldat->recs[i].key = tldat->keys[x];
                break;
            }

        tldat->recs[i].key = ((3 - (tldat->recs[i].key & 3)) << 2) |
            (tldat->recs[i].key & 0xfffffffffffffff0LL) |
            ((tldat->recs[i].key >> 2) & 3);
    }

    free(hdb_buf);
    return tldat;
}

int tldat_extract(tldat_t *tldat, uint64_t idx, const char *outpath)
{
    struct stat st;
    if (stat(outpath, &st) == 0)
        if (!S_ISDIR(st.st_mode))
        {
            tldat->err = ERR_INV_OUTDIR;
            return tldat->err;
        }
    else
    {
        tldat->err = ERR_INV_OUTDIR;
        return tldat->err;
    }

    uint8_t *buf = (uint8_t *)calloc(tldat->recs[idx].zsize, sizeof(char));
    FILE *dat = NULL, *outfile = NULL;
    char *datpath = (char *)calloc(strlen(tldat->path)+32, sizeof(char));
    strcpy(datpath, tldat->path);
    strcat(datpath, "/TLFILE.TLDAT");
    if (dat = fopen(datpath, "rb"))
    {
        fseek(dat, tldat->recs[idx].offs, SEEK_SET);
        fread(buf, 1, tldat->recs[idx].zsize, dat);
        fclose(dat);
        tldecrypt(tldat->recs[idx].key, buf, tldat->recs[idx].zsize);

        char *outfilepath = (char *)calloc(strlen(outpath)+10, sizeof(char));
        strcpy(outfilepath, outpath);
        strcat(outfilepath, "/");
        if (tldat->recs[idx].size == tldat->recs[idx].zsize) // uncompressed?
        {
            strncat(outfilepath, tldat->recs[idx].type, 8);
            if (outfile = fopen(outfilepath, "wb+x"))
            {
                fwrite(buf, 1, tldat->recs[idx].zsize, outfile);
                fclose(outfile);
            }
            else
                tldat->err = ERR_WRITE;
        }
        else
        {
            uint32_t tlzc;
            memmove(&tlzc, buf, 4);

            if (tlzc == 0x435a4c54) // 'TLZC'
            {
                uint32_t tver, tzsize, tsize;
                memmove(&tver, &buf[4], 4);
                memmove(&tzsize, &buf[8], 4);
                memmove(&tsize, &buf[12], 4);

                if (tver == 0x201) // Abyss
                {
                    char dummy[8];
                    z_stream zs;

                    memmove(dummy, &buf[16], 8);
                    memset(&zs, 0, sizeof(z_stream));
                    tzsize -= 24;
                    if ((tldat->err = inflateInit(&zs)) == NO_ERROR)
                    {
                        uint8_t *inbuf = (uint8_t *)calloc(tzsize, sizeof(char));
                        uint8_t *outbuf = (uint8_t *)calloc(tsize, sizeof(char));

                        memmove(inbuf, &buf[24], tzsize);
                        zs.avail_in = tzsize;
                        zs.next_in = inbuf;
                        zs.avail_out = tsize;
                        zs.next_out = outbuf;
                        if ((tldat->err = inflate(&zs, Z_NO_FLUSH)) == Z_STREAM_END)
                        {
                            strncat(outfilepath, dummy, 8);
                            if (outfile = fopen(outfilepath, "wb+x"))
                            {
                                fwrite(outbuf, 1, tsize, outfile);
                                fclose(outfile);
                            }
                            else
                                tldat->err = ERR_WRITE;
                        }

                        inflateEnd(&zs);
                        free(inbuf);
                        free(outbuf);
                    }
                }
                else if (tver == 0x401) // Symphonia, unsupported, unknown lzma and encryption
                {
                    char dummy[8];

                    memmove(dummy, &buf[16], 8);
                    tzsize -= 24;

                    strncat(outfilepath, dummy, 8);
                    if (outfile = fopen(outfilepath, "wb+x"))
                    {
                        fwrite(buf, 1, tldat->recs[idx].zsize, outfile);
                        fclose(outfile);
                    }
                    else
                        tldat->err = ERR_WRITE;
                }
                else
                {
                    char dummy[9];
                    uint32_t chunk_size, nchunks;
                }
            }
            else
            {
                strcat(outfilepath, "unsupported-");
                strncat(outfilepath, tldat->recs[idx].type, 8);
                if (outfile = fopen(outfilepath, "wb+x"))
                {
                    fwrite(buf, 1, tldat->recs[idx].size, outfile);
                    fclose(outfile);
                }
                else
                    tldat->err = ERR_WRITE;
            }
        }

        free(outfilepath);
        free(buf);
    }
    else
        tldat->err = ERR_FNF_TLDAT;
    free(datpath);
}

void tldat_free(tldat_t *tldat)
{
    if (tldat)
    {
        free(tldat->crcs);
        free(tldat->keys);
        free(tldat->recs);
        free(tldat->path);
        free(tldat);
    }
}