local TERM_COLOR_CODE_WHITE = '\27[37;1m'
local TERM_COLOR_CODE_GREEN = '\27[32;1m'
local TERM_COLOR_CODE_RED   = '\27[31;1m'
local TERM_COLOR_CODE_GREY  = '\27[30;0m'

local ASSERT_ERROR_TYPE = [[
Types not compatible:
    Expected: %s (%s)
    Gotten:   %s (%s)
]]

local ASSERT_ERROR_VALUE = [[
Values of type '%s' not equal:
    Expected: %s
    Gotten:   %s
]]

local ASSERT_ERROR_TABLE_VALUE = [[
Values in tables not equal:
    For key:  %s

    Expected: %s (%s)
    Gotten:   %s (%s)
]]

local function assert_equal (expected, gotten)
    if type(gotten) ~= type(expected) then
        error(ASSERT_ERROR_TYPE:format(tostring(gotten), type(gotten), tostring(expected), type(expected)))
    elseif type(expected) == 'table' and gotten ~= expected then
        for key, expected_value in pairs(expected) do
            if expected_value ~= gotten[key] then
                error(ASSERT_ERROR_TABLE_VALUE:format(key, expected_value, type(expected_value), gotten[key], type(gotten[key])))
            end
        end
    elseif gotten ~= expected then
        error(ASSERT_ERROR_VALUE:format(type(gotten), gotten, expected))
    end
end

local libraries_tests_and_wrappers = {
    assert_equal = assert_equal
}

for k, v in pairs(_G) do
    libraries_tests_and_wrappers[k] = v
end

local function indent_string (str, indent)
    return indent .. str:gsub('\n', '\n'..indent)
end

local TEST_SUITE_TRACEBACK = "\n\t\t[C]: in function 'xpcall'\n\t\t./TestSuite.lua"

local function find_cutoff_index_for_traceback_string (str)
    local index = str:find(TEST_SUITE_TRACEBACK, 1, true)
    return index
end

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

local TestSuite = {}
      TestSuite.__index = TestSuite

function TestSuite.new (module_name)
	local  new_test_suite = setmetatable({}, TestSuite)
           new_test_suite.name = module_name
           new_test_suite.submodules = {}
           new_test_suite.tests = {}
           new_test_suite.custom_env = {}
    return new_test_suite
end

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

function TestSuite:setEnviroment (new_enviroment)
    self.custom_env = new_enviroment
end

function TestSuite:createUniqueEnviroment ()
    local new_enviroment = {}
    for k, v in pairs(libraries_tests_and_wrappers) do
        new_enviroment[k] = v
    end
    for k, v in pairs(self.custom_env) do
        new_enviroment[k] = v
    end
    return new_enviroment
end

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

function TestSuite:addModules (modules)
    for _, module in ipairs(modules) do
        table.insert(self.submodules, module)
    end
end

function TestSuite:addTest (test_name, extra_info, test_func)
    if type(extra_info) == 'function' then
        extra_info, test_func = test_func, extra_info
    end

    table.insert(self.tests, {
        name = test_name,
        test = test_func,
        extra_info = extra_info or {}
    })
end

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

function TestSuite:runSubmodules (prefix, indent)
    local total_errors, total_tests = 0, 0
    for _, module in ipairs(self.submodules) do
        local errors, tests = module:runTests(prefix, indent)
        total_errors, total_tests = total_errors + errors, total_tests + tests
    end
    return total_errors, total_tests
end

local function setup_debug_hooks (extra_info)
    assert(not (extra_info.max_time and extra_info.max_lines), "TEST CONFIG ERROR: Only one line hook allowed pr. test!")
    local trace
    if extra_info.max_time then
        local stop_time = os.clock() + extra_info.max_time
        trace = function (l)
            if os.clock() >= stop_time then
                debug.sethook(nil, 'l')
                error("Test timed out! This is usually symptom of an infinite loop!")
            end
        end
    elseif extra_info.max_lines then
        local line_countdown = 10 + extra_info.max_lines
        trace = function (l)
            line_countdown = line_countdown - 1
            if line_countdown <= 0 then
                debug.sethook(nil, 'l')
                error("Test exteeded allowed number of lines! This is usually symptom of an infinite loop! (Or too low an estimate!)")
            end
        end
    end
    debug.sethook(trace, 'l')
end

