package main

import "core:image"
import "core:image/png"
import "core:image/qoi"
import "core:os/os2"
import "core:fmt"
import "core:slice"
import "core:mem"

import thumbhash "../"

decode_example :: proc() {
    t: mem.Tracking_Allocator
    mem.tracking_allocator_init(&t, context.allocator)
    alloc := mem.tracking_allocator(&t)

    // The hardcoded hash is that of the flower.png file
    w, h, pixels, err := thumbhash.decode({0x93, 0x4A, 0x6, 0x2D, 0x6, 0x92, 0x56, 0xC3, 0x74, 0x5, 0x58, 0x67, 0xDA, 0x8A, 0xB6, 0x67, 0x94, 0x90, 0x51, 0x7, 0x19}, alloc)
    if err != nil {
        fmt.eprintln(err)
        return
    }

    img, ok := image.pixels_to_image(slice.reinterpret([]image.RGBA_Pixel, pixels[:]), cast(int)w, cast(int)h)
    if !ok {
        fmt.eprintln("Couldn't create an image")
        return
    }

    save_err := qoi.save_to_file("yup.qoi", &img)
    if save_err != nil {
        fmt.eprintln(save_err)
        return
    }

    delete(pixels)

    for _, leak in t.allocation_map {
        fmt.printfln("%v leaked %m", leak.location, leak.size)
    }
    for bad_free in t.bad_free_array {
        fmt.printfln("%v allocation %p was freed badly", bad_free.location, bad_free.memory)
    }
    fmt.printfln("Peak Mem: %m", t.peak_memory_allocated)
    fmt.printfln("Total Mem: %m", t.total_memory_allocated)
    fmt.printfln("Current Mem: %m", t.current_memory_allocated)
    fmt.printfln("Total Allocs: %v, Total Frees: %v:", t.total_allocation_count, t.total_free_count)
    fmt.printfln("Total Mem Free'd: %m", t.total_memory_freed)
}

convert_rgb_to_rgba :: proc(rgb_data: []u8) -> []u8 {
    assert(len(rgb_data) % 3 == 0, "RGB data length must be a multiple of 3")

    rgba_data := make([dynamic]u8, 0, (len(rgb_data) / 3) * 4)

    for i := 0; i < len(rgb_data); i += 3 {
        r := rgb_data[i + 0]
        g := rgb_data[i + 1]
        b := rgb_data[i + 2]
        a := u8(255) // Fully opaque

        append(&rgba_data, r)
        append(&rgba_data, g)
        append(&rgba_data, b)
        append(&rgba_data, a)
    }

    return rgba_data[:]
}

encode_example :: proc() {
    if len(os2.args) < 3 {
        fmt.eprintln("Please provide a path to a png")
        return
    }

    img, err := image.load_from_file(os2.args[2])
    if err != nil {
        fmt.eprintln(err)
        return
    }

    buf := img.pixels.buf[:]
    if img.channels == 3 {
        buf = convert_rgb_to_rgba(img.pixels.buf[:])
    }

    t: mem.Tracking_Allocator
    mem.tracking_allocator_init(&t, context.allocator)
    alloc := mem.tracking_allocator(&t)

    thash, hash_err := thumbhash.encode(img.width, img.height, buf, alloc)
    defer delete(thash)
    if hash_err != nil {
        fmt.eprintln(hash_err)
        return
    }

    fmt.printfln("Hash: %X", thash)

    delete(thash)

    for _, leak in t.allocation_map {
        fmt.printfln("%v leaked %m", leak.location, leak.size)
    }
    for bad_free in t.bad_free_array {
        fmt.printfln("%v allocation %p was freed badly", bad_free.location, bad_free.memory)
    }
    fmt.printfln("Peak Mem: %m", t.peak_memory_allocated)
    fmt.printfln("Total Mem: %m", t.total_memory_allocated)
    fmt.printfln("Current Mem: %m", t.current_memory_allocated)
    fmt.printfln("Total Allocs: %v, Total Frees: %v:", t.total_allocation_count, t.total_free_count)
    fmt.printfln("Total Mem Free'd: %m", t.total_memory_freed)
}

main :: proc() {
    if len(os2.args) < 2 {
        fmt.eprintln("Usage: program \"decode|encode\" <arg>")
        return
    }

    if os2.args[1] == "decode" {
        decode_example()
    } else if os2.args[1] == "encode" {
        encode_example()
    } else {
        fmt.eprintln("Unknown command")
    }
}