-- HUD Overlay Black-Key Script for OBS (separate Create / Update buttons)
-- Black-to-transparent keying: black pixels become fully transparent, brightness is used as alpha
-- Principle: alpha = max(R, G, B); pure black -> alpha 0 (transparent), the brighter the pixel the more opaque
--   For a "pure green on pure black" HUD this is simpler and more stable than greenness keying:
--   no green judgment at all, so anti-aliased edges and dim segments key out naturally.
-- Requires: obs-shaderfilter plugin (installed)
-- Install: OBS -> Tools -> Scripts -> + -> select this file -> pick a mode
--   "Create Scene"   : builds the scene from scratch (deletes and recreates both media sources)
--   "Update Scene"   : updates the existing scene in place (video paths, color, keying parameters)

obs = obslua

local SETTING_FRONT_PATH  = "front_path"
local SETTING_BACK_PATH   = "back_path"
local SETTING_COLOR       = "hud_color"     -- preset color index (1 = keep original green)
local SETTING_BLACK_LOW   = "black_low"
local SETTING_GAIN        = "alpha_gain"
local SETTING_OPACITY     = "opacity"

local SCENE_NAME  = "HUD Overlay Black Key"
local FRONT_NAME  = "HUD - Front BlackKey"
local BACK_NAME   = "Center IG - Back BlackKey"
local FILTER_NAME = "Black Key"

-- Preset colors; index 1 keeps the original green from the video
local COLOR_PRESETS = {
	{ name = "Keep original color (video green)",        r = 0.0, g = 1.0,  b = 0.0,  original = true  },
	{ name = "Red (recommended, complementary to teal)", r = 1.0, g = 0.0,  b = 0.0,  original = false },
	{ name = "Magenta",                                  r = 1.0, g = 0.0,  b = 1.0,  original = false },
	{ name = "Orange",                                   r = 1.0, g = 0.5,  b = 0.0,  original = false },
	{ name = "Yellow",                                   r = 1.0, g = 1.0,  b = 0.0,  original = false },
	{ name = "White",                                    r = 1.0, g = 1.0,  b = 1.0,  original = false },
	{ name = "Bright green (neon green)",                r = 0.25, g = 1.0, b = 0.1,  original = false },
}

local front_path = ""
local back_path  = ""
local color_index = 1
local black_low  = 0.0
local alpha_gain = 1.0
local opacity    = 1.0

local fit_retries = 0

-- Black-to-transparent keying shader:
--   k = max(R, G, B); black -> 0, pure green -> 1.0
--   alpha = rescale k above black_low, then apply gain; soft alpha preserves anti-aliased edges
--   Output is premultiplied: either the original pixel color or the specified HUD color
local SHADER_TEXT = [==[
uniform float black_low <
	string label = "Black threshold";
	string widget_type = "slider";
	float minimum = 0.0;
	float maximum = 0.5;
	float step = 0.001;
> = 0.0;

uniform float alpha_gain <
	string label = "Alpha gain";
	string widget_type = "slider";
	float minimum = 0.5;
	float maximum = 4.0;
	float step = 0.01;
> = 1.0;

uniform bool use_original <
	string label = "Keep original color";
> = true;

uniform float hud_r <
	string label = "HUD color R";
	string widget_type = "slider";
	float minimum = 0.0;
	float maximum = 1.0;
	float step = 0.01;
> = 0.0;

uniform float hud_g <
	string label = "HUD color G";
	string widget_type = "slider";
	float minimum = 0.0;
	float maximum = 1.0;
	float step = 0.01;
> = 1.0;

uniform float hud_b <
	string label = "HUD color B";
	string widget_type = "slider";
	float minimum = 0.0;
	float maximum = 1.0;
	float step = 0.01;
> = 0.0;

uniform float opacity <
	string label = "Opacity";
	string widget_type = "slider";
	float minimum = 0.0;
	float maximum = 1.0;
	float step = 0.01;
> = 1.0;

float4 mainImage(VertData v_in) : TARGET
{
	float4 px = image.Sample(textureSampler, v_in.uv);

	// The OBS filter chain input is premultiplied; un-premultiply first (media sources have alpha=1, so usually no effect)
	px.rgb *= (px.a > 0.0) ? (1.0 / px.a) : 0.0;

	// Brightness key: max channel is 0 for black, 1.0 for pure green
	float k = max(px.r, max(px.g, px.b));

	// Everything below black_low is fully transparent; rescale the rest back to [0, 1]
	float a = clamp((k - black_low) / max(1.0 - black_low, 0.001), 0.0, 1.0);
	a = clamp(a * alpha_gain, 0.0, 1.0) * opacity;

	float3 col = use_original ? px.rgb : float3(hud_r, hud_g, hud_b);

	// Premultiplied output: color x alpha
	return float4(col * a, a);
}
]==]

