package thumbhash

import "core:math"

Error :: enum {
    None,
    Dimensions_Too_Large,
    Invalid_RGBA_Image,
    Hash_Too_Short,
}

hash_to_approximate_aspect_ratio :: proc(hash: []byte) -> (ratio: f32, err: Error) {
    if len(hash) < 5 {
        return 0.0, .Hash_Too_Short
    }

    has_alpha := hash[2] & 0x80 != 0
    l_max := has_alpha ? 5 : 7
    l_min := hash[3] & 7
    is_landscape := hash[4] & 0x80 != 0
    lx := is_landscape ? u8(l_max) : l_min
    ly := is_landscape ? l_min : u8(l_max)

    return f32(lx) / f32(ly), nil
}

@(private)
decode_channel :: proc(hash: []byte, nx, ny: uint, scale: f32, ac_start: int, ac_index: ^int) -> [dynamic]f32 {
    ac := make([dynamic]f32, 0, nx * ny)

    for cy in 0..<ny {
        cx: uint = cy > 0 ? 0 : 1

        for cx * ny < nx * (ny - cy) {
            fac := f32((hash[ac_start + (ac_index^ >> 1)] >> uint((ac_index^ & 1) << 2)) & 15)

            append(&ac, (fac / 7.5 - 1.0) * scale)
            ac_index^ += 1
            cx += 1
        }
    }

    return ac
}

decode :: proc(hash: []byte, allocator := context.allocator) -> (w: uint, h: uint, pixels: [dynamic]byte, err: Error) {
    context.allocator = allocator

    ratio := hash_to_approximate_aspect_ratio(hash) or_return

    header24 := cast(u32)hash[0] |
        (cast(u32)hash[1] << 8)  |
        (cast(u32)hash[2] << 16)

    header16 := cast(u16)hash[3] | (cast(u16)hash[4] << 8)
    l_dc := f32(header24 & 63) / 63.0
    p_dc := f32((header24 >> 6) & 63) / 31.5 - 1.0
    q_dc := f32((header24 >> 12) & 63) / 31.5 - 1.0
    l_scale := f32((header24 >> 18) & 31) / 31.0
    has_alpha := (header24 >> 23) != 0
    p_scale := f32((header16 >> 3) & 63) / 63.0
    q_scale := f32((header16 >> 9) & 63) / 63.0
    is_landscape := (header16 >> 15) != 0
    l_max := has_alpha ? 5 : 7
    lx := cast(uint)max(3, is_landscape ? l_max : int(header16 & 7))
    ly := cast(uint)max(3, is_landscape ? int(header16 & 7) : l_max)
    a_dc, a_scale: f32 = 1.0, 1.0
    if has_alpha {
        header8 := hash[5]
        a_dc = f32(header8 & 15) / 15.0
        a_scale = f32(header8 >> 4) / 15.0
    }

    ac_start := has_alpha ? 6 : 5
    ac_index := 0
    l_ac := decode_channel(hash, lx, ly, l_scale, ac_start, &ac_index)
    p_ac := decode_channel(hash, 3, 3, p_scale * 1.25, ac_start, &ac_index)
    q_ac := decode_channel(hash, 3, 3, q_scale * 1.25, ac_start, &ac_index)
    a_ac: [dynamic]f32
    if has_alpha {
        a_ac = decode_channel(hash, 5, 5, a_scale, ac_start, &ac_index)
    }

    defer delete(l_ac)
    defer delete(p_ac)
    defer delete(q_ac)
    defer delete(a_ac)

    width := cast(uint)math.round(ratio > 1.0 ? 32 : 32 * ratio)
    height := cast(uint)math.round(ratio > 1.0 ? 32 / ratio : 32)

    pixels = make([dynamic]byte, 0, width * height * 4)
    fx := []f32{0, 0, 0, 0, 0, 0, 0}
    fy := []f32{0, 0, 0, 0, 0, 0, 0}

    for y in 0..<height {
        for x in 0..<width {
            l := l_dc
            p := p_dc
            q := q_dc
            a := a_dc

            for cx in 0..<max(lx, has_alpha ? 5 : 3) {
                fx[cx] = math.cos(math.PI / f32(width) * (f32(x) + 0.5) * f32(cx))
            }
            for cy in 0..<max(ly, has_alpha ? 5 : 3) {
                fy[cy] = math.cos(math.PI / f32(height) * (f32(y) + 0.5) * f32(cy))
            }

            j := 0
            for cy in 0..<ly {
                cx: uint = cy > 0 ? 0 : 1
                fy2 := fy[cy] * 2.0
                for cx * ly < lx * (ly - cy) {
                    l += l_ac[j] * fx[cx] * fy2
                    j += 1
                    cx += 1
                }
            }

            j = 0
            for cy in 0..<3 {
                cx := cy > 0 ? 0 : 1
                fy2 := fy[cy] * 2.0
                for cx < 3 - cy {
                    f := fx[cx] *fy2
                    p += p_ac[j] * f
                    q += q_ac[j] * f
                    j += 1
                    cx += 1
                }
            }

            if has_alpha {
                j = 0
                for cy in 0..<5 {
                    cx := cy > 0 ? 0 : 1
                    fy2 := fy[cy] * 2.0
                    for cx < 5 - cy {
                        a += a_ac[j] * fx[cx] * fy2
                        j += 1
                        cx += 1
                    }
                }
            }

            b := l - 2.0 / 3.0 * p
            r := (3.0 * l - b + q) / 2.0
            g := r - q

            append(&pixels, byte(clamp(r, 0.0, 1.0) * 255.0))
            append(&pixels, byte(clamp(g, 0.0, 1.0) * 255.0))
            append(&pixels, byte(clamp(b, 0.0, 1.0) * 255.0))
            append(&pixels, byte(clamp(a, 0.0, 1.0) * 255.0))
        }
    }

    return width, height, pixels, nil
}

