require 'erb'
require 'yaml'

require 'rake/clean'
require 'rake/testtask'

require 'fileutils'

require_relative 'zemu/config'

load '../z80-libraries/tasks.rake'

BUILTINS = Dir.glob("builtin/*").map { |b| [b, File.basename(b)] }

MAX_ALLOCS = 200000

CLEAN.include(
    "**/*.noi",
    "**/*.lk",
    "**/*.rst",
    "**/*.rel",
    "**/*.lst",
    "**/*.o",
    "**/*.bin",
    "**/*.hex",
    "**/*.diss",
    "**/*.lib",
    "**/*.lis",
    "**/*.map",
    "**/*.sym",
    "**/*.exe",
    "**/*.log",
    "**/*.coverage")

# Compiles a single C file into an SDCC .rel (object) file.
def compile(source, output, defines, includes)
    cmd = "sdcc -c -mz80 --sdcccall 1 --opt-code-speed --max-allocs-per-node #{MAX_ALLOCS} --std-sdcc99 "
    cmd += "-o\"#{output}\" "
    
    includes.each { |i| cmd += "-I#{i} " }
    defines.each { |d| cmd += "-D#{d} " }

    cmd += "--Werror "
    
    cmd += "#{source}"
    puts cmd
    system(cmd)
    FileUtils.rm("#{File.dirname(output)}/#{File.basename(output, ".*")}.asm")
end

def assemble(source, output)
    cmd = "sdasz80 -plosgffw "
    cmd += "#{output} "
    cmd += source
    puts cmd
    system(cmd)
end

rule ".rel" => ".c" do |task|
    compile(task.source, task.name, %w(Z80), [File.dirname(task.source), LIB_INCLUDE])
end

rule ".debug.rel" => ".c" do |task|
    compile(task.source, task.name, %w(Z80 DEBUG), [File.dirname(task.source), LIB_INCLUDE])
end

rule ".rel" => ".asm" do |task|
    assemble(task.source, task.name)
end

rule ".debug.rel" => ".asm" do |task|
    assemble(task.source, task.name)
end

def get_addr(map, label)
    if /([0-9a-fA-F]+)\s+#{label}/ =~ File.read(map)
        return $1.to_i(16)
    else
        abort("Could not find label '#{label}' in map file '#{map}'")
    end
end

def valid_memory_map(map_file)
    code_start = get_addr(map_file, "s__CODE")
    code_end   = code_start + get_addr(map_file, "l__CODE")
    data_start = get_addr(map_file, "s__DATA")
    data_end   = data_start + get_addr(map_file, "l__DATA")

    # Assumes data is always located after code.
    if data_start < code_end
        return false
    end

    true
end

def compile_sdcc(output_name, object_files, crt0, code_seg, data_seg)
    basename = File.basename(output_name, ".bin")

    cmd = "sdcc --verbose -mz80 --sdcccall 1 --opt-code-speed --max-allocs-per-node #{MAX_ALLOCS} --std-sdcc99 "
    cmd += "--no-std-crt0 "
    cmd += "-Wl-b_CODE=0x#{code_seg.to_s(16)} "
    cmd += "-Wl-b_DATA=0x#{data_seg.to_s(16)} "
    cmd += "-o #{basename}.hex "
    cmd += "-L#{LIB} "
    cmd += "#{crt0} "
    cmd += object_files.join(" ")
    cmd += " stdlib.lib"
    
    puts cmd
    success = system("#{cmd} > link.log 2>&1")

    unless success
        abort("Link failed:\n#{File.read("link.log")}")
    end

    # Check memory layout is valid.
    # sdld will happily assign two objects of different types to the same memory
    # which causes us a headache if we don't notice!
    unless valid_memory_map("#{basename}.map")
        abort("Invalid memory map for program '#{basename}'!")
    end

    system("objcopy --gap-fill 0x76 -Iihex -Obinary #{basename}.hex #{output_name}")
    
    size = File.size(output_name)
    size_kb = size / 1024.0
    puts "#{output_name}: #{size_kb.round(1)}KB"

    true
end

