from amaranth import *
from amaranth.sim import Simulator, Settle, Delay
from enum import Enum, unique

from hdl.utils import *
from hdl.lib.in_out_buff import InOutBuff   # used for timing analysis
from hdl.core.alu import AluOpCodes
from hdl.core.jump_ctl import JumpOpCodes
from hdl.config import MACHINE_CODE_CSV

import csv

MAX_INSTR = 2**8    # do not change this, this is not configurable
CL_NAMES = ['mode', 'valid_op', 'wb_en', 'alu_op', 'is_imm', 'jump_en', 'mem_wr_en', 'mem_rd_en', 'mem_rw_word']

class Control_LUT(Elaboratable):
    def __init__(self, csv_file, cl_names, **kargs):    # cl_names is a list of control line in order of msb to lsb
        self.csv_file = csv_file
        self.cl_names = cl_names
        self.sim = kargs.get("sim", False)

    def elaborate(self, platform=None):
        pass
        #TODO: implment this to create signals with the same names as the control lines specified

    def create_lut(self):
        csv_raw = self._read_csv(self.csv_file)
        control_list = self._parse_csv(csv_raw)
        mem_width, slice_mapping = self._get_slice_mapping(control_list)

        print('debug')
        print(CL_NAMES)
        print(mem_width, slice_mapping)
        print(control_list[10])
        print(self._create_entry(control_list[10]['data'], mem_width, slice_mapping))

        # self.mem_init = self._create_mem_init(self.csv_parsed)
        # self.mem     = Memory(width=len(self.cl_names) + 1, depth=MAX_INSTR, init=self.mem_init, simulate=self.sim)       # one extra bit for valid opcode

    # read csv file and return a list of dictionaries
    def _read_csv(self, csv_file):
        with open(csv_file, 'r') as f:
            reader = csv.DictReader(f)
            return list(reader)
        
    def _get_slice_mapping(self, control_list):

        # helper function to count the number of bits in a number
        def count_bits(num:int):
            return len(bin(num)[2:])     # remove the '0b' prefix
        
        # find the max bit length of the control lines per entry
        max_width_list = [count_bits(cl) for cl in control_list[0]['data']]
        for entry in control_list:
            for i, cl in enumerate(entry['data']):
                max_width_list[i] = max(max_width_list[i], count_bits(cl))

        # calculate how many bits for each entry
        mem_width = sum(max_width_list)
        slice_mapping = []
        accum = 0

        # generate the slice mapping
        for i, cl in enumerate(max_width_list):
            slice_mapping.append(slice(accum, accum + max_width_list[i]))
            accum += max_width_list[i]
        
        return mem_width, slice_mapping
    

    def _create_entry(self, control, mem_width, slice_mapping):

        # check that the slices map to the correct number of bits
        assert(slice_mapping[-1].stop == mem_width), 'slice mapping is not continuous, this is an internal error'

        def bin_stripped(num):
            return bin(num)[2:]
        
        def pad_zero(num_str, width):
            return '0' * (width - len(num_str)) + num_str
        
        # create the entry with all zeros
        entry_str_list = ['0'] * mem_width

        # fill in the control lines, lsb left, msb right
        for (cl, s) in zip(control, slice_mapping):
            entry_str_list[s.start:s.stop] = list(reversed(list(pad_zero(bin_stripped(cl), s.stop - s.start))))     # signals read from left to right, so reverse the string

        # conversion to int requires lsb right, msb left
        entry_str = ''.join(reversed(entry_str_list))

        # convert to int
        entry = int('0b' + entry_str, 2)
        return entry

    def _parse_csv(self, csv_raw):
        csv_parsed = []

        for row in csv_raw:
            tmp = {}
            tmp['opcode'] = int(row['opcode'])
            tmp['data'] = []

            for cl in self.cl_names:
                try:
                    tmp['data'].append(int(row[cl]))
                except KeyError:
                    print(f'control line "{cl}" not found in csv file')
                    raise

            csv_parsed.append(tmp)
        return csv_parsed
    
    def _alloc_mem(self, csv_parsed):
        pass
            
    def _create_init(self, csv_parsed):
        mem_init = []

        for i in range(MAX_INSTR):
            mem_init.append(0)     # init to 0, this sets all op codes to invalid

        for row in csv_parsed:
            assert(0 <= row['opcode'] < MAX_INSTR), "opcode out of range"
            assert((mem_init[row['opcode']] & 0x1) != 1), 'duplicate opcode, check the .csv file'

            for cl in row['data']:
                mem_init[row['opcode']] = cl

        return mem_init


class Control(Elaboratable):
    def __init__(self, **kargs):
        self.instr_op = Signal(4, reset_less=True)

        # out opcodes
        self.alu_op = Signal(e2s(AluOpCodes))
        self.jump_ctl_op = Signal(e2s(JumpOpCodes))

        # syncronous control signals
        self.valid_op = Signal()
        self.wb_en = Signal()
        self.alu_op = Signal(e2s(AluOpCodes))
        self.is_imm = Signal()
        self.jump_en = Signal()
        self.mem_wr_en = Signal()
        self.mem_rd_en = Signal()
        self.mem_rw_word = Signal()


        self.stall = Signal()
        self.int_sig = Signal()
        self.iret = Signal()
        self.call = Signal()
        self.jump = Signal()

        # register control in
        self.int_en = Signal(1)
        self.user_mode = Signal(1)

        ports_in = [self.instr_op, self.alu_op, self.jump_ctl_op, self.wr_en, self.stall, self.int_sig, self.iret, self.call, self.jump]
        ports_out = []
        self.ports = {'in': ports_in, 'out': ports_out}

        self.sim = kargs.get("sim", False)

    def elaborate(self, platform=None):
        m = Module()

        # dummy sync for simulation only needed if there is no other sequential logic
        if self.sim == True:
            dummy = Signal()
            m.d.sync += dummy.eq(~dummy)

        ...
        
        return m


# test addition
# def test_hdl():
#     dut = Control(sim=True)
#     def proc():
#         yield from step         #step clock 
#         yield Settle()          #needed if for combinatorial logic
#         yield dut.something     #read value
#     sim(dut, proc)

def test_control_mem_init():
    lut = Control_LUT(MACHINE_CODE_CSV, CL_NAMES, sim=True)
    lut.create_lut()
       
if __name__ == '__main__':
    test_control_mem_init()
    # hdl = InOutBuff(Control())
    # cmd(hdl)