package blurhash

import "core:strings"
import "core:math"

@(rodata, private)
ENC_TABLE := [83]byte {
    '0', '1', '2', '3', '4', '5', '6', '7', '8', '9',
    'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J',
    'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T',
    'U', 'V', 'W', 'X', 'Y', 'Z', 'a', 'b', 'c', 'd',
    'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n',
    'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x',
    'y', 'z', '#', '$', '%', '*', '+', ',', '-', '.',
    ':', ';', '=', '?', '@', '[', ']', '^', '_', '{',
    '|', '}', '~',
}

@(rodata, private)
DEC_TABLE := [128]int {
    -1, -1, -1, -1, -1, -1, -1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1,
    -1, -1, -1, 62, 63, 64, -1, -1,
    -1, -1, 65, 66, 67, 68, 69, -1,
    0, 1, 2, 3, 4, 5, 6, 7,
    8, 9, 70, 71, -1, 72, -1, 73,
    74, 10, 11, 12, 13, 14, 15, 16,
    17, 18, 19, 20, 21, 22, 23, 24,
    25, 26, 27, 28, 29, 30, 31, 32,
    33, 34, 35, 75, -1, 76, 77, 78,
    -1, 36, 37, 38, 39, 40, 41, 42,
    43, 44, 45, 46, 47, 48, 49, 50,
    51, 52, 53, 54, 55, 56, 57, 58,
    59, 60, 61, 79, 80, 81, 82, -1,
}

Error :: enum {
    None,

    Invalid_BlurHash,
    Invalid_Component_Count,
}

@(private)
base83_decode_byte :: proc(r: byte) -> int {
    return DEC_TABLE[r]
}

@(private)
base83_decode_str :: proc(str: string) -> int {
    value := 0
    for r in str {
        value = value * 83 + DEC_TABLE[r]
    }
    return value
}

@(private)
base83_decode :: proc{base83_decode_str, base83_decode_byte}

@(private)
base83_encode :: proc(value, length: int, allocator := context.allocator) -> string {
    b := strings.builder_make_none(allocator)

    pow :: proc(base, exp: int) -> int {
        base := base
        exp := exp
        result := 1
        for exp > 0 {
            if exp % 2 == 1 {
                result *= base
            }
            base *= base
            exp /= 2
        }
        return result
    }

    for i in 1..=length {
        digit := (value / pow(83, length - i)) % 83
        strings.write_byte(&b, ENC_TABLE[digit])
    }

    return strings.to_string(b)
}

@(private)
is_valid_blurhash :: proc(blurhash: string) -> bool {
    if len(blurhash) < 6 {
        return false
    }

    size_flag := base83_decode(blurhash[0])
    num_y := (size_flag / 9) + 1
    num_x := (size_flag % 9) + 1

    if len(blurhash) != 4 + 2 * num_x * num_y {
        return false
    }

    return true
}

@(private)
linear_to_srgb :: proc(value: f32) -> int {
    v := max(0, min(1, value))
    if v <= 0.0031308 {
        return int(v * 12.92 * 255 + 0.5)
    }

    return int((1.055 * math.pow(v, 1 / 2.4) - 0.055) * 255 + 0.5)
}

@(private)
srgb_to_linear :: proc(value: int) -> f32 {
    v := f32(value) / 255.0
    if v <= 0.04045 {
        return v / 12.92
    }

    return math.pow((v + 0.055) / 1.055, 2.4)
}

@(private)
decode_dc :: proc(value: int) -> (r, g, b: f32) {
    return srgb_to_linear(value >> 16), srgb_to_linear((value >> 8) & 0xFF), srgb_to_linear(value & 0xFF)
}

@(private)
sign_pow :: proc(value, exp: f32) -> f32 {
    return math.copy_sign(math.pow(abs(value), exp), value)
}

@(private)
decode_ac :: proc(value: int, max_value: f32) -> (r, g, b: f32) {
    quant_r := int(math.floor(f32(value) / (19 * 19)))
    quant_g := int(math.floor(f32(value) / 19)) % 19
    quant_b := value % 19

    r = sign_pow((f32(quant_r) - 9) / 9, 2.0) * max_value
    g = sign_pow((f32(quant_g) - 9) / 9, 2.0) * max_value
    b = sign_pow((f32(quant_b) - 9) / 9, 2.0) * max_value

    return
}

