-- lx-bundle:true

--- Highly experimental WAV reader and worker-analysis proof of concept
-- Pure Lua analysis is used here to test work in a separate Lua state on a real OS thread
-- @module plugins.wavpreview.wav

local wav = {}

-- Keep enough waveform buckets for wide preview panels
local WAVEFORM_BUCKETS = 2048
-- Bound work so long recordings remain practical in this optional pure-Lua preview
local MAX_WAVEFORM_FRAMES = 262144
-- Average a few small transforms instead of freezing on one transform for the whole file
local FFT_SIZE = 1024
local MAX_FFT_WINDOWS = 16
local FFT_DECIBELS = 60
local EXTENSIBLE_TAG = 0xfffe
local PCM_TAG = 1
local FLOAT_TAG = 3
local SUBFORMAT_GUID_TAIL = "\0\0\16\0\128\0\0\170\0\56\155\113"

--- Reject a byte range that extends beyond its source
-- @param data string source bytes
-- @param position number one-based byte position
-- @param size number byte count
-- @param description string range name for errors
local function check_range(data, position, size, description)
	if position < 1 or size < 0 or position + size - 1 > #data then
		error(string.format(
			"%s is outside the WAV data: position=%d size=%d file_size=%d",
			description,
			position,
			size,
			#data
		), 3)
	end
end

--- Read one little-endian unsigned 16-bit value
-- @param data string source bytes
-- @param position number one-based byte position
-- @return number value
local function u16le(data, position)
	check_range(data, position, 2, "uint16")

	return string.unpack("<I2", data, position)
end

--- Read one little-endian unsigned 32-bit value
-- @param data string source bytes
-- @param position number one-based byte position
-- @return number value
local function u32le(data, position)
	check_range(data, position, 4, "uint32")

	return string.unpack("<I4", data, position)
end

--- Resolve ordinary and extensible WAV encoding tags
-- @param chunk string complete format chunk
-- @return number encoding_tag
-- @return number|nil valid_bits
local function get_encoding_tag(chunk)
	local tag = u16le(chunk, 1)

	if tag ~= EXTENSIBLE_TAG then
		return tag
	end
	if #chunk < 40 or u16le(chunk, 17) < 22 then
		error("truncated WAV extensible format")
	end
	if chunk:sub(29, 40) ~= SUBFORMAT_GUID_TAIL then
		error("unsupported WAV extensible subformat")
	end

	return u32le(chunk, 25), u16le(chunk, 19)
end

--- Parse one WAV format chunk
-- @param chunk string format payload
-- @return table format_info
local function parse_format(chunk)
	if #chunk < 16 then
		error("truncated WAV format chunk")
	end

	local encoding_tag, valid_bits = get_encoding_tag(chunk)
	local channels = u16le(chunk, 3)
	local sample_rate = u32le(chunk, 5)
	local block_align = u16le(chunk, 13)
	local bits_per_sample = u16le(chunk, 15)
	local bytes_per_sample = bits_per_sample // 8
	local encoding_name

	if encoding_tag == PCM_TAG and (bits_per_sample == 8 or bits_per_sample == 16
		or bits_per_sample == 24 or bits_per_sample == 32)
	then
		encoding_name = "PCM"
	elseif encoding_tag == FLOAT_TAG and (bits_per_sample == 32 or bits_per_sample == 64) then
		encoding_name = "Float"
	else
		error(string.format(
			"unsupported WAV encoding: format=%d bits=%d",
			encoding_tag,
			bits_per_sample
		))
	end
	if channels < 1 or sample_rate < 1 then
		error("invalid WAV channel count or sample rate")
	end
	if block_align ~= channels * bytes_per_sample then
		error("unsupported WAV sample layout")
	end

	return {
		encoding_tag = encoding_tag,
		encoding_name = encoding_name,
		channels = channels,
		sample_rate = sample_rate,
		block_align = block_align,
		bits_per_sample = bits_per_sample,
		valid_bits = valid_bits,
		bytes_per_sample = bytes_per_sample,
	}
end

--- Decode one PCM sample into the common minus-one to one range
-- @param data string complete file bytes
-- @param position number sample byte position
-- @param bits number stored sample width
-- @return number sample
local function read_pcm_sample(data, position, bits)
	if bits == 8 then
		return (data:byte(position) - 128) / 128
	end
	if bits == 16 then
		return string.unpack("<i2", data, position) / 32768
	end
	if bits == 24 then
		local low, middle, high = data:byte(position, position + 2)
		local value = low + middle * 256 + high * 65536

		if value >= 0x800000 then
			value = value - 0x1000000
		end

		return value / 8388608
	end

	return string.unpack("<i4", data, position) / 2147483648