------------------------------------------------------------
-- Utility functions
------------------------------------------------------------

local function remove_source_by_name(name)
	local s = obs.obs_get_source_by_name(name)
	if s ~= nil then
		obs.obs_source_remove(s)
		obs.obs_source_release(s)
	end
end

local function create_media_source(name, path)
	local s = obs.obs_data_create()
	obs.obs_data_set_string(s, "local_file", path)
	obs.obs_data_set_bool(s, "is_local_file", true)
	obs.obs_data_set_bool(s, "looping", true)
	obs.obs_data_set_bool(s, "restart_on_activate", true)
	obs.obs_data_set_bool(s, "close_when_inactive", false)
	local src = obs.obs_source_create("ffmpeg_source", name, s, nil)
	obs.obs_data_release(s)
	return src
end

-- Point an existing media source at a new file without recreating it
local function update_media_source_path(src, path)
	local s = obs.obs_source_get_settings(src)
	obs.obs_data_set_string(s, "local_file", path)
	obs.obs_source_update(src, s)
	obs.obs_data_release(s)
end

local function apply_key_filter(src)
	local old = obs.obs_source_get_filter_by_name(src, FILTER_NAME)
	if old ~= nil then
		obs.obs_source_filter_remove(src, old)
		obs.obs_source_release(old)
	end

	local preset = COLOR_PRESETS[color_index] or COLOR_PRESETS[1]

	local fs = obs.obs_data_create()
	obs.obs_data_set_bool(fs, "from_file", false)
	obs.obs_data_set_string(fs, "shader_text", SHADER_TEXT)
	obs.obs_data_set_double(fs, "black_low", black_low)
	obs.obs_data_set_double(fs, "alpha_gain", alpha_gain)
	obs.obs_data_set_bool(fs, "use_original", preset.original)
	obs.obs_data_set_double(fs, "hud_r", preset.r)
	obs.obs_data_set_double(fs, "hud_g", preset.g)
	obs.obs_data_set_double(fs, "hud_b", preset.b)
	obs.obs_data_set_double(fs, "opacity", opacity)

	-- First check whether the obs-shaderfilter plugin is loaded
	if obs.obs_source_get_display_name("shader_filter") == nil then
		obs.obs_data_release(fs)
		obs.script_log(obs.LOG_WARNING,
			"HUD Overlay Black Key: obs-shaderfilter plugin not detected; keying was not applied. Please fully quit OBS (including the system tray) and restart it.")
		return
	end

	local f = obs.obs_source_create("shader_filter", FILTER_NAME, fs, nil)
	obs.obs_data_release(fs)

	if f ~= nil then
		obs.obs_source_filter_add(src, f)
		obs.obs_source_release(f)
	end
end