decode :: proc(blurhash: string, width, height: int, punch := 1.0, allocator := context.allocator) -> ([]byte, Error) {
    context.allocator = allocator

    if !is_valid_blurhash(blurhash) {
        return nil, .Invalid_BlurHash
    }
    punch := punch
    if punch < 1 {
        punch = 1
    }

    size_flag := base83_decode(blurhash[0])
    num_y := (size_flag / 9) + 1
    num_x := (size_flag % 9) + 1

    r, g, b := 0, 0, 0

    quantized_max_value := base83_decode(blurhash[1])
    max_value := f32(quantized_max_value + 1) / 166.0

    colors_size := num_x * num_y
    colors := make([][3]f32, colors_size)

    for iter in 0..<colors_size {
        if iter == 0 {
            value := base83_decode(blurhash[2:6])
            if value == -1 {
                return nil, .Invalid_BlurHash
            }

            colors[iter] = compress_values(decode_dc(value))
        } else {
            value := base83_decode(blurhash[4 + iter * 2:6 + iter * 2])
            if value == -1 {
                return nil, .Invalid_BlurHash
            }

            colors[iter] = compress_values(decode_ac(value, max_value * f32(punch)))
        }
    }

    CHANNELS :: 3
    bytes_per_row := width * CHANNELS
    pixels := make([]byte, height * width * CHANNELS)

    for y in 0..<height {
        for x in 0..<width {
            r, g, b: f32 = 0, 0, 0

            for j in 0..<num_y {
                for i in 0..<num_x {
                    basics := math.cos((math.PI * f32(x) * f32(i)) / f32(width)) * math.cos((math.PI * f32(y) * f32(j)) / f32(height))
                    idx := i + j * num_x
                    r += colors[idx][0] * basics
                    g += colors[idx][1] * basics
                    b += colors[idx][2] * basics
                }
            }

            int_r := linear_to_srgb(r)
            int_g := linear_to_srgb(g)
            int_b := linear_to_srgb(b)

            pixels[CHANNELS * x + 0 + y * bytes_per_row] = u8(clamp(int_r, 0, 255))
            pixels[CHANNELS * x + 1 + y * bytes_per_row] = u8(clamp(int_g, 0, 255))
            pixels[CHANNELS * x + 2 + y * bytes_per_row] = u8(clamp(int_b, 0, 255))
        }
    }

    return pixels, nil
}

@(private)
multiply_basis_function :: proc(x_component, y_component, width, height: int, pixels: []byte, bytes_per_row: int) -> (r, g, b: f32) {
    normalization := x_component == 0 && y_component == 0 ? 1 : 2

    for y in 0..<height {
        for x in 0..<width {
            basis := math.cos(math.PI * f32(x_component) * f32(x) / f32(width)) * math.cos(math.PI * f32(y_component) * f32(y) / f32(height))
            r += basis * srgb_to_linear(int(pixels[3 * x + 0 + y * bytes_per_row]))
            g += basis * srgb_to_linear(int(pixels[3 * x + 1 + y * bytes_per_row]))
            b += basis * srgb_to_linear(int(pixels[3 * x + 2 + y * bytes_per_row]))
        }
    }

    scale := f32(normalization) / f32(width * height)

    return r*scale,g*scale,b*scale
}

@(private)
encode_dc :: proc(dc: [3]f32) -> int {
    rounded_r := linear_to_srgb(dc.r)
    rounded_g := linear_to_srgb(dc.g)
    rounded_b := linear_to_srgb(dc.b)

    return (rounded_r << 16) + (rounded_g << 8) + rounded_b
}

@(private)
encode_ac :: proc(ac: [3]f32, max_value: f32) -> int {
    quant_r := int(max(0, min(18, math.floor(sign_pow(ac.r / max_value, 0.5) * 9 + 9.5))))
    quant_g := int(max(0, min(18, math.floor(sign_pow(ac.g / max_value, 0.5) * 9 + 9.5))))
    quant_b := int(max(0, min(18, math.floor(sign_pow(ac.b / max_value, 0.5) * 9 + 9.5))))

    return quant_r * 19 * 19 + quant_g * 19 + quant_b
}

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

    if x_components < 1 || x_components > 9 {
        return "", .Invalid_Component_Count
    }
    if y_components < 1 || y_components > 9 {
        return "", .Invalid_Component_Count
    }

    builder := strings.builder_make_len_cap(0, 2 + 4 + (9 * 9 - 1) * 2)
    factors := make([dynamic][3]f32, 0, x_components * y_components)

    for y in 0..<y_components {
        for x in 0..<x_components {
            r, g, b := multiply_basis_function(x, y, width, height, pixels, bytes_per_row)
            append(&factors, [3]f32{r, g, b})
        }
    }

    dc := factors[0]
    ac := factors[1:]

    size_flag := (x_components - 1) + (y_components - 1) * 9
    strings.write_string(&builder, base83_encode(size_flag, 1))

    max_value: f32

    if len(ac) > 0 {
        actual_max_value: f32 = 0
        for v in ac {
            actual_max_value = max(actual_max_value, max(abs(v[0]), abs(v[1]), abs(v[2])))
        }

        quantized_max_value := int(max(0, min(82, math.floor(actual_max_value * 166 - 0.5))))
        max_value = f32(quantized_max_value + 1) / 166.0
        strings.write_string(&builder, base83_encode(quantized_max_value, 1))
    } else {
        max_value = 1
        strings.write_string(&builder, base83_encode(0, 1))
    }

    strings.write_string(&builder, base83_encode(encode_dc(dc), 4))

    for factor in ac {
        strings.write_string(&builder, base83_encode(encode_ac(factor, max_value), 2))
    }

    return strings.to_string(builder), nil
}