path: root/ar/.config/mpv/scripts/subtitle-search.lua

--[[
Based on sub-search script by kelciour (https://github.com/kelciour/mpv-scripts/blob/master/sub-search.lua)

Differences from the original script:

- Searches in a subtitle file active as a primary subtitle instead of attempting to find subtitle files matching video name
- Outputs all search results in OSD list instead of jumping between them with a hotkey (the closest subtitle is selected by default)
- Supports searching unicode text (subtitles should be encoded as utf8, please re-encode your subtitles if you get no results searching for unicode text)
- Embedded console replaced with more recent variant from mpv sources (to support unicode input)
- Takes into account current `sub-delay` value
- Can search in embedded subtitles (requires ffmpeg to be installed to extract subtitles from video files)
- Can search subtitles for youtube videos (requires ffmpeg to be installed to fetch remote subtitles)
- Supports `.srt`, `.vtt` and `.sub` (microdvd) subtitle formats
- Can use special phrase "*" to show all subtitle lines
- Use `ctrl+shift+f` shortcut to show all subtitle lines simultaneously and dynamically highlight the current line
- Press `Ctrl+Shift+Enter` in result list to adjust `sub-delay` so that selected subtitle line is displayed at the current position

Requires `script-modules/utf8` repository, `script-modules/scroll-list.lua`, `script-modules/sha1.lua`, `script-modules/utf8_data.lua` and `script-modules/input-console.lua` to work.

You can clone `script-modules/utf8` repository with the following command (assuming you are in mpv config directory): `git clone git@github.com:Stepets/utf8.lua.git script-modules/utf8`

Usage:
    Press Ctrl + F, print something and press Enter.
Example:
    'You are playing Empire Strikes Back and press Ctrl+F, type "I am you father" + Enter
    and voilá, the scene pops up.'
--]]

package.path = package.path .. ";" .. mp.command_native({ "expand-path", "~~/script-modules/?.lua" })

local mp = require("mp")
local utils = require("mp.utils")
local msg = require("mp.msg")
local input_console = require("input-console")
local result_list = require("scroll-list")
local utf8 = require("utf8/init")
local utf8_data = require("utf8_data")
local sha1 = require("sha1")

utf8.config = {
	conversion = {
		uc_lc = utf8_data.utf8_uc_lc,
		lc_uc = utf8_data.utf8_lc_uc,
	},
}

utf8:init()

table.insert(result_list.keybinds, {
	"ENTER",
	"jump_to_result",
	function()
		local selected_index = result_list.selected
		if selected_index == nil then
			return
		end

		local selected = result_list.list[selected_index]
		mp.commandv("seek", selected.time, "absolute+exact")
	end,
	{},
})
table.insert(result_list.keybinds, {
	"Ctrl+Shift+ENTER",
	"sync_to_result",
	function()
		local selected_index = result_list.selected
		if selected_index == nil then
			return
		end

		local selected = result_list.list[selected_index]
		local old_delay = mp.get_property_native("sub-delay")
		local delay = -(selected.original_time - mp.get_property_native("time-pos"))
		mp.set_property_native("sub-delay", delay)
	end,
	{},
})

function sub_time_to_seconds(time, sep)
	if time:match("%d%d:%d%d" .. sep .. "%d%d%d") then
		time = "00:" .. time
	end

	local major, minor = time:match("(%d%d:%d%d:%d%d)" .. sep .. "(%d%d%d)")
	local hours, mins, secs = major:match("(%d%d):(%d%d):(%d%d)")
	return hours * 3600 + mins * 60 + secs + minor / 1000
end

local subs_cache = {}

function open_file(path)
	local f, err = io.open(path, "r")
	if f and err == nil then
		return f
	end

	return nil
end

function is_supported_network_protocol(url)
	local protocols = { "http", "https" }

	for _, protocol in pairs(protocols) do
		if url:sub(1, #protocol + 3) == protocol .. "://" then
			return true
		end
	end

	return false
end

function get_sub_filename_async(track_name, on_done)
	local active_track = mp.get_property_native("current-tracks/" .. track_name)
	if active_track == nil then
		on_done(nil)
		return
	end

	local is_external = active_track.external
	local external_filename = active_track["external-filename"]

	-- youtube subtitles specified with edl format
	if is_external and external_filename and external_filename:sub(1, 6) == "edl://" then
		download_subtitle_async(external_filename:match("https://.*"), on_done)
		return
	end

	if is_external and external_filename and is_supported_network_protocol(external_filename) then
		download_subtitle_async(external_filename, on_done)
		return
	end

	if is_external and external_filename then
		on_done(external_filename)
		return
	end

	if is_external == false then
		extract_subtitle_track_async(active_track, on_done)
		return
	end

	on_done(nil)
end

function get_path_to_extract_sub(uniq_sub_id)
	local sub_filename = sha1.hex(uniq_sub_id)
	return utils.join_path(get_temp_dir(), "mpv-subtitle-search-extracted-" .. sub_filename .. ".srt")
end

function download_subtitle_async(url, on_done)
	local sub_path = get_path_to_extract_sub(mp.get_property_native("path") .. "#" .. url)

	if subs_cache[sub_path] then
		on_done(sub_path)
		return
	end

	local extract_overlay = mp.create_osd_overlay("ass-events")
	extract_overlay.data = "{\\a3\\fs20}Fetching remote subtitles, wait..."
	extract_overlay:update()

	mp.command_native_async({
		name = "subprocess",
		capture_stdout = true,
		args = {
			"ffmpeg",
			"-y",
			"-hide_banner",
			"-loglevel",
			"error",
			"-i",
			url,
			"-vn",
			"-an",
			"-c:s",
			"srt",
			sub_path,
		},
	}, function(ok)
		if not ok then
			extract_overlay.data = "{\\a3\\fs20\\c&HFF&}Extraction failed"
			extract_overlay:update()

			mp.add_timeout(2, function()
				extract_overlay:remove()
			end)

			on_done(nil)
		else
			extract_overlay:remove()

			on_done(sub_path)
		end
	end)
end

function extract_subtitle_track_async(track, on_done)
	if track.external then
		on_done(nil)
		return
	end

	local video_file = mp.get_property_native("path")
	local working_dir = mp.get_property_native("working-directory")
	local full_path = utils.join_path(working_dir, video_file)

	local track_index = track["ff-index"]
	local sub_path = get_path_to_extract_sub(full_path .. "#" .. track_index)

	-- check if file already exists
	if open_file(sub_path) then
		msg.info("Reusing extracted subtitle track from " .. sub_path)

		on_done(sub_path)
		return
	end

	msg.info("Extracting embedded subtitle track to " .. sub_path)

	local extract_overlay = mp.create_osd_overlay("ass-events")
	extract_overlay.data = "{\\a3\\fs20}Extracting embedded subtitles, wait..."
	extract_overlay:update()

	mp.command_native_async({
		name = "subprocess",
		capture_stdout = true,
		args = {
			"ffmpeg",
			"-y",
			"-hide_banner",
			"-loglevel",
			"error",
			"-i",
			full_path,
			"-map",
			"0:" .. track_index,
			"-vn",
			"-an",
			"-c:s",
			"srt",
			sub_path,
		},
	}, function(ok)
		if not ok then
			extract_overlay.data = "{\\a3\\fs20\\c&HFF&}Extraction failed"
			extract_overlay:update()

			mp.add_timeout(2, function()
				extract_overlay:remove()
			end)

			on_done(nil)
		else
			extract_overlay:remove()

			on_done(sub_path)
		end
	end)
end

function get_temp_dir()
	local temp_dir = os.getenv("TMPDIR")
	if temp_dir == nil then
		temp_dir = os.getenv("TEMP")
	end

	if temp_dir == nil then
		temp_dir = os.getenv("TMP")
	end

	if temp_dir == nil then
		temp_dir = "/tmp"
	end

	return temp_dir
end

function get_lines(input)
	local lines = {}

	local tail = 1
	for head = 1, #input do
		local ch = input:sub(head, head)
		if ch == "\n" then
			table.insert(lines, input:sub(tail, head - 1))
			tail = head + 1
		elseif head == #input then
			table.insert(lines, input:sub(tail, head))
		end
	end

	return lines
end

function trim(s)
	return s:gsub("^%s*(.-)%s*$", "%1")
end

function parse_vtt_sub(data)
	local result = {}
	local state = "header"

	local cur_line = {}
	for _, line in ipairs(get_lines(data)) do
		line = trim(line)
		if state == "header" then
			if line == "" then
				state = "body"
			end
		elseif state == "body" then
			if line == "" then
				state = "header"
			elseif line:match("^NOTE") or line:match("^STYLE") then
				state = "comment"
			else
				local time_text = line:match("^(%d%d:%d%d:%d%d%.%d%d%d)") or line:match("^(%d%d:%d%d%.%d%d%d)")
				if time_text then
					cur_line.time = sub_time_to_seconds(time_text, ".")
					state = "waiting_text"
				else
					state = "body"
				end
			end
		elseif state == "comment" then
			if #line == 0 then
				state = "body"
			end
		elseif state == "waiting_text" then
			if #line == 0 or line == nil then
				if cur_line.text ~= nil then
					table.insert(result, cur_line)
				end

				cur_line = {}
				state = "body"
			else
				line = remove_tags(line)
				if cur_line.text then
					cur_line.text = cur_line.text .. "\n" .. line
				else
					cur_line.text = line
				end
			end
		end
	end

	return result
end

function remove_tags(text)
	function remove_tag(tag_to_remove)
		return string.gsub(text, "</?" .. tag_to_remove .. ">", "")
	end

	text = remove_tag("b")
	text = remove_tag("i")
	text = remove_tag("u")
	text = remove_tag("ruby")
	text = remove_tag("rt")

	-- remove class tag
	text = remove_tag("c")
	text = string.gsub(text, "<c.[^>]*>", "")

	-- remove voice tag
	text = remove_tag("v")
	text = string.gsub(text, "<v [^>]*>", "")

	-- remove karaoke karaoke tags
	text = string.gsub(text, "</?%d%d:%d%d.%d%d%d>", "")
	text = string.gsub(text, "</?%d%d:%d%d:%d%d.%d%d%d>", "")

	-- remove font tag
	text = string.gsub(text, '<font color="#?[%d%a]+">', "")
	text = string.gsub(text, "</font>", "")

	return text
end

-- detects only most common encodings
function get_encoding_from_bom(data)
	-- utf8
	local bom = data:sub(1, 3)
	if bom == "\xEF\xBB\xBF" then
		return "utf-8"
	end

	-- utf16
	bom = data:sub(1, 2)
	if bom == "\xFF\xFE" or bom == "\xFE\xFF" then
		return "utf-16"
	end

	-- utf32
	bom = data:sub(1, 4)
	if bom == "\xFF\xFE\x00\x00" or bom == "\x00\x00\xFE\xFF" then
		return "utf-32"
	end

	return nil
end

function is_microdvd_sub(data)
	return data:match("{%d+}{%d+}")
end

function parse_microdvd_sub(data)
	local result = {}
	local lines = get_lines(data)

	-- if the first line contains only number, it's a subtitle fps
	local subtitle_fps = tonumber(lines[1])
	if subtitle_fps == nil or subtitle_fps == 0 then
		subtitle_fps = mp.get_property_native("container-fps")
		if subtitle_fps == nil or subtitle_fps == 0 then
			subtitle_fps = 24
		end
	end

	msg.info("Using " .. subtitle_fps .. "fps for microdvd subtitle")

	for _, line in ipairs(lines) do
		local time_text = line:match("^{(%d+)}{(%d+)}")
		if time_text then
			local start_frame = tonumber(time_text:match("^(%d+)"))

			local text = line:match("^{%d+}{%d+}(.*)")
			text = text:gsub("|", " ")
			if text then
				table.insert(result, {
					time = frame_to_secs(start_frame, subtitle_fps),
					text = text,
				})
			end
		end
	end

	return result
end

function frame_to_secs(frame, subtitle_fps)
	return frame / subtitle_fps
end

function parse_sub(data)
	bom_encoding = get_encoding_from_bom(data)
	if bom_encoding ~= nil then
		if bom_encoding == "utf-8" then
			data = data:sub(3)
		else
			local error_overlay = mp.create_osd_overlay("ass-events")
			error_overlay.data = "{\\a3\\fs20\\c&HFF&}Unsupported subtitle encoding: "
				.. bom_encoding
				.. ", please re-encode subtitle file to utf-8 to search"
			error_overlay:update()

			msg.error(
				"Unsupported subtitle encoding: "
					.. bom_encoding
					.. ", please re-encode subtitle file to utf-8 to search"
			)

			mp.add_timeout(10, function()
				error_overlay:remove()
			end)

			return {}
		end
	end

	data = string.gsub(data, "\r\n", "\n")

	if data:sub(1, 6) == "WEBVTT" then
		return parse_vtt_sub(data)
	end

	if is_microdvd_sub(data) then
		return parse_microdvd_sub(data)
	end

	local result = {}
	local state = "waiting_index"
	local cur_line = {}
	for _, line in ipairs(get_lines(data)) do
		line = trim(line)
		if state == "waiting_index" then
			if cur_line.text then
				table.insert(result, cur_line)
				cur_line = {}
			end

			if line:match("^%d+$") then
				state = "waiting_time"
			end
		elseif state == "waiting_time" then
			local time_text = line:match("^(%d%d:%d%d:%d%d,%d%d%d) ")
			if time_text then
				cur_line.time = sub_time_to_seconds(time_text, ",")
				state = "waiting_text"
			else
				state = "waiting_index"
			end
		elseif state == "waiting_text" then
			line = remove_tags(line)
			if #line == 0 then
				if cur_line.text then
					table.insert(result, cur_line)
				end
				cur_line = {}
				state = "waiting_index"
			elseif cur_line.text then
				cur_line.text = cur_line.text .. " " .. line
			else
				cur_line.text = line
			end
		end
	end

	if cur_line.text then
		table.insert(result, cur_line)
	end

	return result
end

function load_sub(path, prefix)
	if not path then
		return nil
	end

	local cached = subs_cache[path]
	if cached then
		return cached
	end

	local f = open_file(path)
	if not f then
		return nil
	end

	local data = f:read("*all")
	f:close()

	local sub = {
		prefix = prefix,
		lines = parse_sub(data),
	}
	subs_cache[path] = sub
	return sub
end

function make_nocase_pattern(s)
	local result = ""
	for _, code in utf8.codes(s) do
		local c = utf8.char(code)
		result = result .. string.format("[%s%s]", utf8.lower(c), utf8.upper(c))
	end
	return result
end

-- highlight found text with colored text in ass syntax
function highlight_match(text, match_text, style_reset)
	local match_start, match_end = utf8.find(utf8.lower(text), utf8.lower(match_text))
	if match_start == nil then
		return text
	end

	local before = result_list.ass_escape(utf8.sub(text, 1, match_start - 1))
	local match = result_list.ass_escape(utf8.sub(text, match_start, match_end))
	local after = result_list.ass_escape(utf8.sub(text, match_end + 1))

	if style_reset == "" then
		style_reset = "{\\c&HFFFFFF&}"
	end

	return before .. "{\\c&HFF00&}" .. match .. style_reset .. after
end

function adjust_sub_time(time)
	local delay = mp.get_property_native("sub-delay")
	if delay == nil then
		return time
	end
	return time + delay
end

function divmod(a, b)
	return math.floor(a / b), a % b
end

function format_time(time)
	decimals = 3
	sep = "."
	local s = time
	local h, s = divmod(s, 60 * 60)
	local m, s = divmod(s, 60)

	local second_format = string.format("%%0%d.%df", 2 + (decimals > 0 and decimals + 1 or 0), decimals)

	return string.format("%02d" .. sep .. "%02d" .. sep .. second_format, h, m, s)
end

function get_subs_to_search_in_async(on_done)
	local result = {}

	get_sub_filename_async("sub", function(primary_filename)
		local sub = load_sub(primary_filename, "P")
		if sub then
			table.insert(result, sub)
		end

		get_sub_filename_async("sub2", function(secondary_filename)
			sub = load_sub(secondary_filename, "S")
			if sub then
				table.insert(result, sub)
			end

			on_done(result)
		end)
	end)
end

function update_search_results_async(query, live)
	get_subs_to_search_in_async(function(subs)
		if #subs == 0 then
			mp.osd_message("External subtitles not found")
			return
		end

		result_list.list = {
			{
				sub = nil,
				time = mp.get_property_native("time-pos"),
				ass = "Original position",
			},
		}
		result_list.selected = 1
		result_list.live = live

		local closest_lower_index = 1
		local closest_lower_time = nil
		local cur_time = mp.get_property_native("time-pos")

		local pat = "(" .. make_nocase_pattern(query) .. ")"
		for _, sub in ipairs(subs) do
			for _, sub_line in ipairs(sub.lines) do
				if query == "*" or utf8.match(sub_line.text, pat) then
					local sub_time = adjust_sub_time(sub_line.time)

					table.insert(result_list.list, {
						sub = sub,
						original_time = sub_line.time,
						time = sub_time + 0.01, -- to ensure that the subtitle is visible
						formatter = function(style_reset)
							local sub_text = result_list.ass_escape(format_time(sub_time) .. ": ")
								.. highlight_match(sub_line.text, query, style_reset)

							if #subs > 1 then
								sub_text = "[" .. sub.prefix .. "] " .. sub_text
							end

							return sub_text
						end,
					})

					if sub_time <= cur_time and (closest_lower_time == nil or closest_lower_time < sub_time) then
						closest_lower_time = sub_time
						closest_lower_index = #result_list.list
					end
				end
			end
		end

		result_list.selected = closest_lower_index
		result_list.header = 'Search results for "' .. query .. '"\\N ------------------------------------'
		result_list.header = result_list.header
			.. "\\NENTER to jump to subtitle, Ctrl+Shift+Enter to adjust subtitle timing to selected line"

		result_list:update()
		result_list:open()
	end)
end

mp.register_script_message("start-search", function()
	if input_console.is_repl_active() then
		input_console.set_active(false)
	else
		input_console.set_enter_handler(function(query)
			update_search_results_async(query, false)
		end)
		input_console.set_active(true)
	end
end)

mp.register_script_message("show-all-lines", function()
	update_search_results_async("*", true)
end)

local function get_current_subtitle_index(list, pos)
	local closest_lower_index = 1
	local closest_lower_time = nil
	for i, item in ipairs(list) do
		if item.time <= pos and (closest_lower_time == nil or closest_lower_time < item.time) then
			closest_lower_time = item.time
			closest_lower_index = i
		end
	end
	return closest_lower_index
end

mp.observe_property("time-pos", "native", function(_, pos)
	if not result_list.hidden and result_list.live and pos ~= nil then
		local index = get_current_subtitle_index(result_list.list, pos)
		if index > 1 then
			result_list.selected = index
			result_list:update()
		end
	end
end)