#!/usr/bin/env v

import os

// generates ML-DSA test vectors by building Go's reference implementation
// and running gen.go with crypto/internal/fips140/mldsa
//
// Usage:
//   v run testdata/gen.vsh
//   v run testdata/gen.vsh <go-source-path> # uses existing Go tree

const go_repo = 'https://github.com/golang/go.git'
const go_ref = 'master'
const go_mldsa_subpath = os.join_path('src', 'crypto', 'internal', 'fips140', 'mldsa')
const go_gen_pkg = os.join_path('.', go_mldsa_subpath, 'gen')

fn main() {
	unbuffer_stdout()
	script_dir := os.dir(@FILE)

	mut go_root := os.join_path(os.temp_dir(), 'go')
	if os.args.len > 1 {
		go_root = os.args[1]
		if !os.exists(os.join_path(go_root, go_mldsa_subpath)) {
			eprintln('error: ${go_root} does not look like a Go source tree')
			exit(1)
		}
		println('> Using existing Go tree at ${go_root}')
	} else {
		if os.exists(os.join_path(go_root, 'src')) {
			println('> Reusing cached Go tree at ${go_root}')
			run_or_exit('git -C ${go_root} pull --quiet')
		} else {
			println('> Cloning Go source tree to ${go_root}...')
			run_or_exit('git clone --depth 1 --branch ${go_ref} ${go_repo} ${go_root}')
		}
	}

	go_bin := os.join_path(go_root, 'bin', 'go')
	if !os.exists(go_bin) {
		println('> Building Go toolchain...')
		run_or_exit('cd ${go_root}/src && bash make.bash')
	}
	println('> Go binary: ${go_bin}')

	gen_dst_dir := os.join_path(go_root, go_mldsa_subpath, 'gen')
	os.mkdir_all(gen_dst_dir) or {
		eprintln('error: mkdir failed: ${err}')
		exit(1)
	}
	os.cp(os.join_path(script_dir, 'gen.go'), os.join_path(gen_dst_dir, 'main.go')) or {
		eprintln('error: copy failed: ${err}')
		exit(1)
	}

	println('> Generating test vectors...')
	res := run_or_exit('cd ${go_root} && GOROOT=${go_root} ${go_bin} run ${go_gen_pkg}')

	os.rmdir_all(gen_dst_dir) or { eprintln('warning: rmdir failed: ${err}') }

	vectors_path := os.join_path(script_dir, 'vectors.json')
	os.write_file(vectors_path, res) or {
		eprintln('error: write failed: ${err}')
		exit(1)
	}
	println('> Wrote ${vectors_path}')
}

fn run_or_exit(cmd string) string {
	res := os.execute_opt(cmd) or {
		eprintln('error: ${err}')
		exit(1)
	}
	if res.exit_code != 0 {
		eprintln('error: ${cmd}\n${res.output}')
		exit(1)
	}
	return res.output
}