import sys
sys.dont_write_bytecode = True
import math

from amaranth import *
from amaranth.lib import wiring
from amaranth.lib.wiring import In, Out
from amaranth.lib.enum import Enum
from amaranth.back import verilog

##############################
######### Oscillator #########
##############################
class Sn76489_Oscillator(wiring.Component):
    def __init__(self):
        print("### Sn76489_Oscillator constructor", file=sys.stderr)
        super().__init__({
            "I_clk_gate": In(1),
            "I_freq": In(10),
            "O_voice": Out(1)
        })
    
    def elaborate(self, platform):
        m = Module()

        counter = Signal(10)
        out = Signal(1)

        with m.If(self.I_clk_gate):
            m.d.sync += counter.eq(counter - 1)
            with m.If(counter == 0):
                m.d.sync += [
                    out.eq(~out),
                    counter.eq(self.I_freq)
                ]

        # for frequency values of 0 and 1 output a constant 1
        voice = out | self.I_freq[1:10] == 0
        m.d.comb += self.O_voice.eq(voice)

        return m


###################################
######### Noise generator #########
###################################
class Sn76489_Noise(wiring.Component):
    def __init__(self):
        print("### Sn76489_Noise constructor", file=sys.stderr)
        super().__init__({
            "I_clk_gate": In(1),
            "I_ctrl": In(3),
            "I_freq": In(10),
            "I_reset_voice": In(1),
            "O_voice": Out(1),
            "O_reset_ack": Out(1)
        })

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

        reset = Signal(1)
        counter = Signal(10)
        shiftreg = Signal(16, init=0x8000)
        flipbit = Signal(1)
        rate = self.I_ctrl[0:2]
        white_noise = self.I_ctrl[2]

        with m.If(self.I_clk_gate):
            m.d.sync += counter.eq(counter - 1)
            with m.If(counter == 0):
                m.d.sync += flipbit.eq(~flipbit)
                with m.Switch(rate):
                    with m.Case(0b00):
                        m.d.sync += counter.eq(0x10)
                    with m.Case(0b01):
                        m.d.sync += counter.eq(0x20)
                    with m.Case(0b10):
                        m.d.sync += counter.eq(0x40)
                    with m.Case(0b11):
                        m.d.sync += counter.eq(self.I_freq)
                
                with m.If(flipbit == 0):
                    m.d.sync += [
                        shiftreg[0:15].eq(shiftreg[1:16]),
                        shiftreg[15].eq(Mux(white_noise, shiftreg[3] ^ shiftreg[0], shiftreg[0]))
                    ]

        with m.If(self.I_reset_voice != reset):
            m.d.sync += [
                shiftreg.eq(0x8000),
                reset.eq(self.I_reset_voice)
            ]


        m.d.comb += [
            self.O_reset_ack.eq(reset),
            self.O_voice.eq(shiftreg[0])
        ]

        return m

#########################
######### Mixer #########
#########################
class Sn76489_Mixer(wiring.Component):
    def __init__(self):
        print("### Sn76489_Mixer constructor", file=sys.stderr)
        super().__init__({
            "I_voices": In(4),
            "I_attenuation": In(4 * 4),
            "O_pcm": Out(8)
        })
    
    def voice_value(self, m, voice, att):
        value = Signal(6)
        values = [63, 59, 55, 50, 46, 42, 38, 34, 29, 25, 21, 17, 13, 8, 4]
        with m.Switch(Cat(att, voice)):
            for i in range(len(values)):
                with m.Case(0b10000 + i):
                    m.d.comb += value.eq(values[i])
            with m.Default():
                m.d.comb += value.eq(0)
        
        return value
    
    def elaborate(self, platform):
        m = Module()

        values = []
        for i in range(4):
            voice = self.I_voices[i]
            att = self.I_attenuation[i*4:(i*4)+4]
            values.append(self.voice_value(m, voice, att))
        
        pcm1 = values[0] + values[1]    # 6 bits + 6 bits = 7 bits
        pcm2 = values[2] + values[3]    # 6 bits + 6 bits = 7 bits
        pcm = pcm1 + pcm2               # 7 bits + 7 bits = 8 bits

        m.d.comb += self.O_pcm.eq(pcm)

        return m

