module main

import os
import os.cmdline
import v.util.recompilation

const args_ = arguments()
const is_debug = args_.contains('-debug')

// support a renamed `v` executable too:
const vexe = os.getenv_opt('VEXE') or { @VEXE }

const vroot = os.dir(vexe)
const vself_flags_with_values = ['-o', '-os', '-cc', '-gc', '-cf', '-cflags', '-d', '-define']

fn main() {
	// make testing `v up` easier, by providing a way to force `v self` to fail,
	// to test the fallback logic:
	if os.getenv('VSELF_SHOULD_FAIL') != '' {
		eprintln('v self failed')
		exit(1)
	}
	vexe_name := os.file_name(vexe)
	short_v_name := vexe_name.all_before('.')

	recompilation.must_be_enabled(vroot,
		'Please install V from source, to use `${vexe_name} self` .')
	os.chdir(vroot)!
	os.setenv('VCOLORS', 'always', true)
	repeat_count, mut args := extract_repeat_count(args_[1..].filter(it != 'self'))
	if args.len == 0 || ('-cc' !in args && '-prod' !in args && '-parallel-cc' !in args) {
		// compiling by default, i.e. `v self`:
		uos := os.user_os()
		uname := os.uname()
		if uos == 'macos' && uname.machine == 'arm64' {
			// Apple silicon, like m1, m2 etc
			// Use tcc by default for V, since tinycc is much faster and also
			// it already supports compiling many programs like V itself, that do not depend on inlined objective-C code
			args << '-cc tcc'
		} else if uos == 'linux' && uname.machine in ['arm64', 'aarch64'] {
			// Bundled TCC can hang while bootstrapping V on Linux ARM64, so
			// prefer the system compiler for self-builds there.
			args << ['-cc', os.getenv_opt('CC') or { 'cc' }]
		}
	}
	if !has_gc_arg(args) {
		args << ['-gc', 'none']
	}
	jargs := args.join(' ')
	obinary := cmdline.option(args, '-o', '')
	sargs := if obinary != '' { jargs } else { '${jargs} -o v2' }
	options := if args.len > 0 { '(${sargs})' } else { '' }
	final_binary := if obinary != '' { obinary } else { 'v2' }
	pgo_cc_kind := pgo_compiler_kind(args)
	for run_idx in 0 .. repeat_count {
		run_label := if repeat_count > 1 { ' [${run_idx + 1}/${repeat_count}]' } else { '' }
		println('V self compiling${run_label} ${options}...')
		cmd := '${os.quoted_path(vexe)} ${sargs} ${os.quoted_path('cmd/v')}'
		mut used_pgo := false
		if pgo_cc_kind != '' {
			used_pgo = compile_with_pgo(vroot, vexe, args, final_binary, pgo_cc_kind)
			if !used_pgo {
				eprintln('PGO self-build failed; falling back to a regular self-build.')
			}
		}
		if !used_pgo {
			if !try_compile(cmd) {
				bootstrap_self_build(vroot, clone_args(args), final_binary) or {
					eprintln('cannot compile to `${vroot}`: \n${err.msg()}')
					exit(1)
				}
			}
		}
		if obinary == '' {
			backup_old_version_and_rename_newer(short_v_name) or { panic(err.msg()) }
		}
	}
	if obinary != '' {
		return
	}
	println('V built successfully as executable "${vexe_name}".')
}

fn repeat_count_arg(arg string) int {
	if arg.len < 2 || arg[0] != `x` {
		return 0
	}
	for ch in arg[1..].bytes() {
		if !ch.is_digit() {
			return 0
		}
	}
	count := arg[1..].int()
	return if count > 0 { count } else { 0 }
}