@(private)
encode_channel :: proc(channel: []f32, nx, ny: uint, width, height: int, allocator := context.allocator) -> (f32, [dynamic]f32, f32) {
    context.allocator = allocator

    dc := f32(0.0)
    ac := make([dynamic]f32, 0, nx * ny / 2)
    scale := f32(0.0)
    fx := make([]f32, width)
    defer delete(fx)

    for cy in 0..<ny {
        cx: uint = 0

        for cx * ny < nx * (ny - cy) {
            f := f32(0.0)

            for x in 0..<width{
                fx[x] = math.cos(math.PI / f32(width) * f32(cx) * (f32(x) + 0.5))
            }

            for y in 0..<height {
                fy := math.cos(math.PI / f32(height) * f32(cy) * (f32(y) + 0.5))
                for x in 0..<width {
                    f += channel[x + y * width] * fx[x] * fy
                }
            }

            f /= f32(width * height)

            if cx > 0 || cy > 0 {
                append(&ac, f)
                scale = max(abs(f), scale)
            } else {
                dc = f
            }

            cx += 1
        }
    }

    if scale > 0.0 {
        for &_ac in ac {
            _ac = 0.5 + 0.5 / scale * _ac
        }
    }

    return dc, ac, scale
}

encode :: proc(width, height: int, pixels: []byte, allocator := context.allocator) -> ([dynamic]byte, Error) {
    context.allocator = allocator

    if width > 100 || height > 100 {
        return nil, .Dimensions_Too_Large
    }

    if len(pixels) != width * height * 4 {
        return nil, .Invalid_RGBA_Image
    }

    avg_r, avg_g, avg_b, avg_a: f32 = 0.0, 0.0, 0.0, 0.0

    j := 0
    for i in 0..<width*height {
        alpha := f32(pixels[j + 3]) / 255.0
        avg_r += alpha / 255.0 * f32(pixels[j])
        avg_g += alpha / 255.0 * f32(pixels[j + 1])
        avg_b += alpha / 255.0 * f32(pixels[j + 2])
        avg_a += alpha

        j += 4
    }

    if avg_a > 0.0 {
        avg_r /= avg_a
        avg_g /= avg_a
        avg_b /= avg_a
    }

    has_alpha := avg_a < f32(width * height)
    l_limit := has_alpha ? 5 : 7
    lx := max(cast(uint)math.round(f32(l_limit * width) / f32(max(width, height))), 1)
    ly := max(cast(uint)math.round(f32(l_limit * height) / f32(max(width, height))), 1)

    l := make([dynamic]f32, 0, width * height)
    p := make([dynamic]f32, 0, width * height)
    q := make([dynamic]f32, 0, width * height)
    a := make([dynamic]f32, 0, width * height)

    defer delete(l)
    defer delete(p)
    defer delete(q)
    defer delete(a)

    for i := 0; i < len(pixels); i += 4 {
        alpha := f32(pixels[i + 3]) / 255.0
        r := avg_r * (1.0 - alpha) + alpha / 255.0 * f32(pixels[i])
        g := avg_g * (1.0 - alpha) + alpha / 255.0 * f32(pixels[i + 1])
        b := avg_b * (1.0 - alpha) + alpha / 255.0 * f32(pixels[i + 2])
        append(&l, (r + g + b) / 3.0)
        append(&p, (r + g) / 2.0 - b)
        append(&q, r - g)
        append(&a, alpha)
    }

    l_dc, l_ac, l_scale := encode_channel(l[:], max(lx, 3), max(ly, 3), width, height)
    p_dc, p_ac, p_scale := encode_channel(p[:], 3, 3, width, height)
    q_dc, q_ac, q_scale := encode_channel(q[:], 3, 3, width, height)
    a_dc, a_ac, a_scale := f32(1.0), [dynamic]f32{}, f32(1.0)

    defer delete(l_ac)
    defer delete(p_ac)
    defer delete(q_ac)
    defer delete(a_ac)

    if has_alpha {
        a_dc, a_ac, a_scale = encode_channel(a[:], 5, 5, width, height)
    }

    is_landscape := width > height
    header24: u32 = cast(u32)math.round(63.0 * l_dc)      | 
        ((cast(u32)math.round(31.5 + 31.5 * p_dc)) << 6)  |
        ((cast(u32)math.round(31.5 + 31.5 * q_dc)) << 12) |
        ((cast(u32)math.round(31.0 * l_scale)) << 18)

    if has_alpha {
        header24 |= 1 << 23
    }

    header16: u16 = cast(u16)(is_landscape ? ly : lx) |
        ((cast(u16)math.round(63.0 * p_scale)) << 3)  |
        ((cast(u16)math.round(63.0 * q_scale)) << 9)

    if is_landscape {
        header16 |= 1 << 15
    }

    hash := make([dynamic]byte, 0, 25)
    append(&hash, byte(header24 & 0xFF))
    append(&hash, byte((header24 >> 8) & 0xFF))
    append(&hash, byte(header24 >> 16))
    append(&hash, byte(header16 & 0xFF))
    append(&hash, byte(header16 >> 8))

    is_odd := false
    if has_alpha {
        append(&hash, cast(byte)math.round(15.0 * a_dc) | ((cast(byte)math.round(15.0 * a_scale)) << 4))
    }

    for ac in ([][]f32{l_ac[:], p_ac[:], q_ac[:]}) {
        for f in ac {
            u := cast(byte)math.round(15.0 * f)
            if is_odd {
                hash[len(hash) - 1] |= u << 4
            } else {
                append(&hash, u)
            }
            is_odd = !is_odd
        }
    }

    if has_alpha {
        for f in a_ac {
            u := cast(byte)math.round(15.0 * f)
            if is_odd {
                hash[len(hash) - 1] |= u << 4
            } else {
                append(&hash, u)
            }
            is_odd = !is_odd
        }
    }

    return hash, nil
}