# Helper function for running unit tests.
# Delete test exe and main.c if they exist.
def unit_test(directory)
    # Do nothing if there is no unit test directory
    unless Dir.exist? "#{directory}/unit_test"
        puts "No unit tests for #{directory}"
        return
    end
    
    exe_name = "test_" + directory.gsub("/", "_") + ".exe"

    if File.exist? exe_name
        FileUtils.rm exe_name
    end

    if File.exist? "#{directory}/unit_test/main.c"
        FileUtils.rm "#{directory}/unit_test/main.c"
    end

    tests = []

    # Get tests from test files.
    Dir.glob("#{directory}/unit_test/*.c").each do |f|
        File.readlines(f).each do |l|
            if /int (test_\S+)\(\)/ =~ l
                tests << $1
            end
        end
    end

    # Filter based on environment variable (if any)
    unless ENV["UNIT_TEST_NAME"].nil?
        tests = tests.select { |t| t.include?(ENV["UNIT_TEST_NAME"]) }
    end

    # Create the templated main.
    template = File.read("#{directory}/unit_test/main.c.erb")
    erb = ERB.new(template)
    File.write("#{directory}/unit_test/main.c", erb.result(binding))

    # Compile test/main.c, other test/*.c, and kernel files for which
    # a file exists under test.
    # Include every source file but reset.asm, which is the startup code.
    src = []

    Dir.glob("#{directory}/unit_test/*.c").each do |f|
        next if f == "#{directory}/unit_test/main.c"

        src << f
    end

    Dir.glob("#{directory}/*.c").each do |f|
        next if f == "#{directory}/main.c"
        
        src << f
    end

    # Add mock files
    Dir.glob("#{directory}/unit_test/mock/*.c").each do |f|
        src << f
    end

    cmd = "gcc "
    cmd += "-Wall -Werror -Wno-unused-value -Wno-unknown-pragmas -Wno-discarded-qualifiers -Wno-format "
    cmd += "-g -O1 "
    cmd += "-DUNIT_TEST "
    cmd += "-I#{directory} "
    cmd += "-I#{directory}/unit_test/mock "
    cmd += "-o #{exe_name} "
    cmd += src.join(" ")
    cmd += " #{directory}/unit_test/main.c"

    system(cmd)

    # Now run tests.
    if (!system("./#{exe_name}"))
        puts "TESTS FAILED"
    else
        puts "TESTS PASSED"
    end
end

VIRTDISK = "/mnt/VIRTDISK"

desc 'Start an instance of Z80-OS in the Zemu interactive debugger.'
task 'debug' => ['build:kernel', 'build:command'] + BUILTINS.map{ |b| "build:builtin:#{b[1]}" } do
    # Abort if the destination directory does not exist.
    unless File.exist?("disk_copy.bin") && Dir.exist?(VIRTDISK)
        abort("Virtual disk not mounted at '#{VIRTDISK}'")
    end

    FileUtils.rm(Dir.glob(File.join(VIRTDISK, "*.exe")))
    while !Dir.glob(File.join(VIRTDISK, "*.exe")).empty?
        sleep(1)
    end

    FileUtils.cp("command.exe", VIRTDISK)

    BUILTINS.each do |_, name|
        FileUtils.cp("#{name}.exe", VIRTDISK)
    end

    zemu_start()
end

desc 'Debug a test using Zemu.'
task 'debug_test', [:test] => ['build:kernel_debug'] do |t, args|
    test_dir = File.join("kernel", "integration_test", args[:test])
    test_app = File.join(test_dir, "#{args[:test]}.bin")
    test_disk = File.join(test_dir, "disk.bin")
    zemu_start("kernel_debug.bin", test_app, test_disk)
end

