path: root/lib/attr.lua

#!/usr/bin/env lua

local util = require("lib.util")

local extract = function(name, patterns)
	local r

	for _, p in pairs(patterns) do
		r = r or name:match(p)
		name = name:gsub(p, "")
	end

	return name, r
end

local build_dlim = function(name, attrs)
	local dlim, r

	attrs = attrs or {}
	dlim = "[%-%.%s]?"
	local vcodecs = {
		"[M]m[Pp][Ee][Gg]" .. dlim .. "[1234]",
	}
	local acodecs = {
		"[Dd][Tt][Ss]" .. dlim .. "[Hh][Dd]",
		"[Dd][Dd]" .. dlim .. "[57]%.1",
	}
	local sources = {
		"[Ww][Ee][Bb]" .. dlim .. "[Dd][Ll]",
		"[Hh][Dd]" .. dlim .. "[Tt][Vv]",
		"[Hh][Dd]" .. dlim .. "[Tt][Ss]",
	}
	local series = {
		"[Ss]%d%d?" .. dlim .. "[Ee]%d%d?",
		"[Ss]%d%d?" .. dlim .. "[Ee][Pp]" .. dlim .. "%d%d?",
		"Season" .. dlim .. "%d%d?" .. dlim .. "Episode" .. dlim .. "%d%d?",
	}
	local sizes = {
		"%d%d%d" .. dlim .. "[Mm][Bb]",
		"%d%d?%.%d%d?" .. dlim .. "[Gg][Bb]",
	}
	local depths = {
		"1[02]" .. dlim .. "[Bb][Ii][Tt]",
	}

	name, attrs.vcodec = extract(name, vcodecs)
	name, attrs.source = extract(name, sources)
	name, attrs.acodecs = extract(name, acodecs)
	name, attrs.size = extract(name, sizes)
	name, attrs.depth = extract(name, depths)

	name, r = extract(name, series)
	if r then
		attrs.season = tonumber(r:match("%d+"))
		attrs.episode = tonumber(r:match("%d+$"))
	end

	return name, attrs
end

local build_atom = function(name, attrs)
	local r, year

	attrs = attrs or {}
	local vcodecs = {
		"[Aa][Vv]1",
		"[xXHh]26[345]",
		"[Aa][Vv][Cc]",
		"[Hh][Ee][Vv][Cc]",
		"[Xx][Vv][Ii][Dd]",
	}
	local acodecs = {
		"[Oo][Pp][Uu][Ss]",
		"[Aa][Aa][Cc]",
		"[Ee]?[Aa][Cc]3",
		"[Dd][Tt][Ss]",
	}
	local sources = {
		"[Bb][Ll][Uu][Rr][Aa][Yy]",
		"[Bb][Rr][Rr][Ii][Pp]",
		"[Dd][Vv][Dd][Rr][Ii][Pp]",
		"[Ww][Ee][Bb][Rr][Ii][Pp]",
		"[Hh][Dd][Rr][Ii][Pp]",
		"[Rr][Ee][Rr][Ii][Pp]",
	}
	local reses = {
		"2160[Pp]",
		"1440[Pp]",
		"1080[Pp]",
		"720[Pp]",
		"480[Pp]",
		"[Uu][Hh][Dd]",
		"4[Kk]",
	}
	local series = {
		"%d%d[Xx]%d%d",
	}
	local channels = {
		"6[Cc][Hh]",
		"[57]%.1",
	}

	name, attrs.vcodec = extract(name, vcodecs)
	name, attrs.source = extract(name, sources)
	name, attrs.res = extract(name, reses)
	name, attrs.acodecs = extract(name, acodecs)
	name, attrs.channel = extract(name, channels)

	name, r = extract(name, series)
	if r then
		attrs.season = tonumber(r:match("%d+"))
		attrs.episode = tonumber(r:match("%d+$"))
	end

	for y in name:gmatch("%d%d%d%d") do
		year = tonumber(y)
		if year > 1900 and year <= tonumber(os.date("%Y")) then
			attrs.year = y
		end
	end
	if attrs.year then
		name = name:gsub(tostring(attrs.year), "")
	end

	return name, attrs
end

local build_low = function(name, attrs)
	local low_attr, lows

	lows = { "SDH" }

	low_attr = {}
	for _, low in pairs(lows) do
		low_attr[#low_attr + 1] = name:match(low)
		name = name:gsub(low, "")
	end

	attrs = attrs or {}
	if #low_attr > 0 then
		attrs.low = low_attr
	end

	return name, attrs
end

local build_title = function(name, attrs)
	attrs.title = {}
	for w in name:gmatch("%w+") do
		attrs.title[#attrs.title + 1] = w
	end

	if #attrs.title > 1 then
		attrs.scene = attrs.title[#attrs.title]
		attrs.title[#attrs.title] = nil
	end

	return attrs
end

local build = function(name)
	local attrs = {}

	name = build_dlim(name, attrs)
	name = build_atom(name, attrs)
	name = build_low(name, attrs)
	build_title(name, attrs)

	return attrs
end

local weigh = function(a1, a2)
	local key_score, score

	key_score = {
		season = 15,
		episode = 15,
		title = 10,
		source = 7,
		scene = 5,
		vcodec = 3,
		acodec = 3,
		rese = 2,
		default = 1,
	}

	score = 0
	for k, v in pairs(a1) do
		if not a2[k] or k == "episode" then
			goto continue
		end

		if k == "title" then
			for _, title in pairs(v) do
				if util.array_search(a2.title, title) then
					score = score + key_score.title
				end
			end
		elseif k == "season" then
			if a1.season == a2.season and a1.episode == a2.episode then
				score = score + key_score.season + key_score.episode
			end
		else
			if v == a2[k] then
				score = score + (key_score[k] or key_score.default)
			end
		end

		::continue::
	end

	return score
end

local fuzzy = function(name, tab)
	local name_attr, high, score

	high = {
		score = 0,
		name = next(tab),
	}

	name_attr = build(name)
	for k in pairs(tab) do
		score = weigh(name_attr, build(k))
		if score > high.score then
			high.score = score
			high.name = k
		end
	end

	return tab[high.name]
end

return {
	build = build,
	fuzzy = fuzzy,
}