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[]'