module lz

const stream_magic = [u8(0x56), 0x4c, 0x5a, 0x31] // VLZ1

struct MatchProfile {
    window      int
    min_match   int
    max_match   int
    max_literal int
}

const match_hash_bits = 16
const match_hash_size = 1 << match_hash_bits
const max_match_candidates = 64

fn wrap_payload(format Format, source []u8, payload []u8) []u8 {
    mut out := []u8{cap: stream_magic.len + 8 + payload.len}
    out << stream_magic
    out << u8(format)
    encode_uvarint(mut out, u64(source.len))
    out << payload
    return out
}

fn unwrap_payload(data []u8, format Format) !([]u8, i64) {
    if data.len < stream_magic.len + 2 {
        return error('invalid lz stream: too short')
    }
    if data[..stream_magic.len] != stream_magic {
        return error('invalid lz stream: bad magic')
    }
    wire_format := data[stream_magic.len]
    if wire_format != u8(format) {
        return error('invalid lz stream: format mismatch')
    }
    decoded_len_u64, mut pos, ok := decode_uvarint(data, stream_magic.len + 1)
    if !ok {
        return error('invalid lz stream: bad length')
    }
    if decoded_len_u64 > u64(max_int) {
        return error('invalid lz stream: decoded length too large')
    }
    decoded_len := i64(decoded_len_u64)
    if pos > data.len {
        return error('invalid lz stream: truncated payload')
    }
    return data[pos..], decoded_len
}

fn compress_with_profile(data []u8, profile MatchProfile, format Format) []u8 {
    if data.len == 0 {
        return wrap_payload(format, data, []u8{})
    }
    mut payload := []u8{cap: data.len}
    mut literals := []u8{cap: profile.max_literal}
    mut last_match := []int{len: match_hash_size, init: -1}
    mut prev_match := []int{len: data.len, init: -1}
    mut pos := 0
    for pos < data.len {
        offset, length := find_best_match(data, pos, profile, last_match, prev_match)
        if length >= profile.min_match {
            flush_literals(mut payload, mut literals)
            emit_match(mut payload, offset, length, profile.min_match)
            for i := pos; i < pos + length; i++ {
                index_match_position(data, i, mut last_match, mut prev_match)
            }
            pos += length
        } else {
            literals << data[pos]
            if literals.len == profile.max_literal {
                flush_literals(mut payload, mut literals)
            }
            index_match_position(data, pos, mut last_match, mut prev_match)
            pos++
        }
    }
    flush_literals(mut payload, mut literals)
    return wrap_payload(format, data, payload)
}

fn decompress_with_profile(data []u8, profile MatchProfile, format Format) ![]u8 {
    payload, expected_len := unwrap_payload(data, format)!
    mut out := []u8{cap: int(expected_len)}
    mut pos := 0
    for pos < payload.len {
        control := payload[pos]
        pos++
        if control & 0x80 == 0 {
            literal_len := int(control & 0x7f) + 1
            if pos + literal_len > payload.len {
                return error('invalid lz stream: truncated literal')
            }
            out << payload[pos..pos + literal_len]
            pos += literal_len
            continue
        }
        match_len := int(control & 0x7f) + profile.min_match
        offset, next_pos, ok := decode_uvarint(payload, pos)
        if !ok {
            return error('invalid lz stream: truncated match offset')
        }
        pos = next_pos
        if offset == 0 || offset > u64(max_i64) || i64(offset) > i64(out.len) {
            return error('invalid lz stream: bad match offset')
        }
        offset_int := int(offset)
        base := out.len - offset_int
        for i in 0 .. match_len {
            out << out[base + i]
        }
    }
    if i64(out.len) != expected_len {
        return error('invalid lz stream: length mismatch')
    }
    return out
}

fn find_best_match(data []u8, pos int, profile MatchProfile, last_match []int, prev_match []int) (int, int) {
    if pos + profile.min_match > data.len {
        return 0, 0
    }
    max_len := if pos + profile.max_match < data.len { profile.max_match } else { data.len - pos }
    mut best_len := 0
    mut best_offset := 0
    hash_idx := match_hash(data, pos)
    mut candidates_checked := 0
    mut i := last_match[hash_idx]
    for i >= 0 && candidates_checked < max_match_candidates {
        offset := pos - i
        if offset > profile.window {
            break
        }
        mut current_len := 0
        for current_len < max_len && data[i + current_len] == data[pos + current_len] {
            current_len++
        }
        if current_len > best_len {
            best_len = current_len
            best_offset = offset
            if best_len == max_len {
                break
            }
        }
        i = prev_match[i]
        candidates_checked++
    }
    return best_offset, best_len
}

fn index_match_position(data []u8, pos int, mut last_match []int, mut prev_match []int) {
    if pos + 2 >= data.len {
        return
    }
    hash_idx := match_hash(data, pos)
    prev_match[pos] = last_match[hash_idx]
    last_match[hash_idx] = pos
}

fn match_hash(data []u8, pos int) int {
    v := (u32(data[pos]) << 16) | (u32(data[pos + 1]) << 8) | u32(data[pos + 2])
    return int((v * u32(2654435761)) >> (32 - match_hash_bits))
}

fn flush_literals(mut payload []u8, mut literals []u8) {
    if literals.len == 0 {
        return
    }
    payload << u8(literals.len - 1)
    payload << literals
    literals.clear()
}

fn emit_match(mut payload []u8, offset int, length int, min_match int) {
    payload << u8(0x80 | u8(length - min_match))
    encode_uvarint(mut payload, u64(offset))
}

fn encode_uvarint(mut out []u8, value u64) {
    mut v := value
    for v >= 0x80 {
        out << u8(v & 0x7f | 0x80)
        v >>= 7
    }
    out << u8(v)
}

fn decode_uvarint(data []u8, start int) (u64, int, bool) {
    mut value := u64(0)
    mut shift := u32(0)
    mut pos := start
    for pos < data.len && shift <= 63 {
        b := data[pos]
        pos++
        value |= u64(b & 0x7f) << shift
        if b & 0x80 == 0 {
            return value, pos, true
        }
        shift += 7
    }
    return 0, start, false
}