module lz

// compress_lzw compresses data using a pure-V LZW dictionary stream.
pub fn compress_lzw(data []u8) ![]u8 {
	mut payload := []u8{}
	mut dict := map[string]int{}
	for i in 0 .. 256 {
		dict[[u8(i)].bytestr()] = i
	}
	mut next_code := 256
	mut word := ''

	for b in data {
		symbol := [b].bytestr()
		candidate := word + symbol
		if candidate in dict {
			word = candidate
			continue
		}
		if word != '' {
			encode_uvarint(mut payload, u64(dict[word]))
		}
		dict[candidate] = next_code
		next_code++
		word = symbol
	}

	if word != '' {
		encode_uvarint(mut payload, u64(dict[word]))
	}

	return wrap_payload(.lzw, data, payload)
}

// decompress_lzw decompresses data produced by compress_lzw.
pub fn decompress_lzw(data []u8) ![]u8 {
	payload, expected_len := unwrap_payload(data, .lzw)!
	if payload.len == 0 {
		if expected_len == i64(0) {
			return []u8{}
		}
		return error('invalid lzw stream: missing codes')
	}

	mut dict := map[int][]u8{}
	for i in 0 .. 256 {
		dict[i] = [u8(i)]
	}
	mut next_code := 256
	mut pos := 0

	first_code, next_pos, ok := decode_uvarint(payload, pos)
	if !ok {
		return error('invalid lzw stream: bad initial code')
	}
	pos = next_pos
	if int(first_code) !in dict {
		return error('invalid lzw stream: unknown initial code')
	}
	mut word := dict[int(first_code)].clone()
	mut out := word.clone()

	for pos < payload.len {
		code_u64, new_pos, ok_code := decode_uvarint(payload, pos)
		if !ok_code {
			return error('invalid lzw stream: bad code')
		}
		pos = new_pos
		code := int(code_u64)
		mut entry := []u8{}
		if code in dict {
			entry = dict[code].clone()
		} else if code == next_code {
			entry = word.clone()
			entry << word[0]
		} else {
			return error('invalid lzw stream: unknown code')
		}

		out << entry
		mut new_entry := word.clone()
		new_entry << entry[0]
		dict[next_code] = new_entry
		next_code++
		word = entry.clone()
	}

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