fn extract_repeat_count(args []string) (int, []string) {
	mut repeat_count := 1
	mut filtered := []string{cap: args.len}
	mut should_skip_repeat_check := false
	for arg in args {
		if should_skip_repeat_check {
			filtered << arg
			should_skip_repeat_check = false
			continue
		}
		if arg in vself_flags_with_values {
			filtered << arg
			should_skip_repeat_check = true
			continue
		}
		if repeat_count == 1 {
			count := repeat_count_arg(arg)
			if count > 0 {
				repeat_count = count
				continue
			}
		}
		filtered << arg
	}
	return repeat_count, filtered
}

fn has_gc_arg(args []string) bool {
	for arg in args {
		if arg == '-gc' {
			return true
		}
		if arg.starts_with('-gc=') {
			return true
		}
	}
	return false
}

fn has_profile_cflag(args []string) bool {
	mut skip_next := false
	for i, arg in args {
		if skip_next {
			skip_next = false
			continue
		}
		if arg in ['-cflags', '-cf'] {
			if i + 1 < args.len {
				next_arg := args[i + 1]
				if next_arg.contains('-fprofile') {
					return true
				}
				skip_next = true
			}
			continue
		}
		if (arg.starts_with('-cflags=') || arg.starts_with('-cf=')) && arg.contains('-fprofile') {
			return true
		}
	}
	return false
}

fn pgo_compiler_kind(args []string) string {
	if '-prod' !in args || '-no-prod-options' in args {
		return ''
	}
	if os.user_os() == 'windows' {
		return ''
	}
	if has_profile_cflag(args) {
		return ''
	}
	mut ccompiler := cmdline.option(args, '-cc', '')
	if ccompiler == '' {
		ccompiler = os.getenv_opt('CC') or { 'cc' }
	}
	cc_file_name := os.file_name(ccompiler)
	if cc_file_name.contains('clang') || cc_file_name.contains('gcc')
		|| cc_file_name.contains('g++') || ccompiler == 'cc' {
		cc_ver := os.execute('${os.quoted_path(ccompiler)} --version').output
		if cc_ver.contains('clang') {
			_ := find_llvm_profdata() or { return '' }
			return 'clang'
		}
		if cc_ver.contains('Free Software Foundation') || cc_ver.contains('GCC') {
			return 'gcc'
		}
	}
	if cc_file_name.contains('clang') {
		_ := find_llvm_profdata() or { return '' }
		return 'clang'
	}
	if cc_file_name.contains('gcc') || cc_file_name.contains('g++') {
		return 'gcc'
	}
	return ''
}

fn find_llvm_profdata() !string {
	if profdata := os.find_abs_path_of_executable('llvm-profdata') {
		return profdata
	}
	$if macos {
		xcrun_result := os.execute('xcrun --find llvm-profdata')
		if xcrun_result.exit_code == 0 {
			xcrun_path := xcrun_result.output.trim_space()
			if xcrun_path != '' && os.exists(xcrun_path) {
				return xcrun_path
			}
		}
	}
	return error('can not find llvm-profdata in PATH')
}

fn with_output_arg(args []string, output string) []string {
	mut res := []string{cap: args.len + 2}
	mut skip_next := false
	for i, arg in args {
		if skip_next {
			skip_next = false
			continue
		}
		if arg == '-o' {
			if i + 1 < args.len {
				skip_next = true
			}
			continue
		}
		if arg.starts_with('-o=') {
			continue
		}
		res << arg
	}
	res << ['-o', output]
	return res
}

fn clone_args(args []string) []string {
	mut cloned := []string{cap: args.len}
	for arg in args {
		cloned << arg.clone()
	}
	return cloned
}

fn compose_v_cmd(vexe string, args []string, source string) string {
	mut parts := []string{cap: args.len + 2}
	parts << os.quoted_path(vexe)
	for arg in args {
		parts << os.quoted_path(arg)
	}
	parts << os.quoted_path(source)
	return parts.join(' ')
}

fn run_cmd(cmd string) ! {
	result := os.execute(cmd)
	if result.exit_code != 0 {
		return error(result.output)
	}
	if result.output.len > 0 {
		println(result.output.trim_space())
	}
}

