Module:FamilyTree

Documentation for this module may be created at Module:FamilyTree/doc

local p = {}

local cargo = mw.ext.cargo

local function esc(value)
	if not value then
		return ''
	end
	value = tostring(value)
	value = value:gsub('\\', '\\\\')
	value = value:gsub('"', '\\"')
	return value
end

local function cargoQuery(tables, fields, args)
	args = args or {}
	local ok, result = pcall(function()
		return cargo.query(tables, fields, args)
	end)

	if ok and result then
		return result
	end

	return {}
end

local function trim(s)
	if s == nil then
		return nil
	end
	s = tostring(s)
	s = mw.text.trim(s)
	if s == '' then
		return nil
	end
	return s
end

local function addUnique(list, seen, value)
	value = trim(value)
	if value and not seen[value] then
		seen[value] = true
		table.insert(list, value)
	end
end

local function addSet(set, value)
	value = trim(value)
	if value then
		set[value] = true
	end
end

local function sorted(list)
	table.sort(list, function(a, b)
		return tostring(a):lower() < tostring(b):lower()
	end)
	return list
end

local function formatYear(dateValue)
	dateValue = trim(dateValue)
	if not dateValue then
		return nil
	end
	return tostring(dateValue):match('^(%d%d%d%d)')
end

local function getDisplayName(pageName)
	pageName = trim(pageName)
	if not pageName then
		return nil
	end

	local rows = cargoQuery(
		'Characters',
		'Page,DisplayName',
		{
			where = 'Page="' .. esc(pageName) .. '"',
			limit = 1
		}
	)

	if rows[1] and trim(rows[1].DisplayName) then
		return trim(rows[1].DisplayName)
	end

	return pageName
end

local function getCharacter(pageName)
	pageName = trim(pageName)
	if not pageName then
		return nil
	end

	local rows = cargoQuery(
		'Characters',
		'Page,DisplayName,BirthDate,DeathDate,Status,Gender',
		{
			where = 'Page="' .. esc(pageName) .. '"',
			limit = 1
		}
	)

	return rows[1]
end

local function makeLinkedName(pageName)
	local displayName = getDisplayName(pageName)
	return '[[' .. pageName .. '|' .. displayName .. ']]'
end

local function linkList(list)
	local out = {}
	for _, pageName in ipairs(list) do
		table.insert(out, makeLinkedName(pageName))
	end
	return table.concat(out, '<br>')
end

local function getParents(person)
	person = trim(person)
	if not person then
		return {}, nil
	end

	local rows = cargoQuery(
		'ParentChild',
		'Child,Parent1,Parent2,UnionID,RelationshipType,BirthOrder',
		{
			where = 'Child="' .. esc(person) .. '"',
			limit = 20
		}
	)

	local parents = {}
	local seen = {}

	for _, row in ipairs(rows) do
		addUnique(parents, seen, row.Parent1)
		addUnique(parents, seen, row.Parent2)
	end

	return sorted(parents), rows[1]
end

local function getChildren(person)
	person = trim(person)
	if not person then
		return {}
	end

	local rows = cargoQuery(
		'ParentChild',
		'Child,Parent1,Parent2,UnionID,RelationshipType,BirthOrder',
		{
			where = 'Parent1="' .. esc(person) .. '" OR Parent2="' .. esc(person) .. '"',
			limit = 200
		}
	)

	table.sort(rows, function(a, b)
		local aOrder = tonumber(a.BirthOrder) or 9999
		local bOrder = tonumber(b.BirthOrder) or 9999
		if aOrder == bOrder then
			return tostring(a.Child):lower() < tostring(b.Child):lower()
		end
		return aOrder < bOrder
	end)

	local children = {}
	local seen = {}

	for _, row in ipairs(rows) do
		addUnique(children, seen, row.Child)
	end

	return children
end

local function getPartners(person)
	person = trim(person)
	if not person then
		return {}, {}
	end

	local rows = cargoQuery(
		'Unions',
		'UnionID,Partner1,Partner2,UnionType,Status,MarriageDate,DivorceDate,EngagementDate',
		{
			where = 'Partner1="' .. esc(person) .. '" OR Partner2="' .. esc(person) .. '"',
			limit = 50
		}
	)

	local partners = {}
	local seen = {}

	for _, row in ipairs(rows) do
		local p1 = trim(row.Partner1)
		local p2 = trim(row.Partner2)

		if p1 == person then
			addUnique(partners, seen, p2)
		elseif p2 == person then
			addUnique(partners, seen, p1)
		end
	end

	return sorted(partners), rows