#######################
######### Top #########
#######################
class Sn76489(wiring.Component):
    def __init__(self, clk_target=3579545, clk_infreq=25 * 1000 * 1000):
        print("### Sn76489 constructor", file=sys.stderr)
        super().__init__({
            "I_command": In(8),
            "I_write": In(1),
            "O_pcm": Out(8),
            "O_pcm_new": Out(1)
        })
        self.clk_divider = round((clk_infreq / clk_target) * 16)
        print(f"# SN76489 clk_divider: {self.clk_divider}")
    
    def elaborate(self, platform):
        m = Module()

        # Clock divider stuff
        clk_gate = Signal(1)
        clk_counter = Signal(math.ceil(math.log(self.clk_divider, 2)))
        m.d.sync += [
            clk_gate.eq(0),
            clk_counter.eq(clk_counter - 1)
        ]
        with m.If(clk_counter == 0):
            m.d.sync += [
                clk_gate.eq(1),
                clk_counter.eq(self.clk_divider)
            ]

        # voice control registers
        registers = [
            Signal(10), # register 000 - voice 1 freq
            Signal(4),  # register 001 - voice 1 attenuation
            Signal(10), # register 010 - voice 2 freq
            Signal(4),  # register 011 - voice 2 attenuation
            Signal(10), # register 100 - voice 3 freq
            Signal(4),  # register 101 - voice 3 attenuation
            Signal(3),  # register 110 - noise control
            Signal(4)   # register 111 - noise attenuation
        ]

        # Instantiate three oscillators
        oscillators = []
        voice_outputs = []
        for i in range(3):
            osc = Sn76489_Oscillator()
            m.submodules += osc
            oscillators.append(osc)
            osc_output = Signal(1)
            voice_outputs.append(osc_output)
            m.d.comb += [
                osc.I_clk_gate.eq(clk_gate),
                osc.I_freq.eq(registers[2*i]), # connect freq registers
                osc_output.eq(osc.O_voice)
            ]

        # Instantiate noise generator
        noise = Sn76489_Noise()
        m.submodules += noise
        noise_reset = Signal(1)
        noise_reset_ack = Signal(1)
        m.d.comb += [
            noise.I_clk_gate.eq(clk_gate),
            noise.I_ctrl.eq(registers[6]),
            noise.I_freq.eq(registers[4]),
            noise.I_reset_voice.eq(noise_reset),
            noise_reset_ack.eq(noise.O_reset_ack)
        ]
        voice_outputs.append(noise.O_voice)

        # Instantiate the voice mixer
        mixer = Sn76489_Mixer()
        m.submodules += mixer
        pcm = Signal(8)
        m.d.comb += [
            mixer.I_voices.eq(Cat(voice_outputs[0], voice_outputs[1], voice_outputs[2], voice_outputs[3])),
            mixer.I_attenuation.eq(Cat(registers[1], registers[3], registers[5], registers[7])),
            pcm.eq(mixer.O_pcm)
        ]

        # Update logic for control registers
        update = Signal(1)
        m.d.sync += update.eq(0)
        update_data = Signal(7)
        register = Signal(3)
        with m.If(self.I_write):
            m.d.sync += [
                update.eq(1),
                update_data.eq(Cat(self.I_command[0:6], self.I_command[7]))
            ]
            with m.If(self.I_command[7]):
                m.d.sync += register.eq(self.I_command[4:7])

        with m.If(update):
            with m.Switch(Cat(update_data[6], register)):
                with m.Case(0b0000):
                    m.d.sync += registers[0][4:10].eq(update_data[0:6])
                with m.Case(0b0001):
                    m.d.sync += registers[0][0:4].eq(update_data[0:4])
                with m.Case(0b0010, 0b0011):
                    m.d.sync += registers[1].eq(update_data[0:4])
                with m.Case(0b0100):
                    m.d.sync += registers[2][4:10].eq(update_data[0:6])
                with m.Case(0b0101):
                    m.d.sync += registers[2][0:4].eq(update_data[0:4])
                with m.Case(0b0110, 0b0111):
                    m.d.sync += registers[3].eq(update_data[0:4])
                with m.Case(0b1000):
                    m.d.sync += registers[4][4:10].eq(update_data[0:6])
                with m.Case(0b1001):
                    m.d.sync += registers[4][0:4].eq(update_data[0:4])
                with m.Case(0b1010, 0b1011):
                    m.d.sync += registers[5].eq(update_data[0:4])
                with m.Case(0b1100, 0b1101):
                    m.d.sync += [
                        registers[6].eq(update_data[0:3]),
                        noise_reset.eq(~noise_reset_ack)
                    ]
                with m.Case(0b1110, 0b1111):
                    m.d.sync += registers[7].eq(update_data[0:4])


        m.d.sync += [
            self.O_pcm.eq(pcm),
            self.O_pcm_new.eq(clk_gate)
        ]

        return m


if __name__ == "__main__":
    psg = Sn76489()
    with open("sn76489.v", "w") as f:
        f.write(verilog.convert(psg, name="sn76489"))