end

--- Decode one floating-point sample and replace invalid values with silence
-- @param data string complete file bytes
-- @param position number sample byte position
-- @param bits number stored sample width
-- @return number sample
local function read_float_sample(data, position, bits)
	local sample = string.unpack(bits == 32 and "<f" or "<d", data, position)

	if sample ~= sample or sample == math.huge or sample == -math.huge then
		return 0
	end

	return math.max(-1, math.min(1, sample))
end

--- Decode one channel from one zero-based audio frame
-- @param audio table parsed WAV audio
-- @param frame number zero-based frame index
-- @param channel number one-based channel index
-- @return number sample
function wav.get_sample(audio, frame, channel)
	local position = audio.data_position + frame * audio.block_align
		+ (channel - 1) * audio.bytes_per_sample

	if audio.encoding_tag == PCM_TAG then
		return read_pcm_sample(audio.data, position, audio.bits_per_sample)
	end

	return read_float_sample(audio.data, position, audio.bits_per_sample)
end

--- Mix every channel in one frame into one preview sample
-- @param audio table parsed WAV audio
-- @param frame number zero-based frame index
-- @return number sample
local function get_mixed_sample(audio, frame)
	local sum = 0

	for channel = 1, audio.channels do
		sum = sum + wav.get_sample(audio, frame, channel)
	end

	return sum / audio.channels
end

--- Parse one complete little-endian RIFF WAV file
-- @param data string complete file bytes
-- @return table audio
function wav.parse(data)
	assert(type(data) == "string", "WAV data must be a string")
	check_range(data, 1, 12, "RIFF header")

	if data:sub(1, 4) ~= "RIFF" or data:sub(9, 12) ~= "WAVE" then
		error("not a little-endian RIFF WAV file")
	end

	local position = 13
	local format_info
	local data_position
	local data_size

	while position + 7 <= #data do
		local chunk_id = data:sub(position, position + 3)
		local chunk_size = u32le(data, position + 4)
		local payload_position = position + 8
		check_range(data, payload_position, chunk_size, "WAV chunk " .. chunk_id)

		if chunk_id == "fmt " and not format_info then
			format_info = parse_format(data:sub(payload_position, payload_position + chunk_size - 1))
		elseif chunk_id == "data" and not data_position then
			data_position = payload_position
			data_size = chunk_size
		end

		-- RIFF chunks include one padding byte after odd payload sizes
		position = payload_position + chunk_size + chunk_size % 2
	end

	if not format_info then
		error("WAV format chunk is missing")
	end
	if not data_position then
		error("WAV audio data chunk is missing")
	end

	format_info.data = data
	format_info.data_position = data_position
	format_info.data_size = data_size
	format_info.file_size = #data
	format_info.frame_count = data_size // format_info.block_align
	format_info.duration = format_info.frame_count / format_info.sample_rate

	return format_info
end

--- Read and parse one WAV file
-- @param path string WAV path
-- @return table|nil audio
-- @return string|nil error_message
function wav.read_file(path)
	local handle, open_error = io.open(path, "rb")

	if not handle then
		return nil, open_error
	end

	local data = handle:read("*a")
	handle:close()
	local loaded, audio = pcall(wav.parse, data)

	return loaded and audio or nil, loaded and nil or tostring(audio)
end