namespace 'build' do
    desc "Build boot sector image"
    task 'boot' do
        cmd = "zcc "
        cmd += "+#{CONFIG} -compiler-sccz80 "
        cmd += "-O2 -SO2 "
        cmd += "-L#{LIB} -I#{LIB_INCLUDE} "
        cmd += "-Ca\"-I#{LIB_INCLUDE}\" "
        cmd += "--no-crt "
        cmd += "-o bootsector_temp.bin "
        cmd += "boot/bootsector.asm"

        success = system(cmd)

        if success
            # Pad to 512 bytes.
            bin = File.read("bootsector_temp.bin", mode: "rb")
            abort("Boot sector too large: size is #{bin.size}") if bin.size > 512
            File.open("bootsector.bin", "wb") do |f|
                bin.bytes.each do |b|
                    f.putc b
                end

                (512-bin.size).times do
                    f.putc 0
                end
            end

            system("z88dk-dis -o 0x8000 bootsector.bin > bootsector.diss")
        end
    end

    desc "Build 2nd stage loader image"
    task 'loader' => ['lib:stdlib'] do
        cmd = "zcc "
        cmd += "+#{CONFIG} -compiler-sccz80 "
        cmd += "-DZ88DK "
        cmd += "-O2 -SO2 "
        cmd += "-L#{LIB} -I#{LIB_INCLUDE} "
        cmd += "-Ca\"-I#{LIB_INCLUDE}\" "
        cmd += "-Cl\"-r0x8200\" "
        cmd += "-crt0 loader/reset.asm "
        cmd += "-m "
        cmd += "-o loader_temp.bin "
        cmd += "loader/*.c "
        cmd += "loader/*.asm "

        success = system(cmd)

        if success
            num_sectors = 9

            bin = File.read("loader_temp.bin", mode: "rb")
            abort("Loader image too large: size is #{bin.size} (#{512 * num_sectors} expected)") if bin.size > (512*num_sectors)

            File.open("loader.bin", "wb") do |f|
                bin.bytes.each do |b|
                    f.putc b
                end

                ((512*num_sectors)-bin.size).times do
                    f.putc 0
                end
            end

            system("z88dk-dis -o 0x8200 loader.bin > loader.diss")

            File.open("loader.bin", "rb") do |f|
                num_sectors.times do |i|
                    File.open("loader_#{i+1}.bin", "wb") do |f_out|
                        512.times do
                            f_out.putc f.getc
                        end
                    end
                end
            end
        end
    end

    kernel_dependencies = FileList.new("kernel/*.c", "kernel/*.asm") do |fl|
        fl.exclude(/reset/)
    end.ext(".rel")
    kernel_dependencies_debug = kernel_dependencies.ext(".debug.rel")

    desc "Build kernel image"
    task 'kernel' => ['lib:stdlib', 'kernel/reset.rel'] + kernel_dependencies do
        success = compile_sdcc("kernel.bin", kernel_dependencies, 'kernel/reset.rel', 0x0060, 0x6000)

        if success
            system("z88dk-dis -o 0x0000 kernel.bin > kernel.diss")

            # Make sure that the kernel image is no more than 32Kb in size.
            limit = 32 * 1024
            bin = File.read("kernel.bin", mode: "rb")
            abort("Kernel image is too large: size is #{bin.size} (#{limit} expected)") if bin.size > limit
        end
    end

    desc "Build kernel image in debug mode for tests"
    task 'kernel_debug' => ['lib:stdlib', 'kernel/reset.rel'] + kernel_dependencies_debug do
        success = compile_sdcc("kernel_debug.bin", kernel_dependencies_debug, 'kernel/reset.rel', 0x0060, 0x6000)

        if success
            system("z88dk-dis -o 0x0000 kernel_debug.bin > kernel_debug.diss")

            # Make sure that the kernel image is no more than 32Kb in size.
            limit = 32 * 1024
            bin = File.read("kernel_debug.bin", mode: "rb")
            abort("Kernel image is too large: size is #{bin.size} (#{limit} expected)") if bin.size > limit
        end
    end

    command_dependencies = FileList.new("command/*.asm", "command/*.c") do |fl|
        fl.exclude(/reset/)
    end.ext("rel")

    desc "Build command processor image"
    task 'command' => ['lib:stdlib', PROCESS_CRT0_OBJ] + command_dependencies do
        success = compile_sdcc("command.bin", command_dependencies, PROCESS_CRT0_OBJ, 0x8000, 0xd000)

        if success
            system("z88dk-dis -o 0x8000 command.bin > command.diss")

            # Make sure that the command processor image is no more than 32Kb in size.
            limit = 32 * 1024
            bin = File.read("command.bin", mode: "rb")
            abort("Command processor image is too large: size is #{bin.size} (#{limit} expected)") if bin.size > limit

            # Now make an executable from the binary.
            File.open("command.exe", "wb") do |f|
                f.write([0x0a, 0x80].pack("CC"))
                File.open("command.bin", "rb") do |f2|
                    f2.each_byte do |b|
                        f.write([b].pack("C"))
                    end
                end
            end
        end
    end

    BUILTINS.each do |builtin, builtin_name|
        # Open the YAML file for the builtin program.
        config = YAML.load_file(File.join(builtin, "config.yaml"))

        # Base address for the program.
        code_seg = config['code']
        data_seg = config['data']
        code_seg_hdr = code_seg >> 8

        builtin_dependencies = FileList.new("#{builtin}/*.asm", "#{builtin}/*.c").ext("rel")

        desc "Build executable for builtin program '#{builtin_name}'"
        task "builtin:#{builtin_name}" => ['lib:stdlib', PROCESS_CRT0_OBJ] + builtin_dependencies do
            success = compile_sdcc("#{builtin_name}.bin", builtin_dependencies, PROCESS_CRT0_OBJ, code_seg, data_seg)

            if success
                system("z88dk-dis -o 0x#{code_seg.to_s(16)} #{builtin_name}.bin > #{builtin_name}.diss")

                # Now make an executable from the binary.
                File.open("#{builtin_name}.exe", "wb") do |f|
                    f.write([0x0a, code_seg_hdr].pack("CC"))
                    File.open("#{builtin_name}.bin", "rb") do |f2|
                        f2.each_byte do |b|
                            f.write([b].pack("C"))
                        end
                    end
                end
            end
        end
    end

    desc "Build all"
    task "all" => ["loader", "boot", "kernel", "command"] + Dir.glob("builtin/*").map { |builtin| "builtin:#{File.basename(builtin)}" }