end

local function getUnionBetween(personA, personB)
	personA = trim(personA)
	personB = trim(personB)

	if not personA or not personB then
		return nil
	end

	local rows = cargoQuery(
		'Unions',
		'UnionID,Partner1,Partner2,UnionType,Status,MarriageDate,DivorceDate,EngagementDate',
		{
			where = '(Partner1="' .. esc(personA) .. '" AND Partner2="' .. esc(personB) .. '") OR (Partner1="' .. esc(personB) .. '" AND Partner2="' .. esc(personA) .. '")',
			limit = 1
		}
	)

	return rows[1]
end

local function getSiblings(person)
	person = trim(person)
	if not person then
		return {}
	end

	local targetRows = cargoQuery(
		'ParentChild',
		'Child,Parent1,Parent2,UnionID',
		{
			where = 'Child="' .. esc(person) .. '"',
			limit = 10
		}
	)

	if not targetRows[1] then
		return {}
	end

	local targetUnion = trim(targetRows[1].UnionID)
	local targetP1 = trim(targetRows[1].Parent1)
	local targetP2 = trim(targetRows[1].Parent2)

	local whereParts = {}

	if targetUnion then
		table.insert(whereParts, 'UnionID="' .. esc(targetUnion) .. '"')
	end

	if targetP1 and targetP2 then
		table.insert(whereParts, '(Parent1="' .. esc(targetP1) .. '" AND Parent2="' .. esc(targetP2) .. '")')
	end

	if #whereParts == 0 then
		return {}
	end

	local rows = cargoQuery(
		'ParentChild',
		'Child,Parent1,Parent2,UnionID,RelationshipType,BirthOrder',
		{
			where = table.concat(whereParts, ' OR '),
			limit = 100
		}
	)

	table.sort(rows, function(a, b)
		local aOrder = tonumber(a.BirthOrder) or 9999
		local bOrder = tonumber(b.BirthOrder) or 9999
		if aOrder == bOrder then
			return tostring(a.Child):lower() < tostring(b.Child):lower()
		end
		return aOrder < bOrder
	end)

	local siblings = {}
	local seen = {}

	for _, row in ipairs(rows) do
		local child = trim(row.Child)
		if child and child ~= person then
			addUnique(siblings, seen, child)
		end
	end

	return siblings
end

local function getNeighbors(person)
	local neighbors = {}

	person = trim(person)
	if not person then
		return neighbors
	end

	local unionRows = cargoQuery(
		'Unions',
		'UnionID,Partner1,Partner2',
		{
			where = 'Partner1="' .. esc(person) .. '" OR Partner2="' .. esc(person) .. '"',
			limit = 500
		}
	)

	for _, row in ipairs(unionRows) do
		addSet(neighbors, row.Partner1)
		addSet(neighbors, row.Partner2)
	end

	local parentChildRows = cargoQuery(
		'ParentChild',
		'Child,Parent1,Parent2,UnionID',
		{
			where = 'Child="' .. esc(person) .. '" OR Parent1="' .. esc(person) .. '" OR Parent2="' .. esc(person) .. '"',
			limit = 1000
		}
	)

	for _, row in ipairs(parentChildRows) do
		addSet(neighbors, row.Child)
		addSet(neighbors, row.Parent1)
		addSet(neighbors, row.Parent2)
	end

	neighbors[person] = nil
	return neighbors
end

local function collectConnectedComponent(root)
	local visited = {}
	local queue = {}
	local head = 1

	root = trim(root)
	if not root then
		return visited
	end

	visited[root] = true
	table.insert(queue, root)

	while head <= #queue do
		local current = queue[head]
		head = head + 1

		local neighbors = getNeighbors(current)

		for neighbor, _ in pairs(neighbors) do
			if neighbor and not visited[neighbor] then
				visited[neighbor] = true
				table.insert(queue, neighbor)
			end
		end
	end

	return visited