function TestSuite:runTests (parent_prefix, parent_indent)
    local prefix, indent = self.name, (parent_indent or 0) + 2
    if parent_prefix then
        prefix = parent_prefix .. '.' .. prefix
    else
        io.write('\n## Running tests ##\n\n')
    end

    local nr_errors, nr_tests = self:runSubmodules(prefix, indent)

    local function error_handler (err)
        return err..debug.traceback('', 4)
    end

    local function print_status (left, right, color)
        local term_width = 80
        local color = color or ''
        io.write(TERM_COLOR_CODE_WHITE..left..': '..string.rep(' ', term_width + 1-#left-#right)..color..right..'\n'..TERM_COLOR_CODE_GREY)
    end

    for index, test in ipairs(self.tests) do
        nr_tests = nr_tests + 1
        -- Setup Test Env, and name
        local ext_name = prefix .. '.' .. test.name
        local env = self:createUniqueEnviroment()
        setfenv(test.test, env)

        setup_debug_hooks(test.extra_info)
        -- Call tests
        local success, traceback = xpcall(test.test, error_handler)
        -- Unset line hook, if set earlier
        debug.sethook(nil, 'l')
        -- Write work (or not.)
        if success then
            print_status(ext_name, 'SUCCESS!', TERM_COLOR_CODE_GREEN)
        else
            print_status(ext_name, 'ERROR!', TERM_COLOR_CODE_RED)
            traceback = indent_string(traceback, '\t')
            local stop_index = find_cutoff_index_for_traceback_string(traceback)
            io.write('\n'..TERM_COLOR_CODE_GREY)
            io.write(traceback:sub(1, stop_index))
            io.write('\n\n')
            nr_errors = nr_errors + 1
        end
    end

    if nr_errors == 0 then
        print_status(prefix, 'NO ERRORS! WELL DONE!', TERM_COLOR_CODE_GREEN)
    elseif nr_errors == 1 then
        print_status(prefix, 'A SINGLE ERROR! ALMOST THERE!', TERM_COLOR_CODE_RED)
    else
        print_status(prefix, nr_errors..' ERRORS! GET TO WORK!', TERM_COLOR_CODE_RED)
    end

    if not parent_prefix then
        io.write(TERM_COLOR_CODE_WHITE)
        io.write('\n## All tests run! ##\n\n')

        local width_of_bar       = 70
        local size_of_red_bar    = math.ceil( width_of_bar * nr_errors/nr_tests)
        local size_of_green_bar  = width_of_bar-size_of_red_bar

        io.write('Status:   ')
        io.write(nr_errors == 0 and TERM_COLOR_CODE_GREEN or TERM_COLOR_CODE_RED)
        io.write(string.rep('#', width_of_bar))
        io.write(TERM_COLOR_CODE_WHITE)

        -- Get
        io.write('\n\nSummary:  '..TERM_COLOR_CODE_GREEN..string.rep('#', size_of_green_bar)..TERM_COLOR_CODE_RED..string.rep('#', size_of_red_bar)..'\n\n')
        io.write(TERM_COLOR_CODE_WHITE)
    end

    return nr_errors, nr_tests
end

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

local function get_end_of_function (text, start_i)
    local indent, i, end_i = 0, start_i, text:len()
    while i <= end_i do
        if text:sub(i,i) == '"' then        -- "style" strings
            i = text:find('"', i + 1) + 1
        elseif text:sub(i,i) == "'" then    -- 'style' strings
            i = text:find("'", i + 1) + 1
        elseif text:sub(i,i+1) == '[[' then -- [[style]] strings
            i = text:find(']]', i) + 2
        elseif text:sub(i,i+3) == '--[[' then -- multi line comments
            i = text:find(']]', i) + 2
        elseif text:sub(i,i+1) == '--' then -- single line comments
            i = text:find('\n', i) + 1
        elseif text:sub(i-1, i+8):find('%sfunction%s') or i == start_i and text:sub(i, i+7) == 'function' then
            indent, i = indent + 1, i + 8
        elseif text:sub(i-1, i+5):find('%swhile%s') or i == start_i and text:sub(i, i+4) == 'while' then
            indent, i = indent + 1, i + 4
        elseif text:sub(i-1, i+3):find('%sfor%s') or i == start_i and text:sub(i, i+2) == 'for' then
            indent, i = indent + 1, i + 2
        elseif text:sub(i-1, i+2):find('%sif%s') or i == start_i and text:sub(i, i+1) == 'if' then
            indent, i = indent + 1, i + 2
        elseif text:sub(i-1, i+3):find('%send%s') or (i == end_i - 2) and text:sub(i, i+2) == 'end' then
            indent, i = indent - 1, i + 2
            if indent == 0 then
                return i
            end
        else
            i = i + 1
        end
    end
end

function TestSuite.getLocalFunction (module_name, func_name)
    local file = io.open(module_name .. '.lua')
    local file_text = file:read('*a')
    file:close()

    local function_start_index  = file_text:find('function '..func_name)
    local function_end_index    = get_end_of_function(file_text, function_start_index)

    local new_text = file_text:sub(1, function_end_index):gsub('local function '..func_name, 'return function')
    assert(new_text ~= file_text, ('No function in module "%s" with name "%s" '):format(module_name, func_name))
    local chunk, error_msg = load(new_text)
    if not chunk then
        error(('While loading function in module "%s" with name "%s":\n\n\t%s\n'):format(module_name, func_name, error_msg))
    end
    assert(chunk, error)
    return chunk()
end

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

return TestSuite