module builder

import os

const windows_icon_group_resource_id = 1
const max_windows_icon_dimension = 256

struct WindowsIconImage {
    width        u8
    height       u8
    color_count  u8
    planes       u16
    bit_count    u16
    bytes_in_res u32
    image_data   []u8
}

struct WindowsIconSize {
    width  int
    height int
}

fn (b &Builder) ensure_windows_icon_flag_is_valid() {
    if b.pref.icon_path == '' {
        return
    }
    if b.pref.os != .windows || b.pref.build_mode == .build_module || b.pref.is_o
        || b.pref.is_shared {
        verror('`-icon` is supported only when building Windows executables')
    }
    if b.pref.generate_c_project != '' || b.pref.out_name.ends_with('.c')
        || b.pref.out_name.ends_with('.js') || b.pref.should_output_to_stdout() {
        verror('`-icon` cannot be used when emitting generated C/JS output instead of a Windows executable')
    }
}

fn (mut b Builder) prepare_cross_windows_icon_resource() !string {
    if b.pref.icon_path == '' {
        return ''
    }
    ico_path := b.prepare_windows_icon_ico_path()!
    rc_path := b.get_vtmp_filename(b.pref.out_name, '.icon.rc')
    obj_path := b.get_vtmp_filename(b.pref.out_name, '.icon.o')
    os.write_file(rc_path, '1 ICON ${rc_quoted_string(ico_path)}\n')!
    b.pref.cleanup_files << rc_path
    b.pref.cleanup_files << obj_path
    windres := b.find_windres()!
    cmd := '${os.quoted_path(windres)} -i ${os.quoted_path(rc_path)} -o ${os.quoted_path(obj_path)} -O coff'
    res := os.execute(cmd)
    if res.exit_code != 0 {
        return error('failed to compile Windows icon resource with `${windres}`: ${res.output.trim_space()}')
    }
    return obj_path
}

fn (mut b Builder) prepare_windows_icon_ico_path() !string {
    if b.pref.icon_path == '' {
        return ''
    }
    icon_path := os.real_path(b.pref.icon_path)
    if !os.is_file(icon_path) {
        return error('icon file `${icon_path}` does not exist')
    }
    match os.file_ext(icon_path).to_lower_ascii() {
        '.ico' {
            return icon_path
        }
        '.png' {
            png_bytes := os.read_bytes(icon_path)!
            ico_bytes := png_to_ico_bytes(png_bytes)!
            ico_path := b.get_vtmp_filename(b.pref.out_name, '.icon.ico')
            os.write_file_array(ico_path, ico_bytes)!
            b.pref.cleanup_files << ico_path
            return ico_path
        }
        else {
            return error('`-icon` accepts only `.ico` or `.png` files')
        }
    }
}

fn (b &Builder) find_windres() !string {
    compiler_dir := if b.pref.ccompiler.contains('/') || b.pref.ccompiler.contains('\\') {
        os.dir(b.pref.ccompiler)
    } else {
        ''
    }
    compiler_name := executable_stem(os.file_name(b.pref.ccompiler))
    mut candidates := []string{}
    for suffix in ['-gcc', '-clang', '-cc', '-g++', '-clang++'] {
        if compiler_name.ends_with(suffix) {
            candidates << '${compiler_name[..compiler_name.len - suffix.len]}-windres'
        }
    }
    candidates << 'windres'
    candidates << 'llvm-windres'
    for candidate in candidates {
        for name in [candidate, '${candidate}.exe'] {
            if compiler_dir != '' {
                full_path := os.join_path(compiler_dir, name)
                if os.is_file(full_path) {
                    return full_path
                }
            }
            if resolved := os.find_abs_path_of_executable(name) {
                return resolved
            }
        }
    }
    return error('could not find `windres`, which is needed for `-icon` while cross-compiling to Windows')
}

fn executable_stem(name string) string {
    lower_name := name.to_lower_ascii()
    if lower_name.ends_with('.exe') {
        return name[..name.len - 4]
    }
    return name
}

fn rc_quoted_string(path string) string {
    return '"' + path.replace('\\', '\\\\').replace('"', '\\"') + '"'
}

fn parse_ico_file(path string) ![]WindowsIconImage {
    return parse_ico_bytes(os.read_bytes(path)!)
}