end

local function makeCard(pageName)
	pageName = trim(pageName)
	if not pageName then
		return ''
	end

	local c = getCharacter(pageName)
	local displayName = getDisplayName(pageName)
	local birthYear = c and formatYear(c.BirthDate) or nil
	local deathYear = c and formatYear(c.DeathDate) or nil

	local years = ''
	if birthYear or deathYear then
		years = '<div class="ft-years">' .. (birthYear or '?') .. '–' .. (deathYear or '') .. '</div>'
	end

	return '<div class="ft-card">[[' .. pageName .. '|' .. displayName .. ']]' .. years .. '</div>'
end

local function makePersonRow(people)
	if not people or #people == 0 then
		return ''
	end

	local html = {}
	table.insert(html, '<div class="ft-row">')
	for _, person in ipairs(people) do
		table.insert(html, makeCard(person))
	end
	table.insert(html, '</div>')

	return table.concat(html)
end

local function buildCoupleGroups(people)
	local groups = {}
	local used = {}
	local sortedPeople = {}

	for _, person in ipairs(people or {}) do
		table.insert(sortedPeople, person)
	end
	sorted(sortedPeople)

	for _, person in ipairs(sortedPeople) do
		if not used[person] then
			local partners = getPartners(person)
			local matchedPartner = nil

			for _, partner in ipairs(partners) do
				for _, candidate in ipairs(sortedPeople) do
					if candidate == partner and not used[candidate] then
						matchedPartner = partner
						break
					end
				end
				if matchedPartner then
					break
				end
			end

			if matchedPartner then
				used[person] = true
				used[matchedPartner] = true

				local union = getUnionBetween(person, matchedPartner)
				table.insert(groups, {
					people = { person, matchedPartner },
					marriageYear = union and formatYear(union.MarriageDate) or nil
				})
			else
				used[person] = true
				table.insert(groups, {
					people = { person },
					marriageYear = nil
				})
			end
		end
	end

	return groups
end

local function buildRootCoupleGroups(root, partners)
	local groups = {}

	if not partners or #partners == 0 then
		table.insert(groups, {
			people = { root },
			marriageYear = nil
		})
	else
		for _, partner in ipairs(partners) do
			local union = getUnionBetween(root, partner)
			table.insert(groups, {
				people = { root, partner },
				marriageYear = union and formatYear(union.MarriageDate) or nil
			})
		end
	end

	return groups
end

local function makeCoupleCards(groups)
	if not groups or #groups == 0 then
		return ''
	end

	local html = {}
	table.insert(html, '<div class="ft-row">')

	for _, group in ipairs(groups) do
		local people = group.people or {}
		local marriageYear = group.marriageYear

		if #people == 1 then
			table.insert(html, '<div class="ft-single">')
			table.insert(html, makeCard(people[1]))
			table.insert(html, '</div>')
		else
			table.insert(html, '<div class="ft-couple">')
			table.insert(html, makeCard(people[1]))
			table.insert(html, '<div class="ft-marriage">')
			table.insert(html, marriageYear and ('<div class="ft-marriage-year">' .. marriageYear .. '</div>') or '')
			table.insert(html, '<div class="ft-marriage-line"></div>')
			table.insert(html, '</div>')
			table.insert(html, makeCard(people[2]))
			table.insert(html, '</div>')
		end
	end

	table.insert(html, '</div>')
	return table.concat(html)
end

local function makeSectionTitle(title)
	return '<div class="ft-section-title">' .. title .. '</div>'
end

