#!/usr/bin/env lua

package.path = table.concat({
	"?.lua",
	"?/init.lua",
	"tools/lua/?.lua",
	"tools/lua/?/init.lua",
	package.path,
}, ";")

local json = require("tools.lua.json")
local path = require("tools.lua.path")
local process = require("tools.lua.process")

local root = path.repo_root(arg and arg[0] or "tools/test.lua")
local window_collection = os.getenv("GPU_WINDOW_COLLECTION") or path.join(root, "..", "window")

local function usage()
	io.write([[
usage:
	lua tools/test.lua list
	lua tools/test.lua selftest
	lua tools/test.lua check
	lua tools/test.lua <example> [--backend NAME] [--seconds N] [--frames N] [--json PATH] [--profile PATH] [--visual] [--screenshot PATH] [--no-screenshot] [--dry-run]
	lua tools/test.lua visual <example> [--backend NAME] [--frames N] [--compare PATH] [--no-screenshot] [--proof-readback PATH] [--proof-mode MODE]
	lua tools/test.lua visual-matrix [--backend LIST] [--scene LIST] [--frames N] [--json PATH] [--dry-run]
	lua tools/test.lua windows-visual-matrix [--backend d3d11,d3d12] [--scene LIST] [--frames N] [--json PATH] [--dry-run]
	lua tools/test.lua bench <example|all> [--backend NAME] [--seconds N] [--frames N] [--json PATH]
]])
end

