package png_to_qoi_util

import "core:os"
import "core:fmt"
import "core:mem"
import "core:nbio"
import "core:time"
import "core:bytes"
import "core:thread"
import "core:strings"
import "core:image"
import "core:image/png"
import "core:image/qoi"
import si "core:sys/info"

FileDesc :: struct {
	handle: nbio.Handle,
	src_path, dst_path: string,
	png_data, qoi_data: []u8,
	qoi_bfr: bytes.Buffer,
	read_ok, write_ok: bool,
}

files : [dynamic]FileDesc

thread_pool: thread.Pool
io_loop: ^nbio.Event_Loop
io_pending, io_succeeded, io_failed: int

open_dir :: proc(name: string) {
	fi, err := os.read_all_directory_by_path(name, context.allocator)
	defer os.file_info_slice_delete(fi, context.allocator)

	if err != nil {
		fmt.printf("err %v: failed to read dir %v\n", err, name)
		os.exit(1)
	}

	for file in fi {
		if file.type == .Directory {
			open_dir(file.fullpath)
			continue
		}

		if os.ext(file.name) == ".png" {
			src := strings.clone(file.fullpath)
			dst, _:= strings.replace(file.fullpath,
					".png", ".qoi", 1)
			append(&files, FileDesc{
				src_path = src,
				dst_path = dst,
			})
		}
	}
}

main :: proc() {
	sw: time.Stopwatch
	time.stopwatch_start(&sw)

	root := "."
	if len(os.args) > 1 {
		assert(os.is_dir(os.args[1]), "input must be dir")
		root = os.args[1]
	}

	open_dir(root)
	defer delete(files)

	fmt.printfln("found %d png files", len(files))

	if len(files) == 0 {
		fmt.println("no png files found")
		return
	}

	core_cnt, thread_cnt, ok := si.cpu_core_count()
	if !ok {
		fmt.printf("err: Could not retrieve CPU info.\n")
		os.exit(1)
	}

	if thread_cnt < 1 {
		thread_cnt = 1
	}

	if thread_cnt > len(files) {
		thread_cnt = len(files)
	}

	thread.pool_init(&thread_pool, context.allocator, thread_cnt)
	thread.pool_start(&thread_pool)

	assert(nbio.acquire_thread_event_loop() == nil)
	defer nbio.release_thread_event_loop()
	io_loop = nbio.current_thread_event_loop()
	io_pending = len(files)

	for &file in files {
		nbio.read_entire_file(
			file.src_path,
			&file,
			read_cb,
			context.allocator,
			l = io_loop
		)
	}

	for io_pending > 0 {
		job_err := nbio.tick(50 * time.Millisecond)
		if job_err != .None {
			fmt.println("nbio tick failed:", job_err)
			break
		}
	}

	thread.pool_finish(&thread_pool)
	thread.pool_destroy(&thread_pool)

	time.stopwatch_stop(&sw)
	fmt.println(io_succeeded, "succeeded")
	fmt.println(io_failed, "failed")
	fmt.printfln("Total time: %.3fs",
		time.duration_seconds(time.stopwatch_duration(sw)))
}

read_cb :: proc(
	user_data: rawptr, data: []byte, err: nbio.Read_Entire_File_Error) {
	file := cast(^FileDesc)user_data

	if err.value != nil {
		fmt.printf("io read err: %v\n", err)
		file.read_ok = false
		file_load_done(file, false)
		return
	}

	file.png_data = data
	file.read_ok = true
	fmt.printfln("read complete: %q bytes=%d err=%v",
		file.src_path, len(data), err.value)

	thread.pool_add_task(
		&thread_pool,
		context.allocator,
		png_to_qoi_task,
		file
	)
}