local function fit_items_to_canvas()
	local vi = obs.obs_video_info()
	if obs.obs_get_video_info(vi) == false then
		return false
	end
	local cw, ch = vi.base_width, vi.base_height

	local scene_src = obs.obs_get_source_by_name(SCENE_NAME)
	if scene_src == nil then
		return false
	end
	local scene = obs.obs_scene_from_source(scene_src)
	local items = obs.obs_scene_enum_items(scene)
	local all_ready = true

	for _, item in ipairs(items) do
		local s = obs.obs_sceneitem_get_source(item)
		local w = obs.obs_source_get_width(s)
		local h = obs.obs_source_get_height(s)
		if w > 0 and h > 0 then
			local scale = math.min(cw / w, ch / h)
			local sc = obs.vec2()
			sc.x = scale
			sc.y = scale
			obs.obs_sceneitem_set_scale(item, sc)
			local pos = obs.vec2()
			pos.x = (cw - w * scale) / 2
			pos.y = (ch - h * scale) / 2
			obs.obs_sceneitem_set_pos(item, pos)
			-- Use Lanczos when upscaling line-art footage; sharper than the default bilinear
			if obs.obs_sceneitem_set_scale_filter ~= nil then
				obs.obs_sceneitem_set_scale_filter(item, obs.OBS_SCALE_LANCZOS or 4)
			end
		else
			all_ready = false
		end
	end

	obs.sceneitem_list_release(items)
	obs.obs_source_release(scene_src)
	return all_ready
end

local function fit_retry()
	fit_retries = fit_retries - 1
	if fit_items_to_canvas() or fit_retries <= 0 then
		obs.timer_remove(fit_retry)
	end
end

local function schedule_fit()
	fit_retries = 10
	obs.timer_remove(fit_retry)
	obs.timer_add(fit_retry, 500)
end

------------------------------------------------------------
-- Create scene (from scratch)
------------------------------------------------------------

local function create_scene()
	if front_path == "" or back_path == "" then
		obs.script_log(obs.LOG_WARNING, "HUD Overlay Black Key: please set the foreground and background video paths first")
		return
	end

	local scene = nil
	local scene_src = obs.obs_get_source_by_name(SCENE_NAME)
	if scene_src ~= nil then
		scene = obs.obs_scene_from_source(scene_src)
	else
		scene = obs.obs_scene_create(SCENE_NAME)
		scene_src = obs.obs_scene_get_source(scene)
	end

	remove_source_by_name(FRONT_NAME)
	remove_source_by_name(BACK_NAME)

	local back = create_media_source(BACK_NAME, back_path)
	obs.obs_scene_add(scene, back)

	local front = create_media_source(FRONT_NAME, front_path)
	apply_key_filter(front)
	obs.obs_scene_add(scene, front)

	obs.obs_source_release(back)
	obs.obs_source_release(front)

	obs.obs_frontend_set_current_scene(scene_src)
	obs.obs_source_release(scene_src)

	schedule_fit()

	obs.script_log(obs.LOG_INFO, "HUD Overlay Black Key: scene created, symbol color = " ..
		(COLOR_PRESETS[color_index] or COLOR_PRESETS[1]).name)
end

------------------------------------------------------------
-- Update scene (in place, without recreating sources)
------------------------------------------------------------

local function update_scene()
	if front_path == "" or back_path == "" then
		obs.script_log(obs.LOG_WARNING, "HUD Overlay Black Key: please set the foreground and background video paths first")
		return
	end

	local scene_src = obs.obs_get_source_by_name(SCENE_NAME)
	if scene_src == nil then
		obs.script_log(obs.LOG_WARNING,
			"HUD Overlay Black Key: scene \"" .. SCENE_NAME .. "\" does not exist yet; please use \"Create Scene\" first")
		return
	end

	local back  = obs.obs_get_source_by_name(BACK_NAME)
	local front = obs.obs_get_source_by_name(FRONT_NAME)
	if back == nil or front == nil then
		if back ~= nil then obs.obs_source_release(back) end
		if front ~= nil then obs.obs_source_release(front) end
		obs.obs_source_release(scene_src)
		obs.script_log(obs.LOG_WARNING,
			"HUD Overlay Black Key: media sources are missing; please use \"Create Scene\" to rebuild the scene")
		return
	end

	-- Update video files in place
	update_media_source_path(back, back_path)
	update_media_source_path(front, front_path)

	-- Re-apply the keying filter with the current color / keying parameters
	apply_key_filter(front)

	obs.obs_source_release(back)
	obs.obs_source_release(front)

	obs.obs_frontend_set_current_scene(scene_src)
	obs.obs_source_release(scene_src)

	schedule_fit()

	obs.script_log(obs.LOG_INFO, "HUD Overlay Black Key: scene updated, symbol color = " ..
		(COLOR_PRESETS[color_index] or COLOR_PRESETS[1]).name)
