import subprocess

NORMAL_COLOR = '\033[0m'

REG_STATUS_TO_COLOR = {
    'insert': '\033[32m',
    'change': '\033[33m',
    'remove': '\033[31m',
    'none': '\033[0m',
    'jump': '\033[36m',
}


class AsciiPainter:
    def __init__(self, registers=None, max_stack_size=8, stack_draw_mode='up'):
        self.registers = registers if registers else REGISTERS
        self.states = []
        self.max_stack_size = max_stack_size
        self.stack_draw_mode = stack_draw_mode

    def saveState(self, emu, line_nr: int | None = None) -> None:
        # Init state
        state = []
        self.states.append(state)
        # Write opcode to state
        if line_nr is not None:
            signature = emu.getLineSignature(line_nr)
            for token in signature:
                state.append({'state': 'none', 'val': token})
        #
        for i in range(3 - len(state)):
            state.append({'state': 'none', 'val': ''})
        for i in range(len(state) - 3):
            state.pop()
        state.append({'state': 'none', 'val': ''})
        # Write registers to state
        for reg_name in self.registers:
            if reg_name == '':
                state.append({'state': 'none', 'val': ''})
                continue
            reg_state = emu.regState(reg_name)
            state.append({'state': reg_state, 'val': emu.getVal(reg_name)})
        # Write stack to state
        stack_base_pointer = emu.getVal('%rbp')
        base_sp = emu.getVal('%rsp')
        assert isinstance(emu.getVal('%rsp'), int)
        state.append({'state': 'none', 'val': ''})
        if emu.regState('m' + str(base_sp - 1)) == 'remove':
            state.append(
                {
                    'state': emu.regState('m' + str(base_sp - 1)),
                    'val': emu.stack[base_sp - 1],
                },
            )
        else:
            state.append({'state': 'none', 'val': ''})
        for index in range(0, self.max_stack_size):
            stack_i = base_sp + index
            if stack_i < emu.max_stack_size:
                state.append(
                    {
                        'state': emu.regState('m' + str(stack_i)),
                        'val': emu.stack[stack_i],
                    },
                )
                if stack_base_pointer == index:
                    stack_base_pointer = emu.stack[base_pointer]
                # TODO: Draw stack frame seperators
            else:
                state.append({'state': 'none', 'val': ''})

    def getWidthOfColumns(
        self,
        number_states: int,
        number_states_per_block: int,
        number_state_vars: int,
    ) -> list[int]:
        widths = [0 for i in range(0, number_states_per_block)]
        for base_state in range(0, number_states, number_states_per_block):
            for var_i in range(0, number_state_vars):
                for state_i in range(
                    0,
                    min(number_states_per_block, number_states - base_state),
                ):
                    width = len(str(self.states[base_state + state_i][var_i]['val']))
                    widths[state_i] = max(widths[state_i], width)
        return widths

    def getWidthOfNameColumn(self, number_state_vars: int) -> int:
        widest = 0
        for var_i in range(0, number_state_vars):
            widest = max(widest, len(self.nameOfVar(var_i - 4)))
        return widest

    def to_string(self, emu) -> str:
        number_states = len(self.states)
        number_state_vars = len(self.states[0])

        term_width = int(subprocess.check_output(['tput', 'cols']))
        number_states_per_block: int = term_width // (10 + 2) - 1
        separator = '-' * term_width
        output: list[str] = []

        name_fmt = '{}{:' + str(self.getWidthOfNameColumn(number_state_vars)) + '} '
        column_fmt = [
            '{}{:>' + str(width) + '}  '
            for width in self.getWidthOfColumns(
                number_states,
                number_states_per_block,
                number_state_vars,
            )
        ]

        for base_state in range(0, number_states, number_states_per_block):
            output.append(separator)
            for var_i in range(0, number_state_vars):
                output.append(name_fmt.format(NORMAL_COLOR, self.nameOfVar(var_i - 4)))
                for state_i in range(
                    0,
                    min(number_states_per_block, number_states - base_state),
                ):
                    var = self.states[base_state + state_i][var_i]
                    output.append(
                        column_fmt[state_i].format(
                            REG_STATUS_TO_COLOR[var['state']],
                            var['val'],
                            NORMAL_COLOR,
                        ),
                    )
                output.append('\n')
        output.append(separator)

        return ''.join(output)

    def nameOfVar(self, i):
        if i < 0:
            return ''
        if i < len(self.registers):
            return self.registers[i]
        elif i != len(self.registers):
            state_i = i - len(self.registers) - 2
            return 'TOS{}{}'.format('+' if state_i < 0 else '-', abs(state_i))
        else:
            return ''

    def __str__(self):
        return 'AsciiPainter[]'