-- lx-bundle:true

-- This file is part of Lx, created by Liqube <www.liqube.com>
-- Copyright Liqube Audio. All rights reserved.
-- Lx Interim License. See LICENSE for details.

--- Minimal legacy Word DOC text extractor
-- Reads WordDocument stream plus 0Table/1Table CLX piece table
-- @module plugins.msoutils.doc

local cfb = require("plugins.msoutils.cfb")
local codepage = require("core.text.codepage")

local reader = {}

local function bytez(data, offset)
	return data:byte(offset + 1) or 0
end

local function u16z(data, offset)
	return bytez(data, offset) | (bytez(data, offset + 1) << 8)
end

local function u32z(data, offset)
	return bytez(data, offset)
		| (bytez(data, offset + 1) << 8)
		| (bytez(data, offset + 2) << 16)
		| (bytez(data, offset + 3) << 24)
end

local function i32z(data, offset)
	local value = u32z(data, offset)
	if value >= 0x80000000 then
		return value - 0x100000000
	end
	return value
end

local function substrz(data, offset, len)
	if len <= 0 then
		return ""
	end
	return data:sub(offset + 1, offset + len)
end

-- Primary Windows language ids that share one shipped legacy codepage
local language_codepages = {
	[0x01] = "cp1256", [0x02] = "cp1251", [0x03] = "cp1252", [0x05] = "cp1250",
	[0x06] = "cp1252", [0x07] = "cp1252", [0x08] = "cp1253", [0x09] = "cp1252",
	[0x0a] = "cp1252", [0x0b] = "cp1252", [0x0c] = "cp1252", [0x0d] = "cp1255",
	[0x0e] = "cp1250", [0x0f] = "cp1252", [0x10] = "cp1252", [0x11] = "cp932",
	[0x12] = "cp949", [0x13] = "cp1252", [0x14] = "cp1252", [0x15] = "cp1250",
	[0x16] = "cp1252", [0x18] = "cp1250", [0x19] = "cp1251", [0x1b] = "cp1250",
	[0x1c] = "cp1250", [0x1d] = "cp1252", [0x1e] = "cp874", [0x1f] = "cp1254",
	[0x20] = "cp1256", [0x21] = "cp1252", [0x22] = "cp1251", [0x23] = "cp1251",
	[0x24] = "cp1250", [0x25] = "cp1257", [0x26] = "cp1257", [0x27] = "cp1257",
	[0x29] = "cp1256", [0x2a] = "cp1258", [0x2d] = "cp1252", [0x2f] = "cp1251",
	[0x36] = "cp1252", [0x38] = "cp1252", [0x3e] = "cp1252", [0x3f] = "cp1251",
	[0x44] = "cp1251",
}

--- Choose the Windows codepage used by one Word language id
-- @param lid integer Word language id
-- @return string|nil name shipped codepage name
local function codepage_from_language(lid)
	local exact_chinese = {
		[0x0404] = "cp950",
		[0x0804] = "cp936",
		[0x0c04] = "cp950",
		[0x1004] = "cp936",
		[0x1404] = "cp950",
	}
	if exact_chinese[lid] then
		return exact_chinese[lid]
	end
	if lid == 0x0c1a or lid == 0x1c1a or lid == 0x201a or lid == 0x301a then
		return "cp1251"
	end
	if lid == 0x082c or lid == 0x0843 then
		return "cp1251"
	end
	if lid == 0 then
		return "cp1252"
	end

	return language_codepages[lid & 0x03ff]
end

--- Resolve the codepage for compressed text from FIB language and charset data
-- @param fib table parsed Word file information block
-- @return string|nil name codepage name
-- @return string|nil error_message unsupported metadata error
local function resolve_compressed_codepage(fib)
	if fib.charset ~= 0 or fib.table_charset ~= 0 then
		return nil, string.format(
			"unsupported DOC character set %d/%d",
			fib.charset,
			fib.table_charset
		)
	end

	local language = fib.lid
	local far_east_primary = fib.lid_fe and (fib.lid_fe & 0x03ff) or 0
	if far_east_primary == 0x04 or far_east_primary == 0x11 or far_east_primary == 0x12 then
		language = fib.lid_fe
	end
	local name = codepage_from_language(language)
	if not name then
		return nil, string.format("unsupported DOC language 0x%04x", language)
	end
	local mapping, mapping_error = codepage.resolve(name)
	if not mapping then
		return nil, mapping_error
	end

	return name
end

local function safe_utf8_char(code)
	if not code or code == 0 then
		return ""
	end
	local ok, value = pcall(utf8.char, code)
	return ok and value or ""
end