end

local function on_create_clicked(props, p)
	create_scene()
	return true
end

local function on_update_clicked(props, p)
	update_scene()
	return true
end

------------------------------------------------------------
-- Script UI and settings
------------------------------------------------------------

function script_description()
	return "Black-to-transparent keying: black pixels become fully transparent; pixel brightness (max of R/G/B) is used as alpha.\n" ..
		"For a 'pure green on pure black' HUD this is simpler and more stable than greenness keying; edges and dim segments key out naturally.\n" ..
		"\"Create Scene\": builds the scene from scratch (deletes and recreates both media sources).\n" ..
		"\"Update Scene\": updates the existing scene in place (video paths, symbol color, keying parameters).\n" ..
		"The filter panel also has Black threshold / Alpha gain / HUD color sliders for fine-tuning."
end

function script_properties()
	local props = obs.obs_properties_create()

	obs.obs_properties_add_path(props, SETTING_FRONT_PATH,
		"Foreground video (HUD, green symbols on black)", obs.OBS_PATH_FILE,
		"Video files (*.avi *.mp4 *.mov *.mkv *.webm);;All files (*.*)", "")

	obs.obs_properties_add_path(props, SETTING_BACK_PATH,
		"Background video (bottom layer)", obs.OBS_PATH_FILE,
		"Video files (*.avi *.mp4 *.mov *.mkv *.webm);;All files (*.*)", "")

	local clist = obs.obs_properties_add_list(props, SETTING_COLOR,
		"Symbol color", obs.OBS_COMBO_TYPE_LIST, obs.OBS_COMBO_FORMAT_INT)
	for i, c in ipairs(COLOR_PRESETS) do
		obs.obs_property_list_add_int(clist, c.name, i)
	end

	obs.obs_properties_add_float_slider(props, SETTING_BLACK_LOW,
		"Black threshold (increase if dark noise remains around the symbols)", 0.0, 0.5, 0.001)
	obs.obs_properties_add_float_slider(props, SETTING_GAIN,
		"Alpha gain (increase if the symbols look too transparent)", 0.5, 4.0, 0.01)
	obs.obs_properties_add_float_slider(props, SETTING_OPACITY,
		"Overall opacity", 0.0, 1.0, 0.01)

	obs.obs_properties_add_button(props, "create_btn", "Create Scene", on_create_clicked)
	obs.obs_properties_add_button(props, "update_btn", "Update Scene", on_update_clicked)

	return props
end

function script_defaults(settings)
	obs.obs_data_set_default_string(settings, SETTING_FRONT_PATH,
		"D:\\Test Pictures + Video\\HUD - Front Video.avi")
	obs.obs_data_set_default_string(settings, SETTING_BACK_PATH,
		"D:\\Test Pictures + Video\\Center IG - Back Videos.avi")
	obs.obs_data_set_default_int(settings, SETTING_COLOR, 1)
	obs.obs_data_set_default_double(settings, SETTING_BLACK_LOW, 0.0)
	obs.obs_data_set_default_double(settings, SETTING_GAIN, 1.0)
	obs.obs_data_set_default_double(settings, SETTING_OPACITY, 1.0)
end

function script_update(settings)
	front_path  = obs.obs_data_get_string(settings, SETTING_FRONT_PATH)
	back_path   = obs.obs_data_get_string(settings, SETTING_BACK_PATH)
	color_index = obs.obs_data_get_int(settings, SETTING_COLOR)
	black_low   = obs.obs_data_get_double(settings, SETTING_BLACK_LOW)
	alpha_gain  = obs.obs_data_get_double(settings, SETTING_GAIN)
	opacity     = obs.obs_data_get_double(settings, SETTING_OPACITY)
end