#!/usr/bin/env -S v

import compress.zlib

struct TestVector {
    name string
    data []u8
}

fn main() {
    tmp_root := join_path(temp_dir(), 'v_zlib_interop_${getpid()}')
    mkdir_all(tmp_root) or { panic(err) }
    defer {
        rmdir_all(tmp_root) or {}
    }

    c_bin := compile_c_helper(tmp_root) or {
        eprintln('SKIP: ${err.msg()}')
        exit(2)
    }
    py_script := join_path(@DIR, 'zlib_ref.py')
    ensure_python(py_script) or {
        eprintln('SKIP: ${err.msg()}')
        exit(2)
    }

    vectors := make_test_vectors()
    mut total_checks := 0
    for i, vec in vectors {
        total_checks += run_case(tmp_root, c_bin, py_script, i, vec) or {
            eprintln('FAIL: ${vec.name}: ${err.msg()}')
            exit(1)
            0
        }
        println('ok ${i + 1}/${vectors.len}: ${vec.name} (${vec.data.len} bytes)')
    }
    println('PASS: ${vectors.len} vectors, ${total_checks} cross-checks')
}

fn compile_c_helper(tmp_root string) !string {
    cc := choose_cc()
    if cc == '' {
        return error('no C compiler found (tried cc, gcc, clang)')
    }
    src := join_path(@DIR, 'zlib_ref.c')
    bin := join_path(tmp_root, 'zlib_interop_ref')
    must_succeed('${cc} -O2 ${shell_quote(src)} -lz -o ${shell_quote(bin)}',
        'C zlib helper build failed')!
    return bin
}

fn choose_cc() string {
    for cc in ['cc', 'gcc', 'clang'] {
        if execute('${cc} --version >/dev/null 2>&1').exit_code == 0 {
            return cc
        }
    }
    return ''
}

fn ensure_python(py_script string) ! {
    must_succeed("python3 -c 'import zlib' >/dev/null 2>&1",
        'python3 with zlib module is not available')!
    if !exists(py_script) {
        return error('missing Python helper: ${py_script}')
    }
}

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, c_bin string, py_script string, case_idx int, vec TestVector) !int {
    case_dir := join_path(tmp_root, 'case_${case_idx:02}_${vec.name}')
    mkdir_all(case_dir)!

    v_z := zlib.compress(vec.data)!
    c_z := c_compress(case_dir, c_bin, vec.data)!
    py_z := py_compress(case_dir, py_script, vec.data)!

    mut checks := 0
    producers := {
        'v':  v_z
        'c':  c_z
        'py': py_z
    }
    for producer, compressed in producers {
        v_plain := zlib.decompress(compressed)!
        assert_equal_bytes('v.decompress(${producer}.compress)', vec.data, v_plain)!
        checks++

        c_plain := c_decompress(case_dir, c_bin, producer, compressed)!
        assert_equal_bytes('c.decompress(${producer}.compress)', vec.data, c_plain)!
        checks++

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

fn c_compress(case_dir string, c_bin string, plain []u8) ![]u8 {
    in_path := join_path(case_dir, 'plain.in')
    out_path := join_path(case_dir, 'c.zlib')
    write_file_array(in_path, plain)!
    must_succeed('${shell_quote(c_bin)} compress ${shell_quote(in_path)} ${shell_quote(out_path)}',
        'C zlib compression failed')!
    return read_bytes(out_path)!
}

fn c_decompress(case_dir string, c_bin string, producer string, compressed []u8) ![]u8 {
    in_path := join_path(case_dir, '${producer}.for_c.zlib')
    out_path := join_path(case_dir, '${producer}.from_c.out')
    write_file_array(in_path, compressed)!
    must_succeed('${shell_quote(c_bin)} decompress ${shell_quote(in_path)} ${shell_quote(out_path)}',
        'C zlib decompression failed')!
    return read_bytes(out_path)!
}

fn py_compress(case_dir string, py_script string, plain []u8) ![]u8 {
    in_path := join_path(case_dir, 'plain_py.in')
    out_path := join_path(case_dir, 'py.zlib')
    write_file_array(in_path, plain)!
    must_succeed('python3 ${shell_quote(py_script)} compress ${shell_quote(in_path)} ${shell_quote(out_path)}',
        'Python zlib compression failed')!
    return read_bytes(out_path)!
}

fn py_decompress(case_dir string, py_script string, producer string, compressed []u8) ![]u8 {
    in_path := join_path(case_dir, '${producer}.for_py.zlib')
    out_path := join_path(case_dir, '${producer}.from_py.out')
    write_file_array(in_path, compressed)!
    must_succeed('python3 ${shell_quote(py_script)} decompress ${shell_quote(in_path)} ${shell_quote(out_path)}',
        'Python zlib 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("'", "'\\''")}'"
}