-- Copyright 2025 David Vazgenovich Shakaryan

local utils = require('mp.utils')

-- font size is in units of osd height, which is scaled to 720
local font_size = 20
local colours = {
	default='{\\c}',
	category='{\\c&H99DDFF&}',
	selected='{\\c&HFF00&}',
	title='{\\c&999999&}',
	search_hl='{\\c&FFDD&}',
	search_path='{\\c&666666&}',
	icon_playing='{\\c&HFF6633&}',
	icon_favourite='{\\c&HFF00FF&}',
	bg='{\\alpha&H44&\\c&H&}',
}

local script_dir = mp.get_script_directory()
local stream_prefix = mp.get_opt('iptv_menu.stream_prefix')

local osd = mp.create_osd_overlay('ass-events')
local osd_lines = math.floor((720 / font_size) + 0.5) - 1
local osd_padding = math.floor((720 - (osd_lines * font_size)) / 2)
local osd_cursor_glyph = '{\\p1\\pbo' .. math.floor(font_size / 5) .. '}' ..
	'm 0 0 l ' .. math.ceil(font_size / 32) .. ' 0 ' ..
	math.ceil(font_size / 32) .. ' ' .. font_size ..
	' 0 ' .. font_size .. '{\\p0}'
local osd_bg = mp.create_osd_overlay('ass-events')
osd_bg.z = -1
osd_bg.data = '{\\pos(0,0)}' .. colours.bg ..
	'{\\p1}m 0 0 l 7680 0 7680 720 0 720{\\p0}'

local categories = {}
local streams = {}
local favourites
local playing_stream_id

local depth = 0
local menus = {}
local key_bindings = {}

local function copy_table(t)
	local u = {}
	for k, v in pairs(t) do
		u[k] = v
	end
	return u
end

local function read_json_file(fn)
	local f = io.open(script_dir .. '/' .. fn, 'r')
	if not f then
		return {}
	end
	local json = f:read('*all')
	f:close()
	return utils.parse_json(json)
end

local function write_json_file(fn, data)
	local f = io.open(script_dir .. '/' .. fn, 'w')
	f:write(utils.format_json(data), '\n')
	f:close()
end

