module lz

// compress_lz78 compresses data using a pure-V LZ78 dictionary stream.
pub fn compress_lz78(data []u8) ![]u8 {
    mut payload := []u8{}
    mut dict := map[string]int{}
    mut next_index := 1
    mut word := []u8{}

    for b in data {
        mut candidate := word.clone()
        candidate << b
        candidate_key := candidate.bytestr()
        if candidate_key in dict {
            word = candidate.clone()
            continue
        }

        prefix_index := if word.len == 0 { 0 } else { dict[word.bytestr()] }
        encode_uvarint(mut payload, u64(prefix_index))
        payload << u8(1)
        payload << b
        dict[candidate_key] = next_index
        next_index++
        word.clear()
    }

    if word.len > 0 {
        final_index := dict[word.bytestr()]
        encode_uvarint(mut payload, u64(final_index))
        payload << u8(0)
    }

    return wrap_payload(.lz78, data, payload)
}

// decompress_lz78 decompresses data produced by compress_lz78.
pub fn decompress_lz78(data []u8) ![]u8 {
    payload, expected_len := unwrap_payload(data, .lz78)!
    mut out := []u8{cap: int(expected_len)}
    mut dict := map[int][]u8{}
    mut next_index := 1
    mut pos := 0

    for pos < payload.len {
        prefix, next_pos, ok := decode_uvarint(payload, pos)
        if !ok {
            return error('invalid lz78 stream: bad prefix index')
        }
        pos = next_pos
        if pos >= payload.len {
            return error('invalid lz78 stream: missing suffix flag')
        }
        has_suffix := payload[pos]
        pos++

        mut phrase := if prefix == 0 {
            []u8{}
        } else {
            if int(prefix) !in dict {
                return error('invalid lz78 stream: unknown prefix index')
            }
            dict[int(prefix)].clone()
        }

        if has_suffix == 1 {
            if pos >= payload.len {
                return error('invalid lz78 stream: missing suffix byte')
            }
            phrase << payload[pos]
            pos++
        } else if has_suffix != 0 {
            return error('invalid lz78 stream: bad suffix flag')
        }

        out << phrase
        dict[next_index] = phrase
        next_index++
    }

    if i64(out.len) != expected_len {
        return error('invalid lz78 stream: length mismatch')
    }
    return out
}