#include "mm/memory.h"
#include <stddef.h>
#include <stdint.h>

static inline int ok(int m, int need, int size)
{
    return (m >= 0 && need >= 0 && m <= size && size - m >= need);
}

static inline uint32_t read_color(unsigned char *pal, int idx, int bpp)
{
    // palette entry size = bpp/8
    int step = bpp >> 3;

    unsigned char b = pal[idx * step + 0];
    unsigned char g = pal[idx * step + 1];
    unsigned char r = pal[idx * step + 2];
    unsigned char a = (step == 4) ? pal[idx * step + 3] : 0xFF;

    return ((uint32_t)a << 24) |
           ((uint32_t)r << 16) |
           ((uint32_t)g << 8)  |
           (uint32_t)b;
}

unsigned int *tga_parse(unsigned char *ptr, int size)
{
    if (!ptr || size < 18) return NULL;

    int m = 0;

    int id_len   = ptr[0];
    int cmap_type = ptr[1];
    int img_type  = ptr[2];

    int cmap_start = ptr[3] | (ptr[4] << 8);
    int cmap_len   = ptr[5] | (ptr[6] << 8);
    int cmap_bpp   = ptr[7];

    int w = ptr[12] | (ptr[13] << 8);
    int h = ptr[14] | (ptr[15] << 8);
    int bpp = ptr[16];
    int desc = ptr[17];

    if (w <= 0 || h <= 0 || w > 4096 || h > 4096)
        return NULL;

    if (!ok(18, id_len, size))
        return NULL;

    m = 18 + id_len;
    if (m > size) return NULL;

    int flip_y = !(desc & 0x20);

    size_t pixels = (size_t)w * (size_t)h;
    if (pixels == 0 || pixels > (SIZE_MAX / sizeof(unsigned int)))
        return NULL;

    unsigned int *out = (unsigned int *)kmalloc((pixels + 2) * sizeof(unsigned int));
    if (!out) return NULL;

    out[0] = w;
    out[1] = h;

    size_t i = 0;

    unsigned char *palette = NULL;
    if (cmap_type == 1)
    {
        // palette starts immediately after header+id
        palette = ptr + m;

        size_t pal_size = (size_t)cmap_len * (cmap_bpp >> 3);
        if (!ok(m, pal_size, size)) {
            kfree(out);
            return NULL;
        }

        m += pal_size;
    }

    // =========================
    // TYPE 2 (TRUECOLOR)
    // =========================
    if (img_type == 2)
    {
        int bytespp = bpp >> 3;
        if (bytespp != 3 && bytespp != 4) goto fail;

        for (int y = 0; y < h; y++)
        {
            int row = flip_y ? (h - 1 - y) : y;

            for (int x = 0; x < w; x++)
            {
                if (!ok(m, bytespp, size)) goto fail;

                unsigned char b = ptr[m + 0];
                unsigned char g = ptr[m + 1];
                unsigned char r = ptr[m + 2];
                unsigned char a = (bytespp == 4) ? ptr[m + 3] : 0xFF;

                m += bytespp;

                out[2 + (size_t)row * w + x] =
                    ((uint32_t)a << 24) |
                    ((uint32_t)r << 16) |
                    ((uint32_t)g << 8)  |
                    (uint32_t)b;
            }
        }
    }

    // =========================
    // TYPE 3 (GRAYSCALE)
    // =========================
    else if (img_type == 3)
    {
        for (size_t idx = 0; idx < pixels; idx++)
        {
            if (!ok(m, 1, size)) goto fail;

            unsigned char v = ptr[m++];

            out[2 + idx] =
                0xFF000000 |
                (v << 16) |
                (v << 8)  |
                v;
        }
    }

    // =========================
    // TYPE 1 (PALETTED)
    // =========================
    else if (img_type == 1)
    {
        if (!palette) goto fail;

        for (size_t idx = 0; idx < pixels; idx++)
        {
            if (!ok(m, 1, size)) goto fail;

            unsigned char index = ptr[m++];

            out[2 + idx] = read_color(palette, index, cmap_bpp);
        }
    }

    // =========================
    // TYPE 9/10/11 (RLE)
    // =========================
    else if (img_type == 9 || img_type == 10 || img_type == 11)
    {
        int bytespp = (img_type == 9) ? 1 : (bpp >> 3);
        if (img_type != 9 && bytespp != 3 && bytespp != 4) goto fail;

        while (i < pixels)
        {
            if (!ok(m, 1, size)) goto fail;

            unsigned char packet = ptr[m++];
            int count = (packet & 0x7F) + 1;

            if (count > (int)(pixels - i))
                count = (int)(pixels - i);

            if (packet & 0x80)
            {
                // RLE
                if (!ok(m, bytespp, size)) goto fail;

                uint32_t color;

                if (img_type == 9)
                {
                    color = read_color(palette, ptr[m], cmap_bpp);
                    m += 1;
                }
                else if (img_type == 11)
                {
                    unsigned char v = ptr[m++];
                    color = 0xFF000000 | (v << 16) | (v << 8) | v;
                }
                else
                {
                    unsigned char b = ptr[m + 0];
                    unsigned char g = ptr[m + 1];
                    unsigned char r = ptr[m + 2];
                    unsigned char a = (bytespp == 4) ? ptr[m + 3] : 0xFF;
                    m += bytespp;

                    color =
                        ((uint32_t)a << 24) |
                        ((uint32_t)r << 16) |
                        ((uint32_t)g << 8)  |
                        (uint32_t)b;
                }

                for (int k = 0; k < count; k++)
                    out[2 + i++] = color;
            }
            else
            {
                // RAW
                for (int k = 0; k < count; k++)
                {
                    if (!ok(m, bytespp, size)) goto fail;

                    uint32_t color;

                    if (img_type == 9)
                    {
                        color = read_color(palette, ptr[m++], cmap_bpp);
                    }
                    else if (img_type == 11)
                    {
                        unsigned char v = ptr[m++];
                        color = 0xFF000000 | (v << 16) | (v << 8) | v;
                    }
                    else
                    {
                        unsigned char b = ptr[m + 0];
                        unsigned char g = ptr[m + 1];
                        unsigned char r = ptr[m + 2];
                        unsigned char a = (bytespp == 4) ? ptr[m + 3] : 0xFF;

                        m += bytespp;

                        color =
                            ((uint32_t)a << 24) |
                            ((uint32_t)r << 16) |
                            ((uint32_t)g << 8)  |
                            (uint32_t)b;
                    }

                    out[2 + i++] = color;
                }
            }
        }
    }
    else
    {
        goto fail;
    }

    return out;

fail:
    kfree(out);
    return NULL;
}