local function decode_utf16le(data)
	local out = {}
	local i = 1

	while i < #data do
		local lo, hi = data:byte(i, i + 1)
		local code = lo | (hi << 8)
		i = i + 2

		if code >= 0xd800 and code <= 0xdbff and i < #data then
			local lo2, hi2 = data:byte(i, i + 1)
			local trail = lo2 | (hi2 << 8)
			if trail >= 0xdc00 and trail <= 0xdfff then
				code = 0x10000 + ((code - 0xd800) << 10) + (trail - 0xdc00)
				i = i + 2
			end
		end

		out[#out + 1] = safe_utf8_char(code)
	end

	return table.concat(out)
end

local function parse_fib(word)
	local ident = u16z(word, 0)
	if ident == 0xa5dc then
		local fc_min = u32z(word, 0x18)
		local fc_mac = u32z(word, 0x1c)
		local ccp_text = i32z(word, 0x34)

		if fc_min >= fc_mac or fc_mac > #word then
			return nil, "invalid Word 6/95 text bounds"
		end

		return {
			format = "word6",
			ccp_text = math.max(0, ccp_text),
			fc_min = fc_min,
			fc_mac = fc_mac,
			lid = u16z(word, 0x06),
			lid_fe = 0,
			charset = u16z(word, 0x14),
			table_charset = u16z(word, 0x16),
		}
	end

	if ident ~= 0xa5ec then
		return nil, string.format("unsupported WordDocument FIB magic 0x%04x", ident)
	end

	local pos = 0
	pos = pos + 2 -- wIdent
	pos = pos + 2 -- nFib
	pos = pos + 2 -- unused
	local lid = u16z(word, pos)
	pos = pos + 2
	pos = pos + 2 -- pnNext
	local flags = u16z(word, pos)
	pos = pos + 2

	local encrypted = ((flags >> 8) & 1) ~= 0
	if encrypted then
		return nil, "encrypted DOC files are not supported"
	end

	pos = pos + 2 -- nFibBack
	pos = pos + 4 -- lKey
	pos = pos + 1 -- envr
	pos = pos + 1 -- flags2
	local charset = u16z(word, pos)
	pos = pos + 2
	local table_charset = u16z(word, pos)
	pos = pos + 2
	pos = pos + 4 -- reserved5
	pos = pos + 4 -- reserved6
	pos = pos + 2 -- csw
	local lid_fe = u16z(word, pos + (13 * 2))
	pos = pos + (14 * 2) -- fibRgW
	pos = pos + 2 -- cslw
	pos = pos + 4 -- cbMac
	pos = pos + 4 -- reserved1
	pos = pos + 4 -- reserved2
	local ccp_text = i32z(word, pos)
	pos = pos + 4
	pos = pos + (18 * 4) -- remaining fibRgLw fields through reserved14

	local cb_rg_fc_lcb = u16z(word, pos)
	pos = pos + 2
	if cb_rg_fc_lcb < 34 then
		return nil, "FIB table is too small for CLX"
	end

	local fc_clx_offset = pos + ((34 - 1) * 8)
	local fc_clx = u32z(word, fc_clx_offset)
	local lcb_clx = u32z(word, fc_clx_offset + 4)

	return {
		format = "word97",
		ccp_text = math.max(0, ccp_text),
		which_table = ((flags >> 9) & 1) == 1 and "1Table" or "0Table",
		fc_clx = fc_clx,
		lcb_clx = lcb_clx,
		lid = lid,
		lid_fe = lid_fe,
		charset = charset,
		table_charset = table_charset,
	}
end

--- Parse one bounded Word piece table
-- @param table_stream string 0Table or 1Table bytes
-- @param fc_clx integer CLX byte offset
-- @param lcb_clx integer CLX byte length
-- @return table|nil pieces parsed pieces
-- @return string|nil error_message malformed table error
local function parse_piece_table(table_stream, fc_clx, lcb_clx)
	if fc_clx > #table_stream or lcb_clx > #table_stream - fc_clx then
		return nil, "DOC CLX exceeds the table stream"
	end

	local clx = substrz(table_stream, fc_clx, lcb_clx)
	local pos = 0

	while pos < #clx do
		local marker = bytez(clx, pos)
		pos = pos + 1

		if marker == 0x01 then
			if #clx - pos < 2 then
				return nil, "truncated DOC property block length"
			end
			local cb_prc = u16z(clx, pos)
			pos = pos + 2
			if cb_prc > #clx - pos then
				return nil, "truncated DOC property block"
			end
			pos = pos + cb_prc
		elseif marker == 0x02 then
			if #clx - pos < 4 then
				return nil, "truncated DOC piece table length"
			end
			local lcb = u32z(clx, pos)
			pos = pos + 4
			if lcb < 4 or (lcb - 4) % 12 ~= 0 then
				return nil, "invalid DOC piece table shape"
			end
			if lcb > #clx - pos then
				return nil, "truncated DOC piece table"
			end
			local count = (lcb - 4) // 12
			if count <= 0 then
				return nil, "empty DOC piece table"
			end

			local cps = {}
			for i = 1, count + 1 do
				cps[i] = u32z(clx, pos)
				pos = pos + 4
			end

			local pieces = {}
			for i = 1, count do
				pos = pos + 2 -- pcd flags
				local fc_compressed = u32z(clx, pos)
				pos = pos + 4
				pos = pos + 2 -- prm
				pieces[i] = {
					cp_start = cps[i],
					cp_end = cps[i + 1],
					compressed = (fc_compressed & 0x40000000) ~= 0,
					fc = fc_compressed,
				}
			end

			return pieces
		else
			return nil, "unsupported DOC CLX marker " .. tostring(marker)
		end
	end

	return nil, "DOC piece table not found"
end

local function piece_offset(piece)
	if piece.compressed then
		return (piece.fc & 0x3fffffff) // 2
	end
	return piece.fc
end

local function decode_piece(word, piece, from_cp, to_cp, codepage_name)
	local skip_chars = from_cp - piece.cp_start
	local char_count = to_cp - from_cp
	if char_count <= 0 then
		return ""
	end

	local offset = piece_offset(piece)
	if piece.compressed then
		return codepage.decode(
			substrz(word, offset + skip_chars, char_count),
			codepage_name,
			utf8.char(0xfffd)
		)
	end

	return decode_utf16le(substrz(word, offset + (skip_chars * 2), char_count * 2))
end

local function cleanup_text(text)
	text = text:gsub("\r\n", "\n"):gsub("\r", "\n")
	text = text:gsub("\007", "\t")
	text = text:gsub("\011", "\n")
	text = text:gsub("\012", "\n\n")
	text = text:gsub("[%z\001-\006\008\014-\031]", "")
	text = text:gsub("[ \t]+\n", "\n")
	text = text:gsub("\n\n\n+", "\n\n")
	return text:match("^%s*(.-)%s*$") or text
end

local function load_document(filename)
	local archive, err = cfb.open(filename)
	if not archive then
		return nil, err
	end

	local word, word_err = archive:read_stream("WordDocument")
	if not word then
		return nil, word_err or "WordDocument stream not found"
	end

	local fib, fib_err = parse_fib(word)
	if not fib then
		return nil, fib_err
	end

	if fib.format == "word6" then
		local codepage_name, codepage_error = resolve_compressed_codepage(fib)
		if not codepage_name then
			return nil, codepage_error
		end
		local decoded, _, decode_error = codepage.decode(
			substrz(word, fib.fc_min, fib.fc_mac - fib.fc_min),
			codepage_name,
			utf8.char(0xfffd)
		)
		if not decoded then
			return nil, decode_error
		end

		return {
			text = cleanup_text(decoded),
			fib = fib,
		}
	end

	local table_stream = archive:read_stream(fib.which_table)
		or archive:read_stream(fib.which_table == "1Table" and "0Table" or "1Table")
	if not table_stream then
		return nil, "DOC table stream not found"
	end

	local pieces, pieces_err = parse_piece_table(table_stream, fib.fc_clx, fib.lcb_clx)
	if not pieces then
		return nil, pieces_err
	end
	local compressed_codepage
	for _, piece in ipairs(pieces) do
		if piece.compressed then
			local codepage_error
			compressed_codepage, codepage_error = resolve_compressed_codepage(fib)
			if not compressed_codepage then
				return nil, codepage_error
			end
			break
		end
	end

	local out = {}
	for _, piece in ipairs(pieces) do
		local from_cp = math.max(piece.cp_start, 0)
		local to_cp = math.min(piece.cp_end, fib.ccp_text)
		if to_cp > from_cp then
			local decoded, _, decode_error = decode_piece(word, piece, from_cp, to_cp, compressed_codepage)
			if not decoded then
				return nil, decode_error
			end
			out[#out + 1] = decoded
		end
	end

	return {
		text = cleanup_text(table.concat(out)),
		fib = fib,
	}
end

function reader.read(filename)
	return load_document(filename)
end

function reader.extract_text(filename)
	local doc, err = load_document(filename)
	if not doc then
		return nil, err
	end
	return doc.text
end

function reader.extract_markdown(filename)
	local doc, err = load_document(filename)
	if not doc then
		return nil, err
	end
	return doc.text
end

return reader