fn try_compile(cmd string) bool {
	result := os.execute(cmd)
	if result.exit_code != 0 {
		return false
	}
	if result.output.len > 0 {
		println(result.output.trim_space())
	}
	return true
}

fn compile_with_pgo(vroot string, vexe string, args []string, out_binary string, cc_kind string) bool {
	pgo_workspace := os.join_path(vroot, '.vself_pgo')
	os.rmdir_all(pgo_workspace) or {}
	os.mkdir_all(pgo_workspace) or {
		eprintln('PGO disabled: can not create ${pgo_workspace}: ${err.msg()}')
		return false
	}
	defer {
		os.rmdir_all(pgo_workspace) or {}
	}
	profile_dir := os.join_path(pgo_workspace, 'profile')
	os.mkdir_all(profile_dir) or {
		eprintln('PGO disabled: can not create ${profile_dir}: ${err.msg()}')
		return false
	}
	pgo_binary := os.join_path(pgo_workspace, 'v_pgo_gen')
	training_output := os.join_path(pgo_workspace, 'cmd_v_training.c')
	mut use_profile_flag := '-fprofile-use=${profile_dir}'
	mut llvm_profdata := ''
	mut profile_data := ''
	if cc_kind == 'clang' {
		llvm_profdata = find_llvm_profdata() or {
			eprintln('PGO disabled: can not find `llvm-profdata`.')
			return false
		}
		profile_data = os.join_path(pgo_workspace, 'code.profdata')
		use_profile_flag = '-fprofile-use=${profile_data}'
	}
	mut generate_args := with_output_arg(args, pgo_binary)
	generate_args << ['-cflags', '-fprofile-generate=${profile_dir}']
	generate_cmd := compose_v_cmd(vexe, generate_args, 'cmd/v')
	run_cmd(generate_cmd) or {
		eprintln('PGO step failed while building the instrumented compiler.')
		eprintln(err.msg())
		return false
	}
	training_cmd := '${os.quoted_path(pgo_binary)} -o ${os.quoted_path(training_output)} ${os.quoted_path('cmd/v')}'
	run_cmd(training_cmd) or {
		eprintln('PGO step failed while generating the profiling data.')
		eprintln(err.msg())
		return false
	}
	if cc_kind == 'clang' {
		merge_cmd := '${os.quoted_path(llvm_profdata)} merge -output=${os.quoted_path(profile_data)} ${os.quoted_path(profile_dir)}'
		run_cmd(merge_cmd) or {
			eprintln('PGO step failed while merging the profiling data.')
			eprintln(err.msg())
			return false
		}
	}
	mut final_args := with_output_arg(args, out_binary)
	final_args << ['-cflags', use_profile_flag]
	if cc_kind == 'gcc' {
		final_args << ['-cflags', '-fprofile-correction']
	}
	final_cmd := compose_v_cmd(vexe, final_args, 'cmd/v')
	run_cmd(final_cmd) or {
		eprintln('PGO step failed while building the final compiler binary.')
		eprintln(err.msg())
		return false
	}
	return true
}