local function load_section(prefix, name)
	categories[#categories+1] = {
		id=prefix .. ':0',
		type='category',
		name=name,
		parent_id='root',
	}

	local tmp = read_json_file(prefix .. '_categories.json')
	for _, v in ipairs(tmp) do
		v.type = 'category'
		v.id = prefix .. ':' .. v.category_id
		v.name = v.category_name
		v.parent_id = prefix .. ':' .. v.parent_id
		categories[#categories+1] = v
	end

	tmp = read_json_file(prefix .. '_streams.json')
	for _, v in ipairs(tmp) do
		v.type = 'stream'
		v.id = v.stream_id
		v.parent_id = prefix .. ':' .. v.category_id
		streams[#streams+1] = v
	end
end

local function load_data()
	categories[#categories+1] = {
		id='favourites',
		type='category',
		name='Favourites',
		parent_id='root',
	}

	load_section('live', 'Live TV')
	load_section('movie', 'Movies')

	favourites = read_json_file('favourites.json')
	-- json loading/dumping breaks when the table is empty, so we need a
	-- dummy value to prevent that
	if next(favourites) == nil then
		favourites = {oi=true}
	end
end

local function osd_menu_lines()
	if depth > 1 then
		-- leaves an extra line for padding between titles and options
		return osd_lines - depth
	else
		return osd_lines
	end
end

local function update_osd()
	local out = {}

	if depth > 1 then
		for i = 2, depth do
			local col = colours.title
			if menus[i].search_active then
				col = colours.selected
			end
			out[#out+1] = col .. ' » [' .. menus[i].title .. ']'
		end
		out[#out+1] = ' ' -- space character for correct line height
	end

	local menu = menus[depth]
	for i = menu.view_top, math.min(
			menu.view_top + osd_menu_lines() - 1, #menu.options) do
		local opt = menu.options[i]
		local str = opt.name
		local col = colours.default
		local icons = ''

		if i == menu.cursor and not menu.search_active then
			col = colours.selected
			icons = col .. '› ' .. icons
		elseif opt.type == 'category' then
			col = colours.category
		end

		if opt.matches then
			local buf = ''
			local n = 0
			for _, match in ipairs(opt.matches) do
				buf = buf .. col ..
					str:sub(n + 1, match.start - 1) ..
					colours.search_hl ..
					str:sub(match.start, match.stop)
				n = match.stop
			end
			str = buf .. col .. str:sub(n + 1)
		else
			str = col .. str
		end

		if opt.type == 'category' then
			str = col .. '[' .. str .. ']'
		elseif opt.stream_id == playing_stream_id then
			icons = icons .. colours.icon_playing .. '\226\143\186 '
		end

		if favourites[opt.type .. ':' .. opt.id] then
			icons = icons .. colours.icon_favourite .. '★ '
		end

		if opt.path and #opt.path > 0 then
			local path = colours.search_path
			for i = #opt.path, 1, -1 do
				local node = opt.path[i]
				path = path .. ' « [' .. node.name .. ']'
			end
			str = str .. path
		end

		out[#out+1] = icons .. str
	end

	-- \q2 disables line wrapping
	osd.data = '{\\q2\\fs' .. font_size .. '\\pos(' .. osd_padding ..
		',' .. osd_padding .. ')}' .. table.concat(out, '\\N')
	osd:update()
	osd_bg:update()
end

local function set_cursor(pos, move_view)
	local menu = menus[depth]
	local lines = osd_menu_lines()

	local pos = math.max(1, math.min(pos, #menu.options))
	local top = menu.view_top
	if move_view then
		top = top + pos - menu.cursor
	end

	-- move view to keep selected option visible
	if pos < top then
		top = pos
	elseif pos > top + lines - 1 then
		top = pos - lines + 1
	end

	top = math.max(1, math.min(top, #menu.options - lines + 1))

	menu.cursor = pos
	menu.view_top = top
	update_osd()
end

local function cursor_up()
	set_cursor(menus[depth].cursor - 1)
end

local function cursor_down()
	set_cursor(menus[depth].cursor + 1)
end

local function cursor_start()
	set_cursor(0)
end

local function cursor_end()
	set_cursor(#menus[depth].options)
end

local function cursor_page_up()
	set_cursor(menus[depth].cursor - osd_menu_lines(), true)
end

local function cursor_page_down()
	set_cursor(menus[depth].cursor + osd_menu_lines(), true)
end

local function prev_menu()
	if depth > 1 then
		depth = depth - 1
		update_osd()
	end
end

local function push_menu(t)
	local menu = {
		options={},
		cursor=1,
		view_top=1,
	}

	for k, v in pairs(t) do
		menu[k] = v
	end

	depth = depth + 1
	menus[depth] = menu
end

local function category_menu_options(category_id)
	local options = {}
	for _, arr in ipairs({categories, streams}) do
		for _, v in ipairs(arr) do
			local test
			if category_id == 'favourites' then
				test = favourites[v.type .. ':' .. v.id]
			else
				test = v.parent_id == category_id
			end

			if test then
				options[#options+1] = v
			end
		end
	end
	return options
end

local function push_category_menu(category_id, title)
	push_menu({
		options=category_menu_options(category_id),
		title=title,
	})
	update_osd()
end

local function play_stream(stream_id)
	-- add a per-file option containing the stream_id, allowing it to be
	-- retrieved when a start-file event is received
	mp.commandv('loadfile', stream_prefix .. stream_id, 'replace', -1,
		'script-opt=iptv_menu.tmp.stream_id=' .. stream_id)
end

local function select_option()
	local menu = menus[depth]
	local opt = menu.options[menu.cursor]

	if opt.type == 'category' then
		push_category_menu(opt.id, opt.name)
	else
		play_stream(opt.stream_id)
	end
end

local function favourite_option()
	local menu = menus[depth]
	local opt = menu.options[menu.cursor]

	local key = opt.type .. ':' .. opt.id
	if favourites[key] then
		favourites[key] = nil
	else
		favourites[key] = true
	end

	write_json_file('favourites.json', favourites)
	update_osd()
end

local function split_search_text()
	local menu = menus[depth]
	return menu.search_text:sub(1, menu.search_cursor - 1),
		menu.search_text:sub(menu.search_cursor)
end

local function update_search_osd()
	local bc, ac = split_search_text()
	menus[depth].title = 'Searching: ' .. bc .. osd_cursor_glyph .. ac
	update_osd()
end

local function search_menu_options_build(options, t, path)
	local menu = menus[depth]
	local path = path or {}

	for _, v in ipairs(options) do
		if not t[v.type] then
			t[v.type] = {}
		end

		local v = copy_table(v)
		v.path = path
		t[v.type][#t[v.type]+1] = v

		if v.type == 'category' and v.id ~= 'favourites' then
			local path = copy_table(path)
			path[#path+1] = v
			search_menu_options_build(
				category_menu_options(v.id), t, path)
		end
	end
end

local function search_menu_options(options)
	local t = {}
	search_menu_options_build(options, t)

	-- display categories before streams
	local ret = t.category or {}
	if t.stream then
		for _, v in ipairs(t.stream) do
			ret[#ret+1] = v
		end
	end
	return ret
end

local function update_search_matches()
	local menu = menus[depth]

	if #menu.search_text == 0 then
		menu.options = menu.search_options
		update_search_osd()
		return
	end

	local options = {}
	for _, v in ipairs(menu.search_options) do
		local matches = {}

		local i, j = 0, 0
		while true do
			i, j = v.name:find(menu.search_text, j + 1, true)
			if not i then
				break
			end
			matches[#matches+1] = {start=i, stop=j}
		end

		if #matches > 0 then
			local t = copy_table(v)
			t.matches = matches
			options[#options+1] = t
		end
	end

	menu.options = options
	update_search_osd()
end

local function search_input_text(event)
	if event.event ~= 'down' and event.event ~= 'repeat' then
		return
	end

	local menu = menus[depth]
	local bc, ac = split_search_text()
	menu.search_text = bc .. event.key_text .. ac
	menu.search_cursor = menu.search_cursor + 1
	update_search_matches()
end

local function search_input_bs()
	local menu = menus[depth]
	if menu.search_cursor == 1 then
		return
	end

	local bc, ac = split_search_text()
	menu.search_text = bc:sub(1, -2) .. ac
	menu.search_cursor = menu.search_cursor - 1
	update_search_matches()
end

local function search_input_del()
	local menu = menus[depth]
	if menu.search_cursor == #menu.search_text + 1 then
		return
	end

	local bc, ac = split_search_text()
	menu.search_text = bc .. ac:sub(2)
	update_search_matches()
end

local function set_search_cursor(pos)
	local menu = menus[depth]
	local pos = math.max(1, math.min(#menu.search_text + 1, pos))
	if pos == menu.search_cursor then
		return
	end

	menu.search_cursor = pos
	update_search_osd()
end

local function search_cursor_left()
	set_search_cursor(menus[depth].search_cursor - 1)
end

local function search_cursor_right()
	set_search_cursor(menus[depth].search_cursor + 1)
end

local function search_cursor_start()
	set_search_cursor(1)
end

local function search_cursor_end()
	set_search_cursor(#menus[depth].search_text + 1)
end

local bind_search_keys
local bind_menu_keys

local function start_search()
	local menu = menus[depth]
	if menu.type == 'search' then
		-- resume search
		menu.search_active = true
		menu.search_cursor = #menu.search_text + 1
		menu.cursor = 1
		menu.view_top = 1
		update_search_osd()
	else
		push_menu({
			type='search',
			search_active=true,
			search_options=search_menu_options(menu.options),
			search_text='',
			search_cursor=1,
		})
		update_search_matches()
	end

	bind_search_keys()
end

local function end_search()
	local menu = menus[depth]
	menu.search_active = false
	menu.title = 'Search results: ' .. menu.search_text
	update_osd()
	bind_menu_keys()
end

local function cancel_search()
	menus[depth].search_active = false
	depth = depth - 1
	update_osd()
	bind_menu_keys()
end

local function bind_key(key, func, opts)
	-- unique name is needed for removal
	local i = #key_bindings+1
	local name = 'key' .. i
	key_bindings[i] = name
	mp.add_forced_key_binding(key, name, func, opts)
end

local function unbind_keys()
	for _, key in ipairs(key_bindings) do
		mp.remove_key_binding(key)
	end
	key_bindings = {}
end

function bind_search_keys()
	unbind_keys()

	bind_key('ANY_UNICODE', search_input_text, {complex=true})
	bind_key('BS', search_input_bs, {repeatable=true})
	bind_key('DEL', search_input_del, {repeatable=true})

	bind_key('ENTER', end_search)
	bind_key('ESC', cancel_search)
	bind_key('Ctrl+c', cancel_search)

	bind_key('LEFT', search_cursor_left, {repeatable=true})
	bind_key('RIGHT', search_cursor_right, {repeatable=true})
	bind_key('Ctrl+a', search_cursor_start)
	bind_key('Ctrl+e', search_cursor_end)
end

function bind_menu_keys()
	unbind_keys()

	bind_key('BS', prev_menu)
	bind_key('/', start_search)

	bind_key('ENTER', select_option)
	bind_key('Ctrl+f', favourite_option)

	bind_key('UP', cursor_up, {repeatable=true})
	bind_key('DOWN', cursor_down, {repeatable=true})
	bind_key('HOME', cursor_start)
	bind_key('END', cursor_end)
	bind_key('PGUP', cursor_page_up, {repeatable=true})
	bind_key('PGDWN', cursor_page_down, {repeatable=true})
end

local function toggle_menu()
	osd.hidden = not osd.hidden
	osd:update()
	osd_bg.hidden = osd.hidden
	osd_bg:update()

	if osd.hidden then
		unbind_keys()
	elseif menus[depth].search_active then
		bind_search_keys()
	else
		bind_menu_keys()
	end
end

mp.register_event('start-file', function()
	playing_stream_id = tonumber(mp.get_opt('iptv_menu.tmp.stream_id'))
	update_osd()
end)

mp.register_event('end-file', function()
	playing_stream_id = nil
	update_osd()
end)

mp.add_forced_key_binding('TAB', 'toggle-menu', toggle_menu)
bind_menu_keys()
load_data()
push_category_menu('root')