png_to_qoi_task :: proc(task: thread.Task) {
	file := cast(^FileDesc)task.data

	png_img, png_err := 
			png.load_from_bytes(
				file.png_data,
				{.alpha_add_if_missing,},
				context.allocator
			)

	if png_err != nil {
		fmt.eprintln("png", file.src_path, "failed:", png_err)
		nbio.next_tick_poly(file, on_encode_failed, l = io_loop)
		return
	}
	defer png.destroy(png_img)

	if !normalize_for_qoi(png_img) {
		fmt.eprintfln(
			"qoi normalize failed:" + 
			"%q width=%d height=%d channels=%d depth=%d bytes=%d",
			file.src_path,
			png_img.width,
			png_img.height,
			png_img.channels,
			png_img.depth,
			len(png_img.pixels.buf),
		)
		nbio.next_tick_poly(file, on_encode_failed, l = io_loop)
		return
	}

	qoi_err := qoi.save_to_buffer(&file.qoi_bfr, png_img)
	if qoi_err != nil {
		fmt.eprintln("qoi encode failed:", file.dst_path, qoi_err)
		nbio.next_tick_poly(file, on_encode_failed, l = io_loop)
		return
	}

	delete(file.png_data)
	file.png_data = nil

	file.qoi_data = bytes.buffer_to_bytes(&file.qoi_bfr)
	if len(file.qoi_data) == 0 {
		fmt.eprintln("qoi encode empty output:", file.dst_path)
		nbio.next_tick_poly(file, on_encode_failed, l = io_loop)
		return
	}

	nbio.open_poly(
		file.dst_path,
		file,
		on_open_qoi,
		mode = {.Write, .Create, .Trunc},
		l = io_loop
	)
}

normalize_for_qoi :: proc(img: ^image.Image) -> bool {
	if img == nil {
		return false
	}

	if img.channels < 3 {
		if !image.expand_grayscale(img) {
			return false
		}
	}

	if img.channels == 3 {
		if !image.alpha_add_if_missing(img) {
			return false
		}
	}

	if img.channels != 4 {
		return false
	}

	pixel_count := img.width * img.height
	channel_count := pixel_count * img.channels

	if pixel_count <= 0 {
		return false
	}

	if img.depth == 8 {
		return len(img.pixels.buf) == channel_count
	}

	if img.depth != 16 {
		return false
	}

	if len(img.pixels.buf) != channel_count * 2 {
		return false
	}

	old_pixels := img.pixels
	new_pixels := bytes.Buffer{}

	if resize(&new_pixels.buf, channel_count) != nil {
		return false
	}

	src := mem.slice_data_cast([]u16, old_pixels.buf[:])
	dst := new_pixels.buf[:]

	for i in 0..<channel_count {
		v := u32(src[i])

		// 16-bit 0..65535 -> 8-bit 0..255, rounded.
		dst[i] = u8((v * 255 + 32767) / 65535)
	}

	bytes.buffer_destroy(&old_pixels)

	img.pixels = new_pixels
	img.depth = 8

	return true
}

on_encode_failed :: proc(_: ^nbio.Operation, file: ^FileDesc) {
	file_load_done(file, false)
}

on_open_qoi :: proc(op: ^nbio.Operation, file: ^FileDesc) {
	if op.open.err != .None {
		fmt.eprintfln("open failed %q: %v", file.dst_path, op.open.err)
		file_load_done(file, false)
		return
	}

	file.handle = op.open.handle

	nbio.write_poly(
		file.handle,
		0,
		file.qoi_data,
		file,
		on_write_qoi,
		all = true,
		l = io_loop,
	)
}

on_write_qoi :: proc(op: ^nbio.Operation, file: ^FileDesc) {
	file.write_ok = 
		op.write.err == .None && op.write.written == len(file.qoi_data)

	if !file.write_ok {
		fmt.eprintfln(
			"write failed %q: err=%v written=%d expected=%d",
			file.dst_path,
			op.write.err,
			op.write.written,
			len(file.qoi_data),
		)
	}

	nbio.close_poly(file.handle, file, on_close_qoi, l = io_loop)
}

on_close_qoi :: proc(op: ^nbio.Operation, file: ^FileDesc) {
	close_ok := op.close.err == .None

	if !close_ok {
		fmt.eprintln("close failed:", op.close.err, file.dst_path)
	}

	if file.write_ok && close_ok {
		fmt.printfln("converted %q -> %q", file.src_path, file.dst_path)
		file_load_done(file, true)
	} else {
		file_load_done(file, false)
	}
}

file_load_done :: proc(file: ^FileDesc, ok: bool) {
	if file.png_data != nil {
		delete(file.png_data)
		file.png_data = nil
	}

	bytes.buffer_destroy(&file.qoi_bfr)

	if len(file.src_path) > 0 {
		delete(file.src_path)
		file.src_path = ""
	}

	if len(file.dst_path) > 0 {
		delete(file.dst_path)
		file.dst_path = ""
	}

	file.qoi_data = nil

	if ok {
		io_succeeded += 1
	} else {
		io_failed += 1
	}

	io_pending -= 1
}