local carp = print

local function warn( ... )
	carp( "WARN: %s", string.format( ... ) )
end

-- for efficiency reasons, the coverage linehook stores its data:
--  [ line ] => { [ chunkname1 ] => count, [ chunkname2 ] => count, ... }
-- which is easier to handle in this code if the mapping is flipped:
-- [ chunkname ] => { [ line1 ] => count, [ line2 ] => count, ... }
-- in chunkname that were detected as covered.

local function transformCoverageLines( linesTable, progress, optResults )
	assert( linesTable )
	progress = progress or function() end
	
	local chunks = optResults or { }
	
	local function chunkFilter( chunk )
		-- be more inclusive, since we apparently have several ways of
		-- naming chunks for Lua files.  Scripts menu entries are loaded
		-- with the @fullpath style, and some other namespaces are coming
		-- in with a chunkname of '  filename = fullpath'.  Odd.

		-- doing an endswith match on a long chunk (like say a large table
		-- serialized into a file) is really slow, and we use a lot of such
		-- tables, so fast filter those out by ignoring anything with a
		-- linefeed in it...
		if not string.find( chunk, "\n" ) then
			-- Lua loaded via require has a different convention.  ^#(*&!
			local withoutExt = string.match( chunk, "^[A-Za-z0-9]*$" )
			if withoutExt then
				return withoutExt .. ".lua"
			else
				return string.match( chunk, "[^/]*%.lua$" )
			end
		end
	end
	
	-- match is slow enough to justify caching when dealing with the
	-- kind of data volumes for line coverage
	local results = { }
	local isWorthCovering = function( chunk )
		local result = results[ chunk ]
		if result == nil then
			result = chunkFilter( chunk )
			results[ chunk ] = result
		end
		
		return result
	end
	
	local total = #linesTable
	for line, lineTable in ipairs( linesTable ) do
		progress( line, total )
		for chunk, count in pairs( lineTable ) do
			local simplename = isWorthCovering( chunk )
			
			if simplename then
				if not chunks[ simplename ] then
					chunks[ simplename ] = { }
				end
				chunks[ simplename ][ line ] = count
			else
				-- decruft the line table. it accumulates chunks related to Lua
				-- table archiving and this function slows down.
				lineTable[ chunk ] = nil
			end
		end
	end

	return chunks
end

--------------------------------------------------------------------------------

local function mode( r )
	if not r then
		return nil
	elseif r[ 0 ] then
		return "instruction"
	elseif r[ 1 ] then
		return "line"
	else
		return nil
	end
end

--------------------------------------------------------------------------------

local function iterateInstructionCoverageData( icData, callback, progress )
	progress = progress or function() end
	
	local total = #icData
	for pc, hasht in pairs( icData ) do
		progress( pc, total )
		for hash, baset in pairs( hasht ) do
			for base, count in pairs( baset ) do
				callback( hash, base, pc )
			end
		end
	end
end

--------------------------------------------------------------------------------

-- this translation code is duplicated in RewriteTraceback.lua in order to make
-- it possible for that script to run in an ordinary Lua interpreter on the CLI
-- rather than needing to run in LR itself.

local hashToChunk, chunkToLines

local function loadSymbolsFile( filename, progress )
	progress = progress or function() end
	local chunkHashes = assert( loadfile( filename ) )

	hashToChunk = { }
	chunkToLines = { }

	local function processChunk(t)
		hashToChunk[tostring(t.hash)] = t.name
		chunkToLines[t.name] = t
		progress( t.name )
	end
	
	setfenv(chunkHashes, { chunk = processChunk })
	chunkHashes()
end

local function lineMapper( hash, base, offset )
	local chunk = hashToChunk[ tostring( hash ) ]

	local chunkBases = chunkToLines[ chunk ]

	if chunkBases then
		local baseOffsets = chunkBases[ base ]
		if baseOffsets then
			local line = baseOffsets[ offset ]
			if line then
				return chunk, line
			else
				return nil, "invalid offset"
			end
		else
			return nil, "unknown base"
		end
	else
		return nil, "unknown chunk"
	end
end

local function instructionToLineCoverage( icData, progress, optResults )
	assert( icData )
	progress = progress or function() end
	
	local results = optResults or { }
	
	local misses = false
	
	if not hashToChunk or not chunkToLines then
		warn( "Lua symbols file must be loaded for instruction coverage." )
		return
	else
		iterateInstructionCoverageData( icData, function( hash, base, pc )
			local chunk, line = lineMapper( hash, base, pc )
			if chunk then
				local chunkLines = results[ chunk ]
				if not chunkLines then
					chunkLines = { }
					results[ chunk ] = chunkLines
				end
				chunkLines[ line ] = 1 -- TODO: retain counts
			elseif hash > 8192 then
				warn( "Suspicious instruction coverage entry: %d:%d+%d", hash, base, pc )
				misses = true
			end
		end, progress )
		
		if misses then
			warn( "Lua symbols file may be mismatched." )
		end
	end
		
	return results
end

--------------------------------------------------------------------------------

-- recursive iterator for all instruction to line mappings
local function foreachLineMapping( f, callback )
	callback( f.lines )
	for i, func in ipairs( f.constants.functions ) do
		callback( func.lines )
		foreachLineMapping( func, callback )
	end
end

-- accumulator function for all active lines in all functions
local function allActiveLines( fcd )
	local allLines = { }
	foreachLineMapping( fcd, function( lines )
		for i, line in ipairs( lines ) do
			allLines[ line ] = true
		end
	end )
	return allLines
end

-- decompile a Lua file and return a table mapping each active line to a true value
local function computeActiveLines( filename, chunk )
	if hashToChunk and chunkToLines then
		if chunk and chunkToLines[ chunk ] then
			local activeLines = { }
			for base, lines in pairs( chunkToLines[ chunk ] ) do
				if type( base ) == "number" then
					for _, line in ipairs( lines ) do
						activeLines[ line ] = true
					end
				end
			end
			return activeLines
		end
	else
		local f = loadfile( filename )
	
		if type( f ) == "function" then
			local fd = debug.describefn( f )
			return allActiveLines( fd )
		end
	end

	return { }
end

local function convertCoverageResults( linesTable, progress, optResults )

	local mode = mode( linesTable )
	if mode == "instruction" then
		return instructionToLineCoverage( linesTable, progress, optResults )
	elseif mode == "line" then
		return transformCoverageLines( linesTable, progress, optResults )
	elseif not linesTable or not next( linesTable ) then
		return { }
	else
		error( "coverage is not enabled" )
	end
end

return {
	mode = mode,
	iterateInstructionCoverageData = iterateInstructionCoverageData,
	transformCoverageLines = transformCoverageLines,
	loadSymbolsFile = loadSymbolsFile,
	computeActiveLines = computeActiveLines,
	convertCoverageResults = convertCoverageResults,
}