end

namespace 'install' do
    desc "Install kernel image onto CF-card"
    task 'kernel', [:path] => ['build:kernel'] do |t, args|
        abort("No path specified") if args[:path].nil?
        FileUtils.cp('kernel.bin', File.join(args[:path], 'KERNEL.BIN'))
    end

    desc "Install command processor image onto CF-card"
    task 'command', [:path] => ['build:command'] do |t, args|
        abort("No path specified") if args[:path].nil?
        FileUtils.cp('command.exe', File.join(args[:path], 'COMMAND.EXE'))
    end

    Dir.glob("builtin/*").each do |builtin|
        builtin_name = File.basename(builtin)

        desc "Install executable for builtin program '#{builtin_name}'"
        task "builtin:#{builtin_name}", [:path] => ["build:builtin:#{builtin_name}"] do |t, args|
            abort("No path specified") if args[:path].nil?
            FileUtils.cp("#{builtin_name}.exe", File.join(args[:path], "#{builtin_name.upcase}.EXE"))
        end
    end
end

namespace 'test' do
    Dir.glob("builtin/*").each do |builtin|
        builtin_name = File.basename(builtin)

        desc "Run unit tests for builtin program '#{builtin_name}'"
        task "builtin:#{builtin_name}" do
            unit_test(builtin)
        end
    end

    namespace 'kernel' do
        desc "Run kernel integration tests"
        Rake::TestTask.new 'integration' => ["build:kernel_debug", "kernel/integration_test/reset.rel", "lib:process_crt0"] do |t|
            t.test_files = FileList['kernel/integration_test/test_*.rb'].exclude('kernel/integration_test/test_helper.rb')
        end

        desc "Run kernel unit tests"
        task 'unit' do
            unit_test("kernel")
        end

        desc "Run all kernel tests"
        task 'all' => ['test:kernel:unit', 'test:kernel:integration']
    end
end

class Function
    attr_reader :addr, :size, :label, :mod

    def initialize(addr, size, label, mod)
        @addr = addr
        @size = size
        @label = label
        @mod = mod

        @instrs = []
        @coverage = {}
    end

    def span
        (@addr...@addr+@size)
    end

    def num_instrs
        @instrs.size
    end

    def num_covered
        @instrs.select { |i| !@coverage[i].nil? }.size
    end

    def add_instr(address)
        if !span.include? address
            raise "#{address.to_s(16)} is not in range of function span #{span}"
        end

        @instrs << address unless @instrs.include? address
    end

    def add_coverage(address)
        if !span.include? address
            raise "#{address.to_s(16)} is not in range of function span #{span}"
        end

        # Increment the number of times we've seen this address.
        if @coverage[address].nil?
            @coverage[address] = 1
        else
            @coverage[address] += 1
        end
    end

    def coverage
        (num_covered.to_f / num_instrs.to_f) * 100
    end