--- Build a bounded whole-file envelope from every channel
-- @param audio table parsed WAV audio
-- @return table lows
-- @return table highs
function wav.build_waveform(audio)
	local bucket_count = math.min(WAVEFORM_BUCKETS, math.max(1, audio.frame_count))
	local lows = {}
	local highs = {}
	-- Spread the work budget evenly across the complete recording, not only its beginning
	local frame_step = math.max(1, math.ceil(audio.frame_count / MAX_WAVEFORM_FRAMES))

	for frame = 0, audio.frame_count - 1, frame_step do
		local bucket = math.min(bucket_count, frame * bucket_count // audio.frame_count + 1)

		for channel = 1, audio.channels do
			local sample = wav.get_sample(audio, frame, channel)
			lows[bucket] = math.min(lows[bucket] or sample, sample)
			highs[bucket] = math.max(highs[bucket] or sample, sample)
		end
	end

	for bucket = 1, bucket_count do
		lows[bucket] = lows[bucket] or 0
		highs[bucket] = highs[bucket] or 0
	end

	return lows, highs
end

--- Transform one real-valued signal in place
-- @param real table real values
-- @param imaginary table imaginary values
local function fft(real, imaginary)
	local count = #real
	local reversed = 1

	for index = 1, count do
		if index < reversed then
			real[index], real[reversed] = real[reversed], real[index]
			imaginary[index], imaginary[reversed] = imaginary[reversed], imaginary[index]
		end

		local half = count // 2

		while half >= 1 and reversed > half do
			reversed = reversed - half
			half = half // 2
		end

		reversed = reversed + half
	end

	local size = 2

	while size <= count do
		local angle = -2 * math.pi / size
		local step_real = math.cos(angle)
		local step_imaginary = math.sin(angle)

		for start = 1, count, size do
			local weight_real = 1
			local weight_imaginary = 0

			for offset = 0, size // 2 - 1 do
				local even = start + offset
				local odd = even + size // 2
				local odd_real = real[odd] * weight_real - imaginary[odd] * weight_imaginary
				local odd_imaginary = real[odd] * weight_imaginary + imaginary[odd] * weight_real
				real[odd] = real[even] - odd_real
				imaginary[odd] = imaginary[even] - odd_imaginary
				real[even] = real[even] + odd_real
				imaginary[even] = imaginary[even] + odd_imaginary
				local next_real = weight_real * step_real - weight_imaginary * step_imaginary
				weight_imaginary = weight_real * step_imaginary + weight_imaginary * step_real
				weight_real = next_real
			end
		end

		size = size * 2
	end
end

--- Return the largest power of two suitable for one recording
-- @param frame_count number available frames
-- @return number size
local function get_fft_size(frame_count)
	local size = 1

	while size * 2 <= math.min(FFT_SIZE, frame_count) do
		size = size * 2
	end

	return size
end

--- Average transforms across the recording and normalize them for drawing
-- @param audio table parsed WAV audio
-- @return table spectrum values from zero to one
function wav.build_spectrum(audio)
	local size = get_fft_size(audio.frame_count)

	if size < 2 then
		return { 0 }
	end

	local available_windows = math.max(1, audio.frame_count // size)
	local window_count = math.min(MAX_FFT_WINDOWS, available_windows)
	local power = {}

	for window = 1, window_count do
		local maximum_start = audio.frame_count - size
		local start = window_count == 1 and maximum_start // 2
			or math.floor((window - 1) * maximum_start / (window_count - 1))
		local real = {}
		local imaginary = {}
		local mean = 0

		for index = 1, size do
			real[index] = get_mixed_sample(audio, start + index - 1)
			mean = mean + real[index]
		end

		mean = mean / size

		for index = 1, size do
			-- A Hann window keeps sharp window edges out of the displayed spectrum
			real[index] = (real[index] - mean)
				* (0.5 - 0.5 * math.cos(2 * math.pi * (index - 1) / (size - 1)))
			imaginary[index] = 0
		end

		fft(real, imaginary)

		for bin = 1, size // 2 do
			power[bin] = (power[bin] or 0) + real[bin] ^ 2 + imaginary[bin] ^ 2
		end
	end

	local maximum_power = 0

	for bin = 1, #power do
		power[bin] = power[bin] / window_count
		maximum_power = math.max(maximum_power, power[bin])
	end

	if maximum_power == 0 then
		return power
	end

	for bin = 1, #power do
		local decibels = 10 * math.log(power[bin] / maximum_power, 10)
		power[bin] = math.max(0, 1 + decibels / FFT_DECIBELS)
	end

	return power
end

--- Read and analyze one WAV file inside an experimental isolated worker Lua state
-- The raw sample bytes stay in the worker and only compact preview data returns
-- @param path string WAV path
-- @return table analysis
function wav.analyze_file(path)
	local audio, read_error = wav.read_file(path)

	if not audio then
		error(read_error or "WAV file could not be read", 0)
	end

	local waveform_low, waveform_high = wav.build_waveform(audio)
	local spectrum = wav.build_spectrum(audio)
	-- The main state needs metadata and graphs, never the complete sample payload
	audio.data = nil

	return {
		audio = audio,
		waveform_low = waveform_low,
		waveform_high = waveform_high,
		spectrum = spectrum,
	}
end

return wav