fn parse_ico_bytes(data []u8) ![]WindowsIconImage {
    if data.len < 6 {
        return error('invalid icon file: missing ICO header')
    }
    if read_le_u16(data, 0) != 0 || read_le_u16(data, 2) != 1 {
        return error('invalid icon file: expected an ICO header')
    }
    image_count := int(read_le_u16(data, 4))
    if image_count <= 0 {
        return error('invalid icon file: no icon images were found')
    }
    if data.len < 6 + (image_count * 16) {
        return error('invalid icon file: truncated icon directory')
    }
    mut images := []WindowsIconImage{cap: image_count}
    for i := 0; i < image_count; i++ {
        entry_offset := 6 + (i * 16)
        image_size := int(read_le_u32(data, entry_offset + 8))
        image_offset := int(read_le_u32(data, entry_offset + 12))
        image_end := image_offset + image_size
        if image_offset < 0 || image_size <= 0 || image_offset > data.len || image_end > data.len {
            return error('invalid icon file: icon image ${i + 1} points outside the file')
        }
        images << WindowsIconImage{
            width:        data[entry_offset]
            height:       data[entry_offset + 1]
            color_count:  data[entry_offset + 2]
            planes:       read_le_u16(data, entry_offset + 4)
            bit_count:    read_le_u16(data, entry_offset + 6)
            bytes_in_res: u32(image_size)
            image_data:   data[image_offset..image_end].clone()
        }
    }
    return images
}

fn png_to_ico_bytes(png []u8) ![]u8 {
    size := png_dimensions(png)!
    mut ico := []u8{cap: 22 + png.len}
    append_le_u16(mut ico, 0)
    append_le_u16(mut ico, 1)
    append_le_u16(mut ico, 1)
    ico << u8(if size.width == max_windows_icon_dimension { 0 } else { size.width })
    ico << u8(if size.height == max_windows_icon_dimension { 0 } else { size.height })
    ico << u8(0)
    ico << u8(0)
    append_le_u16(mut ico, 1)
    append_le_u16(mut ico, 32)
    append_le_u32(mut ico, u32(png.len))
    append_le_u32(mut ico, u32(22))
    ico << png
    return ico
}

fn png_dimensions(png []u8) !WindowsIconSize {
    if png.len < 24 {
        return error('invalid PNG icon file: missing PNG header')
    }
    if png[0] != 0x89 || png[1] != `P` || png[2] != `N` || png[3] != `G` || png[4] != 0x0d
        || png[5] != 0x0a || png[6] != 0x1a || png[7] != 0x0a {
        return error('invalid PNG icon file: bad PNG signature')
    }
    if png[12] != `I` || png[13] != `H` || png[14] != `D` || png[15] != `R` {
        return error('invalid PNG icon file: missing IHDR chunk')
    }
    width := int(read_be_u32(png, 16))
    height := int(read_be_u32(png, 20))
    if width <= 0 || height <= 0 || width > max_windows_icon_dimension
        || height > max_windows_icon_dimension {
        return error('PNG icons must be between 1x1 and 256x256 pixels')
    }
    return WindowsIconSize{
        width:  width
        height: height
    }
}

fn build_group_icon_resource(images []WindowsIconImage) []u8 {
    mut data := []u8{cap: 6 + (images.len * 14)}
    append_le_u16(mut data, 0)
    append_le_u16(mut data, 1)
    append_le_u16(mut data, u16(images.len))
    for i, image in images {
        data << image.width
        data << image.height
        data << image.color_count
        data << u8(0)
        append_le_u16(mut data, image.planes)
        append_le_u16(mut data, image.bit_count)
        append_le_u32(mut data, image.bytes_in_res)
        append_le_u16(mut data, u16(i + 1))
    }
    return data
}

fn append_le_u16(mut data []u8, value u16) {
    data << u8(value & 0xff)
    data << u8(value >> 8)
}

fn append_le_u32(mut data []u8, value u32) {
    data << u8(value & 0xff)
    data << u8((value >> 8) & 0xff)
    data << u8((value >> 16) & 0xff)
    data << u8(value >> 24)
}

@[direct_array_access; inline]
fn read_le_u16(data []u8, offset int) u16 {
    return u16(data[offset]) | (u16(data[offset + 1]) << 8)
}

@[direct_array_access; inline]
fn read_le_u32(data []u8, offset int) u32 {
    b0 := u32(data[offset])
    b1 := u32(data[offset + 1])
    b2 := u32(data[offset + 2])
    b3 := u32(data[offset + 3])
    return b0 | (b1 << 8) | (b2 << 16) | (b3 << 24)
}

@[direct_array_access; inline]
fn read_be_u32(data []u8, offset int) u32 {
    b0 := u32(data[offset])
    b1 := u32(data[offset + 1])
    b2 := u32(data[offset + 2])
    b3 := u32(data[offset + 3])
    return (b0 << 24) | (b1 << 16) | (b2 << 8) | b3
}