#!/usr/bin/env -S v

import compress.bzip2

const compression_levels = [1, 6, 9]!

struct TestVector {
	name string
	data []u8
}

fn main() {
	ensure_tools() or {
		eprintln('SKIP: ${err.msg()}')
		exit(2)
	}
	vectors := make_test_vectors()
	mut total_checks := 0
	mut total_runs := 0
	tmp_root := join_path(temp_dir(), 'v_bzip2_interop_${getpid()}')
	mkdir_all(tmp_root) or { panic(err) }
	defer {
		rmdir_all(tmp_root) or {}
	}
	for level in compression_levels {
		for i, vec in vectors {
			total_checks += run_case(tmp_root, level, i, vec) or {
				eprintln('FAIL: ${vec.name} (level=${level}): ${err}')
				exit(1)
				0
			}
			total_runs++
			println('ok ${total_runs}/${vectors.len * compression_levels.len}: ${vec.name} (level=${level}, ${vec.data.len} bytes)')
		}
	}
	println('PASS: ${vectors.len} vectors x ${compression_levels.len} levels, ${total_checks} cross-checks')
}

fn ensure_tools() ! {
	must_succeed('bzip2 --help >/dev/null 2>&1', 'system bzip2 command is not available')!
	must_succeed("python3 -c 'import bz2' >/dev/null 2>&1",
		'python3 with bz2 module is not available')!
}

fn make_test_vectors() []TestVector {
	mut vectors := []TestVector{}
	vectors << TestVector{'empty', []u8{}}
	vectors << TestVector{'ascii_text', 'The quick brown fox jumps over the lazy dog.\n'.repeat(64).bytes()}
	vectors << TestVector{'repeated_byte', []u8{len: 10000, init: `A`}}
	vectors << TestVector{'all_bytes_x4', all_bytes_repeated(4)}
	vectors << TestVector{'lcg_64k', lcg_bytes(65536)}
	return vectors
}

fn run_case(tmp_root string, level int, case_idx int, vec TestVector) !int {
	case_dir := join_path(tmp_root, 'case_l${level}_${case_idx:02}_${vec.name}')
	mkdir_all(case_dir)!

	v_bz2 := bzip2.compress(vec.data, block_size: level)!
	cli_bz2 := cli_compress(case_dir, vec.data, level)!
	py_bz2 := py_compress(case_dir, vec.data, level)!

	mut checks := 0

	producers := {
		'v':   v_bz2
		'cli': cli_bz2
		'py':  py_bz2
	}
	for producer, compressed in producers {
		v_plain := bzip2.decompress(compressed)!
		assert_equal_bytes('v.decompress(${producer}.compress, level=${level})', vec.data, v_plain)!
		checks++

		cli_plain := cli_decompress(case_dir, producer, compressed)!
		assert_equal_bytes('cli.decompress(${producer}.compress, level=${level})', vec.data,
			cli_plain)!
		checks++

		py_plain := py_decompress(case_dir, producer, compressed)!
		assert_equal_bytes('py.decompress(${producer}.compress, level=${level})', vec.data,
			py_plain)!
		checks++
	}
	return checks
}

fn cli_compress(case_dir string, plain []u8, level int) ![]u8 {
	in_path := join_path(case_dir, 'plain.in')
	out_path := join_path(case_dir, 'cli_l${level}.bz2')
	write_file_array(in_path, plain)!
	must_succeed('bzip2 -${level} -c -- ${shell_quote(in_path)} > ${shell_quote(out_path)}',
		'bzip2 compression failed')!
	return read_bytes(out_path)!
}

fn cli_decompress(case_dir string, producer string, compressed []u8) ![]u8 {
	in_path := join_path(case_dir, '${producer}.for_cli.bz2')
	out_path := join_path(case_dir, '${producer}.from_cli.out')
	write_file_array(in_path, compressed)!
	must_succeed('bzip2 -d -c -- ${shell_quote(in_path)} > ${shell_quote(out_path)}',
		'bzip2 decompression failed')!
	return read_bytes(out_path)!
}

fn py_compress(case_dir string, plain []u8, level int) ![]u8 {
	in_path := join_path(case_dir, 'plain_py.in')
	out_path := join_path(case_dir, 'py_l${level}.bz2')
	write_file_array(in_path, plain)!
	py_code := 'import bz2, pathlib, sys; p=pathlib.Path(sys.argv[1]); o=pathlib.Path(sys.argv[2]); l=int(sys.argv[3]); o.write_bytes(bz2.compress(p.read_bytes(), compresslevel=l))'
	must_succeed('python3 -c ${shell_quote(py_code)} ${shell_quote(in_path)} ${shell_quote(out_path)} ${level}',
		'python bz2 compression failed')!
	return read_bytes(out_path)!
}

fn py_decompress(case_dir string, producer string, compressed []u8) ![]u8 {
	in_path := join_path(case_dir, '${producer}.for_py.bz2')
	out_path := join_path(case_dir, '${producer}.from_py.out')
	write_file_array(in_path, compressed)!
	py_code := 'import bz2, pathlib, sys; p=pathlib.Path(sys.argv[1]); o=pathlib.Path(sys.argv[2]); o.write_bytes(bz2.decompress(p.read_bytes()))'
	must_succeed('python3 -c ${shell_quote(py_code)} ${shell_quote(in_path)} ${shell_quote(out_path)}',
		'python bz2 decompression failed')!
	return read_bytes(out_path)!
}

fn all_bytes_repeated(times int) []u8 {
	mut out := []u8{cap: 256 * times}
	for _ in 0 .. times {
		for i in 0 .. 256 {
			out << u8(i)
		}
	}
	return out
}

fn lcg_bytes(n int) []u8 {
	mut out := []u8{len: n}
	mut x := u32(0x12345678)
	for i in 0 .. n {
		x = x * u32(1664525) + u32(1013904223)
		out[i] = u8((x >> 16) & u32(0xff))
	}
	return out
}

fn assert_equal_bytes(label string, expected []u8, got []u8) ! {
	if expected.len != got.len {
		return error('${label}: length mismatch expected=${expected.len} got=${got.len}')
	}
	for i in 0 .. expected.len {
		if expected[i] != got[i] {
			return error('${label}: byte mismatch at offset ${i}')
		}
	}
}

fn must_succeed(command string, context string) ! {
	res := execute(command)
	if res.exit_code != 0 {
		return error('${context}\ncommand: ${command}\nexit_code: ${res.exit_code}\n${res.output}')
	}
}

fn shell_quote(s string) string {
	return "'${s.replace("'", "'\\''")}'"
}