local function parse_options(argv, start)
	local opts = { ["_"] = {} }
	local i = start or 1
	while i <= #argv do
		local a = argv[i]
		if a:sub(1, 2) == "--" then
			local key, value = a:match("^%-%-([^=]+)=(.*)$")
			if not key then
				key = a:sub(3)
				local nextv = argv[i + 1]
				if nextv and nextv:sub(1, 2) ~= "--" then
					value = nextv
					i = i + 1
				else
					value = true
				end
			end
			opts[key:gsub("-", "_")] = value
		else
			opts._[#opts._ + 1] = a
		end
		i = i + 1
	end
	return opts
end

local function list_dirs(parent)
	local cmd
	if path.is_windows then
		cmd = 'dir /b /ad ' .. path.shell_quote(parent)
	else
		cmd = 'find ' .. path.shell_quote(parent) .. ' -mindepth 1 -maxdepth 1 -type d -printf "%f\n" | sort'
	end
	local out, status = process.capture(cmd)
	if status ~= 0 then return {} end
	local dirs = {}
	for line in out:gmatch("[^\r\n]+") do
		if path.exists(path.join(parent, line, "main.odin")) then
			dirs[#dirs + 1] = line
		end
	end
	table.sort(dirs)
	return dirs
end

local function examples()
	return list_dirs(path.join(root, "examples"))
end

local function example_exists(name)
	return path.exists(path.join(root, "examples", name, "main.odin"))
end

local function example_title(name)
	local source_path = path.join(root, "examples", name, "main.odin")
	if not path.exists(source_path) then return nil end
	local source = path.read_file(source_path)
	return source:match('app%.init_window%([^\\n]-"([^"]+)"')
end

local function default_backend()
	if path.is_windows then return "d3d12" end
	return "vulkan"
end

local function runtime_backend_supported(backend)
	if backend == "vulkan" then return true end
	if backend == "opengl" then return path.is_windows or process.uname() == "Linux" end
	if (backend == "d3d11" or backend == "d3d12") and path.is_windows then return true end
	return false
end

local function exe_suffix()
	return path.is_windows and ".exe" or ""
end

local function odin_build_example_cmd(name, backend, out, bench)
	local cmd = {
		"odin", "build",
		path.shell_quote(path.join("examples", name)),
		"-collection:window=" .. path.shell_quote(window_collection),
		"-define:GPU_BACKEND=" .. path.shell_quote(backend),
		"-out:" .. path.shell_quote(out),
	}
	if bench then
		cmd[#cmd + 1] = "-define:GPU_BENCH=true"
	end
	return table.concat(cmd, " ")
end

local function build_example(name, opts)
	local backend = opts.backend or default_backend()
	local bench = opts.seconds or opts.frames or opts.profile or opts.bench
	local out_dir = path.join(root, ".build", bench and "runs" or "examples", backend)
	path.mkdir_p(out_dir)
	local out_name = name
	if bench then
		math.randomseed(os.time())
		local entropy = tostring({}):gsub("[^%w]", "")
		out_name = string.format("%s-%d-%d-%s", name, os.time(), math.random(1000000), entropy)
	end
	local out = path.join(out_dir, out_name .. exe_suffix())
	local cmd = odin_build_example_cmd(name, backend, out, bench)
	return out, process.with_cwd(root, cmd), backend
end

local function result_path(kind, name, backend, ext)
	local dir = path.join(root, ".build", kind, backend)
	path.mkdir_p(dir)
	return path.join(dir, name .. "." .. ext)
end

local function parse_profile_summary(profile_text)
	if not profile_text then return nil end
	local summary = {
		frame_count = tonumber(profile_text:match('"frame_count"%s*:%s*(%d+)')) or 0,
		app_frame_count = tonumber(profile_text:match('"app_frame_count"%s*:%s*(%d+)')) or 0,
		app_wall_seconds = tonumber(profile_text:match('"app_wall_seconds"%s*:%s*([%d%.]+)')) or 0,
		ir2d_lower_ns = tonumber(profile_text:match('"ir2d_lower_ns"%s*:%s*(%d+)')) or 0,
		ir3d_lower_ns = tonumber(profile_text:match('"ir3d_lower_ns"%s*:%s*(%d+)')) or 0,
		compute_lower_ns = tonumber(profile_text:match('"compute_lower_ns"%s*:%s*(%d+)')) or 0,
		draw_indexed_calls = tonumber(profile_text:match('"draw_indexed_calls"%s*:%s*(%d+)')) or 0,
		dispatch_calls = tonumber(profile_text:match('"dispatch_calls"%s*:%s*(%d+)')) or 0,
	}
	if summary.frame_count > 0 then
		summary.draws_per_frame = summary.draw_indexed_calls / summary.frame_count
		summary.ir2d_us_per_frame = summary.ir2d_lower_ns / summary.frame_count / 1000.0
		summary.ir3d_us_per_frame = summary.ir3d_lower_ns / summary.frame_count / 1000.0
		summary.compute_us_per_frame = summary.compute_lower_ns / summary.frame_count / 1000.0
	end
	if summary.app_wall_seconds > 0 and summary.app_frame_count > 0 then
		summary.app_fps = summary.app_frame_count / summary.app_wall_seconds
	end
	return summary
end

local function run_example(name, opts)
	if not example_exists(name) then
		error("unknown example: " .. tostring(name))
	end

	local exe, build_cmd, backend = build_example(name, opts)
	if not runtime_backend_supported(backend) and not opts.dry_run then
		error("backend " .. backend .. " cannot run on this platform")
	end

	if opts.dry_run then
		print(build_cmd)
		print(process.with_cwd(root, path.shell_quote(exe)))
		return { example = name, backend = backend, dry_run = true }
	end

	local status = process.run(build_cmd)
	if status ~= 0 then
		return { example = name, backend = backend, build_status = status, exit_code = status, failed = true }
	end

	local profile = opts.profile or result_path("profiles", name, backend, "profile.json")
	local metrics = result_path("metrics", name, backend, "time.txt")
	if opts.visual and not opts.screenshot and not opts.no_screenshot then
		opts.screenshot = result_path("visual", name, backend, "png")
	end
	if opts.visual and not opts.frames and not opts.seconds then
		opts.frames = "120"
	end
	local visual_json = opts.screenshot and result_path("visual", name, backend, "capture.json") or nil
	local analysis_json = opts.screenshot and result_path("visual", name, backend, "analysis.json") or nil
	local screenshot_ready = opts.screenshot and result_path("visual", name, backend, "ready") or nil
	path.remove_file(profile)
	path.remove_file(metrics)
	path.remove_file(screenshot_ready)
	path.remove_file(visual_json)
	path.remove_file(analysis_json)
	local env = {
		GPU_BACKEND = backend,
		GPU_PROFILE_JSON = profile,
	}
	if opts.proof_output_dir then env.GPU_PROOF_OUTPUT_DIR = opts.proof_output_dir end
	if opts.proof_readback then env.GPU_PROOF_READBACK = opts.proof_readback end
	if opts.proof_headless then env.GPU_PROOF_HEADLESS = opts.proof_headless end
	if opts.proof_mode then env.GPU_PROOF_MODE = opts.proof_mode end
	if opts.proof_backend then env.GPU_PROOF_BACKEND = opts.proof_backend end
	if opts.proof_scene then env.GPU_PROOF_SCENE = opts.proof_scene end
	if os.getenv("GPU_VK_VALIDATION") then env.GPU_VK_VALIDATION = os.getenv("GPU_VK_VALIDATION") end
	if opts.seconds then env.GPU_TEST_SECONDS = opts.seconds end
	if opts.frames then env.GPU_TEST_FRAMES = opts.frames end
	if opts.screenshot then
		env.GPU_SCREENSHOT_READY_FILE = screenshot_ready
		env.GPU_SCREENSHOT_HOLD_SECONDS = opts.screenshot_hold or "0.75"
	end
	local expect_profile = opts.bench or opts.seconds or opts.frames or opts.profile

	local run_cmd = path.shell_quote(exe)
	local uname = process.uname()
	if not path.is_windows and process.command_exists("/usr/bin/time") then
		if uname == "Darwin" then
			run_cmd = "/usr/bin/time -l -o " .. path.shell_quote(metrics) .. " " .. run_cmd
		else
			run_cmd = "/usr/bin/time -v -o " .. path.shell_quote(metrics) .. " " .. run_cmd
		end
	end
	local env_run_cmd = process.env_prefix(env) .. run_cmd
	if opts.screenshot and not path.is_windows and process.command_exists("python3") then
		local title = opts.screenshot_title or example_title(name) or name
		local wait_seconds = opts.screenshot_timeout or "15"
		local delay = opts.screenshot_delay or "0"
		local root_fallback = opts.root_fallback and " --root-fallback" or ""
		local capture_cmd = table.concat({
			"python3",
			path.shell_quote(path.join(root, "tools", "visual_capture.py")),
			"--title", path.shell_quote(title),
			"--output", path.shell_quote(opts.screenshot),
			"--ready-file", path.shell_quote(screenshot_ready),
			"--timeout", path.shell_quote(wait_seconds),
			"--delay", path.shell_quote(delay),
			"--json", path.shell_quote(visual_json),
			root_fallback,
		}, " ")
		env_run_cmd = "sh -c " .. path.shell_quote(
			env_run_cmd .. " & pid=$!; " ..
			capture_cmd .. "; " ..
			"wait $pid"
		)
	end
	run_cmd = process.with_cwd(root, env_run_cmd)
	local started = os.clock()
	local exit_code = process.run(run_cmd)
	local wall_cpu = os.clock() - started
	local time_metrics = process.read_time_v(metrics)
	local profile_data = nil
	if path.exists(profile) then
		profile_data = path.read_file(profile)
	end
	local profile_summary = parse_profile_summary(profile_data)
	local result = {
		example = name,
		backend = backend,
		exit_code = exit_code,
		profile = profile,
		metrics = time_metrics,
		lua_cpu_seconds = wall_cpu,
		profile_present = profile_data ~= nil,
		profile_summary = profile_summary,
		screenshot = opts.screenshot,
		visual_capture = visual_json,
		visual_analysis = analysis_json,
	}
	if opts.screenshot and path.exists(opts.screenshot) and process.command_exists("python3") then
		local analyze_cmd = "python3 " .. path.shell_quote(path.join(root, "tools", "visual_analyze.py")) .. " " .. path.shell_quote(opts.screenshot) .. " --json " .. path.shell_quote(analysis_json)
		if opts.compare then
			analyze_cmd = analyze_cmd .. " --compare " .. path.shell_quote(opts.compare)
		end
		process.run(analyze_cmd)
	end
	if profile_summary and profile_summary.app_fps then
		result.fps = profile_summary.app_fps
		result.fps_source = "app_profile"
	elseif time_metrics.elapsed_seconds and profile_summary and profile_summary.frame_count and profile_summary.frame_count > 0 then
		result.fps = profile_summary.frame_count / time_metrics.elapsed_seconds
		result.fps_source = "process_elapsed"
	end
	if expect_profile and (not profile_summary or (profile_summary.frame_count or 0) == 0) then
		result.failed = true
		result.error = "no rendered frames/profile data"
	end
	if opts.json then
		path.write_file(opts.json, json.encode(result) .. "\n")
	end
	return result
end

local function cmd_list()
	print("examples:")
	for _, name in ipairs(examples()) do
		print("  " .. name)
	end
	print("components:")
	print("  render_ir")
	print("  tests")
	print("  gpu:vulkan")
	print("  gpu:d3d11")
	print("  gpu:d3d12")
	print("  gpu:opengl")
end

local function cmd_check()
	local cmd = process.with_cwd(root, "GPU_SKIP_LUA_SELFTEST=1 " .. path.shell_quote(path.join(root, "scripts", "check.sh")))
	if path.is_windows then
		cmd = process.with_cwd(root, "set GPU_SKIP_LUA_SELFTEST=1&& " .. path.shell_quote(path.join(root, "scripts", "check.sh")))
	end
	return process.run(cmd)
end

local function split_csv(value)
	local out = {}
	if not value or value == true or value == "" then return out end
	for item in tostring(value):gmatch("[^,]+") do
		local trimmed = item:gsub("^%s+", ""):gsub("%s+$", "")
		if trimmed ~= "" then out[#out + 1] = trimmed end
	end
	return out
end

local function windows_visual_matrix_strict_success(result)
	if not result or not result.summary or result.summary.total == 0 then return false end
	if result.summary.fail ~= 0 or result.summary.skip ~= 0 then return false end
	for _, row in ipairs(result.rows or {}) do
		if row.status ~= "pass" then return false end
	end
	return true
end

local function cmd_windows_visual_matrix(argv)
	local opts = parse_options(argv, 2)
	opts.backend = opts.backend or "d3d11,d3d12"
	for _, backend in ipairs(split_csv(opts.backend)) do
		if backend ~= "d3d11" and backend ~= "d3d12" then
			error("windows-visual-matrix only accepts d3d11,d3d12 backends")
		end
	end
	if not path.is_windows and not opts.dry_run then
		error("windows-visual-matrix must run on Windows unless --dry-run is supplied")
	end
	if not opts.dry_run and not opts.json then
		opts.json = path.join(root, ".build", "visual-matrix", "windows-d3d.json")
	end
	local visual_matrix = require("tools.visual_matrix")
	local result = visual_matrix.run(opts, { root = root, window_collection = window_collection })
	print(json.encode(result))
	return windows_visual_matrix_strict_success(result) and 0 or 1
end

local function selftest()
	local visual_matrix = require("tools.visual_matrix")
	assert(json.encode({ b = 2, a = "x" }) == '{"a":"x","b":2}')
	assert(path.basename(path.join("examples", "hello_window")) == "hello_window")
	assert(example_exists("hello_window"))
	assert(not example_exists("__missing__"))
	assert(example_title("torches") == "GPU - Torches & Point Lights")
	local exe, build_cmd, backend = build_example("hello_window", { backend = "vulkan", dry_run = true })
	assert(build_cmd:match("odin") and build_cmd:match("hello_window"))
	assert(backend == "vulkan")
	assert(exe:match("hello_window"))
	local parsed = process.read_time_v(path.join(root, "__no_such_file__"))
	assert(type(parsed) == "table")
	assert(process.parse_elapsed_seconds("0:00.15") == 0.15)
	assert(process.parse_elapsed_seconds("1:02:03") == 3723)
	assert(process.command_exists("python3"))
	assert(process.run("python3 " .. path.shell_quote(path.join(root, "tools", "visual_assert.py")) .. " --selftest") == 0)
	visual_matrix.selftest(root)
	assert(not windows_visual_matrix_strict_success(visual_matrix.run({
		dry_run = true,
		backend = "d3d11,d3d12",
		scene = "stencil_clip",
	}, { root = root })))
	print("lua selftest ok")
end

local function cmd_bench(argv)
	local target = argv[2] or "all"
	local opts = parse_options(argv, 3)
	opts.bench = true
	opts.seconds = opts.seconds or "3"
	local names = {}
	if target == "all" then
		names = examples()
	else
		names = { target }
	end
	local results = {}
	for _, name in ipairs(names) do
		local ok, result = pcall(run_example, name, opts)
		if not ok then
			result = { example = name, failed = true, error = result }
		end
		results[#results + 1] = result
		local status = result.failed and "FAIL" or tostring(result.exit_code or 0)
		local rss = result.metrics and result.metrics.max_rss_kib or nil
		local frames = result.profile_summary and result.profile_summary.frame_count or "-"
		local fps = result.fps and string.format("%.0f", result.fps) or "-"
		local draws = result.profile_summary and result.profile_summary.draws_per_frame and string.format("%.2f", result.profile_summary.draws_per_frame) or "-"
		print(string.format("%-20s exit=%s frames=%s fps=%s draws/frame=%s rss_kib=%s", name, status, tostring(frames), fps, draws, tostring(rss or "-")))
	end
	if opts.json then
		path.write_file(opts.json, json.encode({ results = results }) .. "\n")
	end
end

local function main(argv)
	local cmd = argv[1]
	if not cmd or cmd == "help" or cmd == "--help" then
		usage()
		return 0
	elseif cmd == "list" then
		cmd_list()
		return 0
	elseif cmd == "selftest" then
		selftest()
		return 0
	elseif cmd == "check" then
		return cmd_check()
	elseif cmd == "visual" then
		local target = argv[2]
		if not target then error("visual requires example name") end
		local opts = parse_options(argv, 3)
		opts.visual = true
		local result = run_example(target, opts)
		print(json.encode(result))
		if result.failed then return 1 end
		return result.exit_code or 0
	elseif cmd == "visual-matrix" then
		local visual_matrix = require("tools.visual_matrix")
		return visual_matrix.main(argv, { root = root, window_collection = window_collection })
	elseif cmd == "windows-visual-matrix" then
		return cmd_windows_visual_matrix(argv)
	elseif cmd == "bench" then
		cmd_bench(argv)
		return 0
	else
		local opts = parse_options(argv, 2)
		local result = run_example(cmd, opts)
		if not opts.dry_run then
			print(json.encode(result))
		end
		if result.failed then return 1 end
		return result.exit_code or 0
	end
end

local ok, code_or_err = pcall(main, arg or {})
if not ok then
	io.stderr:write("error: " .. tostring(code_or_err) .. "\n")
	os.exit(1)
end
os.exit(code_or_err or 0)