import "env.memory" memory(4);

global mut base_end: i32 = 0;

export fn load_uw8(module_size: i32) -> i32 {
    let lazy version = 0?0 - 1;
    if version < 0 {
        return module_size;
    }

    let module_end = 0x1e000 + module_size;
    if version & 1 {
        module_end = uncompress(1, 0x1e001);
    } else {
        copy(0x1e000, 0, module_size);
    }
    copy(0, 0x3c200, 8);

    let base_start = 0x3c208;
    let dest = 8;
    let src = 0x1e001;

    loop sections {
        if src < module_end & (base_start >= base_end | src?0 <= base_start?0) {
            let lazy length2 = copy_section(dest, src);
            dest = dest + length2;
            if base_start < base_end & src?0 == base_start?0 {
                base_start = base_start + section_size(base_start);
            }
            src = src + length2;
            branch sections;
        }

        if base_start < base_end {
            let lazy length3 = copy_section(dest, base_start);
            dest = dest + length3;
            base_start = base_start + length3;
            branch sections;
        }
    }

    dest
}

fn section_size(ptr: i32) -> i32 {
    let p = ptr;
    let l: i32;
    let shift: i32;
    loop size {
        let lazy b = (p := p + 1)?0;
        l = l | ((b & 127) << shift);
        shift = shift + 7;
        branch_if b >> 7: size;
    }
    p + 1 - ptr + l
}

fn copy_section(dest: i32, src: i32) -> i32 {
    let lazy length = section_size(src);
    copy(dest, src, length);
    length
}

fn copy(dest: i32, src: i32, len: i32) {
    loop bytes {
        if len > 0 {
            (dest + (len := len - 1))?0 = (src + len)?0;
            branch bytes;
        }
    }
}

// upkr unpacker

global mut upkr_src_ptr: i32 = 0;
global mut upkr_state: i32 = 0;

// uncompress upkr compressed data at `src` into the buffer at `dest`
// returns the end of the uncompressed data
export fn uncompress(src_ptr: i32, dest_ptr: i32) -> i32 {
    upkr_src_ptr = src_ptr;
    upkr_state = 0;

    let offset: i32;

    let i: i32;
    loop init_contexts {
        i?0x3c000 = 0x80;
        branch_if (i := i + 1) < 256 + 1 + 128: init_contexts;
    }

    let prev_was_match: i32;

    block finished {
        loop unpack_loop {
            let lazy is_match = upkr_bit(0);
            if is_match {
                let inline new_offset = if prev_was_match { 1 } else { upkr_bit(256) };
                if new_offset {
                    branch_if !(offset := upkr_length(257) - 1): finished;
                }
                let length = upkr_length(257 + 64);
                loop copy {
                    dest_ptr?0 = (dest_ptr - offset)?0;
                    dest_ptr = dest_ptr + 1;
                    branch_if (length := length - 1): copy;
                }
            } else {
                // literal
                let byte = 1;
                loop literal {
                    branch_if (byte := (byte << 1) | upkr_bit(byte)) < 256: literal;
                }
                dest_ptr?0 = byte;
                dest_ptr = dest_ptr + 1;
            }
            prev_was_match = is_match;
            branch unpack_loop;
        }
    }

    dest_ptr
}

fn upkr_length(context_index: i32) -> i32 {
    let length: i32;
    let bit_pos: i32;
    loop bits {
        if upkr_bit(context_index + bit_pos) {
            length = length | (upkr_bit(context_index + bit_pos + 32) << bit_pos);
            bit_pos = bit_pos + 1;
            branch bits;
        }
    }
    length | (1 << bit_pos)
}

fn upkr_bit(context_index: i32) -> i32 {
    let lazy prob = context_index?0x3c000;

    loop refill {
        if upkr_state < 1<<12 {
            upkr_state = (upkr_state << 8) | upkr_src_ptr?0;
            upkr_src_ptr = upkr_src_ptr + 1;
            branch refill;
        }
    }

    let lazy state_low = upkr_state & 0xff;
    let lazy state_hi = upkr_state >> 8;
    let lazy bit = state_low < prob;

    upkr_state = state_low + select(bit, prob * state_hi, (0x100 - prob) * state_hi - prob);

    context_index?0x3c000 = prob + ((7 + bit * 257 - prob) >> 4);

    bit
}

start fn unpack_base() {
    base_end = uncompress(0, 0x3c200);
}

data 0 {
    file("../target/base.upk")
}