#!/usr/bin/env python3

from shell_helpers import LF
import common
import lkmc.import_path

class Main(common.LkmcCliFunction):
    def __init__(self):
        super().__init__(
            description='''Trace the PIC addresses executed on a Linux kernel boot.

More information at: https://cirosantilli.com/linux-kernel-module-cheat#tracing
'''
        )

    def timed_main(self):
        args = self.get_common_args()
        run = lkmc.import_path.import_path_main('run')
        if self.env['emulator'] == 'gem5':
            args['trace'] = 'Exec,-ExecSymbol,-ExecMicro'
            run(**args)
        elif self.env['emulator'] == 'qemu':
            run_args = args.copy()
            run_args['trace'] = 'exec_tb'
            run_args['quit_after_boot'] = True
            run(**run_args)
            qemu_trace2txt = lkmc.import_path.import_path_main('qemu-trace2txt')
            qemu_trace2txt(**args)
            # Instruction count.
            # We could put this on a separate script, but it just adds more arch boilerplate to a new script.
            # So let's just leave it here for now since it did not add a significant processing time.
            kernel_entry_addr = hex(self.get_elf_entry(self.env['vmlinux']))
            nlines = 0
            nlines_firmware = 0
            with open(self.env['qemu_trace_txt_file'], 'r') as trace_file:
                in_firmware = True
                for line in trace_file:
                    line = line.rstrip()
                    nlines += 1
                    pc = line.split('=')[-1]
                    if pc == kernel_entry_addr:
                        in_firmware = False
                    if in_firmware:
                        nlines_firmware += 1
            print('''\
instructions {}
entry_address {}
instructions_firmware {}
'''.format(
                    nlines,
                    kernel_entry_addr,
                    nlines_firmware
                ),
                end=''
            )

if __name__ == '__main__':
    Main().cli()