function p.connected(frame)
	local args = frame.args
	local parentArgs = frame:getParent() and frame:getParent().args or {}
	local root = trim(args.root or parentArgs.root or args[1] or parentArgs[1])

	if not root then
		return 'Error: no root provided. Use root=Character Name'
	end

	local component = collectConnectedComponent(root)
	local people = {}

	for name, _ in pairs(component) do
		table.insert(people, name)
	end
	sorted(people)

	if #people == 0 then
		return 'No connected people found for ' .. root
	end

	local lines = {}
	table.insert(lines, "'''Connected component for " .. getDisplayName(root) .. "'''")
	table.insert(lines, '* Total people found: ' .. tostring(#people))

	for _, person in ipairs(people) do
		table.insert(lines, '* ' .. makeLinkedName(person))
	end

	return table.concat(lines, '\n')
end

function p.profile(frame)
	local args = frame.args
	local parentArgs = frame:getParent() and frame:getParent().args or {}
	local root = trim(args.root or parentArgs.root or args[1] or parentArgs[1])

	if not root then
		return 'Error: no root provided. Use root=Character Name'
	end

	local parents = getParents(root)
	local siblings = getSiblings(root)
	local partners = getPartners(root)
	local children = getChildren(root)

	local lines = {}
	table.insert(lines, '{| class="wikitable" style="width:100%; max-width:900px;"')
	table.insert(lines, '|-')
	table.insert(lines, '! colspan="2" | Family profile for ' .. getDisplayName(root))
	table.insert(lines, '|-')
	table.insert(lines, '! style="width:20%;" | Person')
	table.insert(lines, '| ' .. makeLinkedName(root))

	table.insert(lines, '|-')
	table.insert(lines, '! Parents')
	table.insert(lines, '| ' .. (#parents > 0 and linkList(parents) or '—'))

	table.insert(lines, '|-')
	table.insert(lines, '! Siblings')
	table.insert(lines, '| ' .. (#siblings > 0 and linkList(siblings) or '—'))

	table.insert(lines, '|-')
	table.insert(lines, '! Partners')
	table.insert(lines, '| ' .. (#partners > 0 and linkList(partners) or '—'))

	table.insert(lines, '|-')
	table.insert(lines, '! Children')
	table.insert(lines, '| ' .. (#children > 0 and linkList(children) or '—'))

	table.insert(lines, '|}')

	return table.concat(lines, '\n')
end

function p.tree(frame)
	local args = frame.args
	local parentArgs = frame:getParent() and frame:getParent().args or {}
	local root = trim(args.root or parentArgs.root or args[1] or parentArgs[1])

	if not root then
		return 'Error: no root provided. Use root=Character Name'
	end

	local parents = getParents(root)
	local partners = getPartners(root)
	local siblings = getSiblings(root)
	local children = getChildren(root)

	local grandparents = {}
	local gpSeen = {}

	for _, parentName in ipairs(parents) do
		local parentParents = getParents(parentName)
		for _, gp in ipairs(parentParents) do
			addUnique(grandparents, gpSeen, gp)
		end
	end
	sorted(grandparents)

	local grandparentGroups = buildCoupleGroups(grandparents)
	local parentGroups = buildCoupleGroups(parents)
	local rootGroups = buildRootCoupleGroups(root, partners)

	local html = {}
	table.insert(html, '<div class="ft-tree">')
	table.insert(html, '<div class="ft-title">Family tree for ' .. getDisplayName(root) .. '</div>')

	if #grandparents > 0 then
		table.insert(html, '<div class="ft-generation">')
		table.insert(html, makeCoupleCards(grandparentGroups))
		table.insert(html, '</div>')
	end

	if #grandparents > 0 and #parents > 0 then
		table.insert(html, '<div class="ft-connector"></div>')
	end

	if #parents > 0 then
		table.insert(html, '<div class="ft-generation">')
		table.insert(html, makeCoupleCards(parentGroups))
		table.insert(html, '</div>')
	end

	if #parents > 0 then
		table.insert(html, '<div class="ft-connector"></div>')
	end

	table.insert(html, '<div class="ft-generation">')
	table.insert(html, makeCoupleCards(rootGroups))
	table.insert(html, '</div>')

	if #siblings > 0 then
		table.insert(html, '<div class="ft-generation ft-extra-generation">')
		table.insert(html, makeSectionTitle('Siblings'))
		table.insert(html, makePersonRow(siblings))
		table.insert(html, '</div>')
	end

	if #children > 0 then
		table.insert(html, '<div class="ft-connector"></div>')
		table.insert(html, '<div class="ft-generation ft-extra-generation">')
		table.insert(html, makeSectionTitle('Children'))
		table.insert(html, makePersonRow(children))
		table.insert(html, '</div>')
	end

	table.insert(html, '</div>')
	return table.concat(html)
end

return p