1 files changed, 44 insertions, 31 deletions

diff --git a/main.lua b/main.lua
index f379dd0..06b926a 100644
--- a/main.lua
+++ b/main.lua
@@ -31,8 +31,7 @@ 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 objects = {}
 local favourites
 local playing_stream_id
 
@@ -64,39 +63,58 @@ local function write_json_file(fn, data)
 	f:close()
 end
 
+local function add_object(v)
+	objects[v.id] = v
+
+	if v.type == 'category' then
+		v.children = {}
+	end
+
+	if not v.parent_id or not objects[v.parent_id] then
+		return
+	end
+
+	local parent_children = objects[v.parent_id].children
+	parent_children[#parent_children+1] = v
+end
+
 local function load_section(prefix, name)
-	categories[#categories+1] = {
-		id=prefix .. ':0',
+	add_object({
+		id=prefix .. ':category: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.id = prefix .. ':' .. v.type .. ':' .. v.category_id
 		v.name = v.category_name
-		v.parent_id = prefix .. ':' .. v.parent_id
-		categories[#categories+1] = v
+		v.parent_id = prefix .. ':' .. v.type .. ':' .. v.parent_id
+		add_object(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
+		v.id = prefix .. ':' .. v.type .. ':' .. v.stream_id
+		v.parent_id = prefix .. ':category:' .. v.category_id
+		add_object(v)
 	end
 end
 
 local function load_data()
-	categories[#categories+1] = {
+	add_object({
+		id='root',
+		type='category',
+	})
+	add_object({
 		id='favourites',
 		type='category',
 		name='Favourites',
 		parent_id='root',
-	}
+	})
 
 	load_section('live', 'Live TV')
 	load_section('movie', 'Movies')
@@ -168,7 +186,7 @@ local function update_osd()
 			icons = icons .. colours.icon_playing .. '\226\143\186 '
 		end
 
-		if favourites[opt.type .. ':' .. opt.id] then
+		if favourites[opt.id] then
 			icons = icons .. colours.icon_favourite .. '★ '
 		end
 
@@ -262,22 +280,18 @@ local function push_menu(t)
 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
+	if category_id == 'favourites' then
+		local options = {}
+		for id in pairs(favourites) do
+			v = objects[id]
+			if v then
 				options[#options+1] = v
 			end
 		end
+		return options
 	end
-	return options
+
+	return objects[category_id].children
 end
 
 local function push_category_menu(category_id, title)
@@ -308,13 +322,12 @@ end
 
 local function favourite_option()
 	local menu = menus[depth]
-	local opt = menu.options[menu.cursor]
+	local id = menu.options[menu.cursor].id
 
-	local key = opt.type .. ':' .. opt.id
-	if favourites[key] then
-		favourites[key] = nil
+	if favourites[id] then
+		favourites[id] = nil
 	else
-		favourites[key] = true
+		favourites[id] = true
 	end
 
 	write_json_file('favourites.json', favourites)