end

desc 'Generate kernel coverage report'
task 'coverage' => ['build:kernel_debug', 'test:kernel:integration'] do
    # Get addresses in disassembly
    disassembly_addrs = []
    File.open("kernel_debug.diss") do |f|
        f.each_line do |l|
            if /;\[([0-9A-Fa-f]{4})\]/ =~ l
                addr = $1.to_i(16)
                disassembly_addrs << addr
            end
        end
    end

    # Parse map file and get a list of functions with their start address and size,
    # plus module if present.
    functions = []
    File.open("kernel_debug.map") do |f|
        prev_addr = 0
        prev_name = "reset"
        prev_module = nil

        f.each_line do |l|
            if /^[0-9A-Fa-f]{8}\s+\S+/ =~ l.strip
                l = l.split
                addr = l[0].to_i(16)

                # Skip over labels
                next if l[1].start_with?("._") || l[1].start_with?("s_") || l[1].start_with?("l_")

                # Skip over data
                next if addr >= 0x6000

                # Add an entry for the previous function now that we know its size.
                functions << Function.new(prev_addr, addr - prev_addr, prev_name, prev_module)
                
                prev_addr = addr
                prev_name = l[1]
                prev_module = l[2]
            end
        end

        # Add the final function.
        functions << Function.new(prev_addr, (disassembly_addrs[-1]+1) - prev_addr, prev_name, prev_module)
    end

    # Figure out each instruction in each function.
    functions_by_addr = {}
    disassembly_addrs.each do |addr|
        # Find the function this address exists in.
        funcs_for_addr = functions.select { |f| f.span.include? addr }
        if funcs_for_addr.size != 1
            raise "No function for address: #{addr.to_s(16)}"
        end

        funcs_for_addr[0].add_instr(addr)
        functions_by_addr[addr] = funcs_for_addr[0]
    end

    # Parse generated coverage files to figure out which addresses we've hit.
    # We keep track of the number of times we see each address, so we could create
    # a heatmap if we want to.
    coverage_addresses = {}
    Dir.glob('**/*.coverage').each do |cov|
        File.open(cov) do |f|
            f.each_line do |l|
                if /\$([0-9A-Fa-f]{4})/ =~ l
                    addr = $1.to_i(16)

                    # Find the function this address exists in.
                    func = functions_by_addr[addr]
                    unless func.nil?
                        func.add_coverage(addr)
                    end
                end
            end
        end
    end

    # Calculate the number of addresses in the program we see in tests.
    num_instrs = functions.map { |f| f.num_instrs }.reduce(:+)
    num_covered = functions.map { |f| f.num_covered }.reduce(:+)

    # Print a coverage percentage.
    percentage = (num_covered.to_f / num_instrs.to_f) * 100
    puts "Coverage: #{percentage.round(2)}%\n"

    # Get modules that aren't nil.
    modules = functions.map { |f| f.mod }.uniq

    modules.each do |mod|
        next if mod.nil?

        mod_functions = functions.select { |f| f.mod == mod }.sort_by { |f| f.coverage }

        # Print coverage report for the module.
        puts "Module '#{mod}':"
        mod_functions.each do |f|
            puts "    %-35s %2.2f%%" % [f.label, f.coverage]
        end
        puts ""
    end

    nil_functions = functions.select { |f| f.mod.nil? }.sort_by { |f| f.coverage }
    puts "Other functions:"
    nil_functions.each do |f|
        puts "    %-35s %2.2f%%" % [f.label, f.coverage]
    end
    puts ""
end

namespace 'benchmark' do
    desc "Run kernel benchmarks"
    task 'kernel' => "build:kernel_debug" do
        Dir.glob("benchmark/kernel/*.rb").each do |b|
            require_relative "#{b}"
            if defined? benchmarks
                benchmarks()
                undef :benchmarks
            end
        end
    end
end

task "test" => ["test:kernel:all"] + BUILTINS.map { |_, name| "test:builtin:#{name}" }

task "install", [:path] => ["install:kernel", "install:command"] + Dir.glob("builtin/*").map { |builtin| "install:builtin:#{File.basename(builtin)}" }