fn bootstrap_self_build(vroot string, args []string, final_binary string) ! {
	bootstrap_prefix := '.vself_bootstrap'
	mut bootstrap_v1 := '${bootstrap_prefix}_v1'
	mut bootstrap_v2 := '${bootstrap_prefix}_v2'
	exe_ext := if os.user_os() == 'windows' { '.exe' } else { '' }
	bootstrap_v1 += exe_ext
	bootstrap_v2 += exe_ext
	os.rm(bootstrap_v1) or {}
	os.rm(bootstrap_v2) or {}
	defer {
		os.rm(bootstrap_v1) or {}
		os.rm(bootstrap_v2) or {}
	}
	vc_source := os.join_path(vroot, 'vc',
		if os.user_os() == 'windows' { 'v_win.c' } else { 'v.c' })
	if !os.exists(vc_source) {
		return error('bootstrap fallback failed: `${vc_source}` is missing')
	}
	cc := os.getenv_opt('CC') or {
		if os.user_os() == 'windows' { 'gcc' } else { 'cc' }
	}
	bootstrap_v1_build_cmd := bootstrap_c_cmd(cc, bootstrap_v1, vc_source)
	run_cmd(bootstrap_v1_build_cmd) or {
		return error('bootstrap fallback failed while building v1.\n${err.msg()}')
	}
	mut bootstrap_args := ['-no-parallel']
	bootstrap_args << with_output_arg(args, bootstrap_v2)
	bootstrap_v1_cmd := os.join_path('.', bootstrap_v1)
	bootstrap_v2_cmd := '${os.quoted_path(bootstrap_v1_cmd)} ${bootstrap_args.join(' ')} ${os.quoted_path('cmd/v')}'
	run_cmd(bootstrap_v2_cmd) or {
		return error('bootstrap fallback failed while building v2.\n${err.msg()}')
	}
	final_args := with_output_arg(args, final_binary)
	bootstrap_v2_cmd_path := os.join_path('.', bootstrap_v2)
	final_cmd := '${os.quoted_path(bootstrap_v2_cmd_path)} ${final_args.join(' ')} ${os.quoted_path('cmd/v')}'
	run_cmd(final_cmd) or {
		return error('bootstrap fallback failed while building the final compiler.\n${err.msg()}')
	}
}

fn bootstrap_c_cmd(cc string, out_binary string, vc_source string) string {
	mut parts := []string{cap: 8}
	parts << os.quoted_path(cc)
	if os.user_os() == 'windows' {
		parts << ['-std=c99', '-municode', '-w', '-o', os.quoted_path(out_binary),
			os.quoted_path(vc_source), '-lws2_32']
	} else {
		parts << ['-std=c99', '-w', '-o', os.quoted_path(out_binary),
			os.quoted_path(vc_source), '-lm', '-lpthread']
	}
	return parts.join(' ')
}

fn list_folder(short_v_name string, bmessage string, message string) {
	if !is_debug {
		return
	}
	if bmessage != '' {
		println(bmessage)
	}
	if os.user_os() == 'windows' {
		os.system('dir ${short_v_name}*.exe')
	} else {
		os.system('ls -lartd ${short_v_name}*')
	}
	println(message)
}

fn backup_old_version_and_rename_newer(short_v_name string) !bool {
	mut errors := []string{}
	short_v_file := if os.user_os() == 'windows' { '${short_v_name}.exe' } else { '${short_v_name}' }
	short_v2_file := if os.user_os() == 'windows' { 'v2.exe' } else { 'v2' }
	short_bak_file := if os.user_os() == 'windows' { 'v_old.exe' } else { 'v_old' }
	v_file := os.real_path(short_v_file)
	v2_file := os.real_path(short_v2_file)
	bak_file := os.real_path(short_bak_file)

	list_folder(short_v_name, 'before:', 'removing ${bak_file} ...')
	if os.exists(bak_file) {
		os.rm(bak_file) or { errors << 'failed removing ${bak_file}: ${err.msg()}' }
	}

	list_folder(short_v_name, '', 'moving ${v_file} to ${bak_file} ...')
	os.mv(v_file, bak_file) or { errors << err.msg() }

	list_folder(short_v_name, '', 'removing ${v_file} ...')
	os.rm(v_file) or {}

	list_folder(short_v_name, '', 'moving ${v2_file} to ${v_file} ...')
	os.mv_by_cp(v2_file, v_file) or { panic(err.msg()) }

	list_folder(short_v_name, 'after:', '')

	if errors.len > 0 {
		eprintln('backup errors:\n  >>  ' + errors.join('\n  >>  '))
	}
	return true
}