#!/usr/bin/env python3
#
# Copyright (c) 2021 Hans Baier <hansfbaier@gmail.com>
# SPDX-License-Identifier: CERN-OHL-W-2.0

from amaranth import *

from usb_protocol.types                       import USBRequestType, USBRequestRecipient, USBTransferType, USBSynchronizationType, USBUsageType, USBDirection, USBStandardRequests
from usb_protocol.types.descriptors.uac2      import AudioClassSpecificRequestCodes
from usb_protocol.emitters                    import DeviceDescriptorCollection
from usb_protocol.emitters.descriptors        import uac2, standard
from usb_protocol.emitters.descriptors        import midi1

class USBDescriptors():
    MAX_PACKET_SIZE_MIDI = 64
    CLOCK_ID             = 1

    def __init__(self, *, ila_max_packet_size: int, use_ila=False) -> None:

        # ILA
        self.USE_ILA             = use_ila
        self.ILA_MAX_PACKET_SIZE = ila_max_packet_size


    def create_usb1_descriptors(self, no_channels: int, max_packet_size: int):
        """ Creates the descriptors for the main USB interface """

        return self.create_descriptors("ADATface (USB1)", no_channels, max_packet_size, self.USE_ILA)


    def create_usb2_descriptors(self, no_channels: int, max_packet_size: int):
        """ Creates the descriptors for the secondary USB interface """
        return self.create_descriptors("ADATface (USB2)", no_channels, max_packet_size)


    def create_descriptors(self, product_id: str, no_channels: int, max_packet_size: int, create_ila=False):
        """ Creates the descriptors for the main USB interface """

        descriptors = DeviceDescriptorCollection()

        with descriptors.DeviceDescriptor() as d:
            d.bcdUSB             = 2.00
            d.bDeviceClass       = 0xEF
            d.bDeviceSubclass    = 0x02
            d.bDeviceProtocol    = 0x01
            d.idVendor           = 0x1209
            d.idProduct          = 0xADA1 if "USB1" in product_id else 0xADA2
            d.iManufacturer      = "OpenAudioGear"
            d.iProduct           = product_id
            d.iSerialNumber      = "0"
            d.bcdDevice          = 0.01

            d.bNumConfigurations = 1

        with descriptors.ConfigurationDescriptor() as configDescr:
            # Interface Association
            interfaceAssociationDescriptor                 = uac2.InterfaceAssociationDescriptorEmitter()
            interfaceAssociationDescriptor.bInterfaceCount = 3 # Audio Control + Inputs + Outputs
            configDescr.add_subordinate_descriptor(interfaceAssociationDescriptor)

            # Interface Descriptor (Control)
            interfaceDescriptor = uac2.StandardAudioControlInterfaceDescriptorEmitter()
            interfaceDescriptor.bInterfaceNumber = 0
            configDescr.add_subordinate_descriptor(interfaceDescriptor)

            # AudioControl Interface Descriptor
            audioControlInterface = self.create_audio_control_interface_descriptor(no_channels)
            configDescr.add_subordinate_descriptor(audioControlInterface)

            self.create_output_channels_descriptor(configDescr, no_channels, max_packet_size)

            self.create_input_channels_descriptor(configDescr, no_channels, max_packet_size)

            midi_interface, midi_streaming_interface = self.create_midi_interface_descriptor()
            configDescr.add_subordinate_descriptor(midi_interface)
            configDescr.add_subordinate_descriptor(midi_streaming_interface)

            if create_ila:
                with configDescr.InterfaceDescriptor() as i:
                    i.bInterfaceNumber = 4

                    with i.EndpointDescriptor() as e:
                        e.bEndpointAddress = USBDirection.IN.to_endpoint_address(4)
                        e.wMaxPacketSize   = self.ILA_MAX_PACKET_SIZE

        return descriptors


    def create_audio_control_interface_descriptor(self, number_of_channels):
        audioControlInterface = uac2.ClassSpecificAudioControlInterfaceDescriptorEmitter()

        # AudioControl Interface Descriptor (ClockSource)
        clockSource = uac2.ClockSourceDescriptorEmitter()
        clockSource.bClockID     = self.CLOCK_ID
        clockSource.bmAttributes = uac2.ClockAttributes.INTERNAL_FIXED_CLOCK
        clockSource.bmControls   = uac2.ClockFrequencyControl.HOST_READ_ONLY
        audioControlInterface.add_subordinate_descriptor(clockSource)

        # streaming input port from the host to the USB interface
        inputTerminal               = uac2.InputTerminalDescriptorEmitter()
        inputTerminal.bTerminalID   = 2
        inputTerminal.wTerminalType = uac2.USBTerminalTypes.USB_STREAMING
        # The number of channels needs to be 2 here in order to be recognized
        # default audio out device by Windows. We provide an alternate
        # setting with the full channel count, which also references
        # this terminal ID
        inputTerminal.bNrChannels   = 2
        inputTerminal.bCSourceID    = 1
        audioControlInterface.add_subordinate_descriptor(inputTerminal)

        # audio output port from the USB interface to the outside world
        outputTerminal               = uac2.OutputTerminalDescriptorEmitter()
        outputTerminal.bTerminalID   = 3
        outputTerminal.wTerminalType = uac2.OutputTerminalTypes.SPEAKER
        outputTerminal.bSourceID     = 2
        outputTerminal.bCSourceID    = 1
        audioControlInterface.add_subordinate_descriptor(outputTerminal)

        # audio input port from the outside world to the USB interface
        inputTerminal               = uac2.InputTerminalDescriptorEmitter()
        inputTerminal.bTerminalID   = 4
        inputTerminal.wTerminalType = uac2.InputTerminalTypes.MICROPHONE
        inputTerminal.bNrChannels   = number_of_channels
        inputTerminal.bCSourceID    = 1
        audioControlInterface.add_subordinate_descriptor(inputTerminal)

        # audio output port from the USB interface to the host
        outputTerminal               = uac2.OutputTerminalDescriptorEmitter()
        outputTerminal.bTerminalID   = 5
        outputTerminal.wTerminalType = uac2.USBTerminalTypes.USB_STREAMING
        outputTerminal.bSourceID     = 4
        outputTerminal.bCSourceID    = 1
        audioControlInterface.add_subordinate_descriptor(outputTerminal)

        return audioControlInterface


    def create_output_streaming_interface(self, c, *, no_channels: int, alt_setting_nr: int, max_packet_size):
        # Interface Descriptor (Streaming, OUT, active setting)
        activeAudioStreamingInterface                   = uac2.AudioStreamingInterfaceDescriptorEmitter()
        activeAudioStreamingInterface.bInterfaceNumber  = 1
        activeAudioStreamingInterface.bAlternateSetting = alt_setting_nr
        activeAudioStreamingInterface.bNumEndpoints     = 2
        c.add_subordinate_descriptor(activeAudioStreamingInterface)

        # AudioStreaming Interface Descriptor (General)
        audioStreamingInterface               = uac2.ClassSpecificAudioStreamingInterfaceDescriptorEmitter()
        audioStreamingInterface.bTerminalLink = 2
        audioStreamingInterface.bFormatType   = uac2.FormatTypes.FORMAT_TYPE_I
        audioStreamingInterface.bmFormats     = uac2.TypeIFormats.PCM
        audioStreamingInterface.bNrChannels   = no_channels
        c.add_subordinate_descriptor(audioStreamingInterface)

        # AudioStreaming Interface Descriptor (Type I)
        typeIStreamingInterface  = uac2.TypeIFormatTypeDescriptorEmitter()
        typeIStreamingInterface.bSubslotSize   = 4
        typeIStreamingInterface.bBitResolution = 24
        c.add_subordinate_descriptor(typeIStreamingInterface)

        # Endpoint Descriptor (Audio out)
        audioOutEndpoint = standard.EndpointDescriptorEmitter()
        audioOutEndpoint.bEndpointAddress     = USBDirection.OUT.to_endpoint_address(1) # EP 1 OUT
        audioOutEndpoint.bmAttributes         = USBTransferType.ISOCHRONOUS  | \
                                                (USBSynchronizationType.ASYNC << 2) | \
                                                (USBUsageType.DATA << 4)
        audioOutEndpoint.wMaxPacketSize = max_packet_size
        audioOutEndpoint.bInterval      = 1
        c.add_subordinate_descriptor(audioOutEndpoint)

        # AudioControl Endpoint Descriptor
        audioControlEndpoint = uac2.ClassSpecificAudioStreamingIsochronousAudioDataEndpointDescriptorEmitter()
        c.add_subordinate_descriptor(audioControlEndpoint)

        # Endpoint Descriptor (Feedback IN)
        feedbackInEndpoint = standard.EndpointDescriptorEmitter()
        feedbackInEndpoint.bEndpointAddress  = USBDirection.IN.to_endpoint_address(1) # EP 1 IN
        feedbackInEndpoint.bmAttributes      = USBTransferType.ISOCHRONOUS  | \
                                               (USBSynchronizationType.NONE << 2)  | \
                                               (USBUsageType.FEEDBACK << 4)
        feedbackInEndpoint.wMaxPacketSize    = 4
        feedbackInEndpoint.bInterval         = 4
        c.add_subordinate_descriptor(feedbackInEndpoint)


    def create_output_channels_descriptor(self, c, no_channels: int, max_packet_size: int):
        #
        # Interface Descriptor (Streaming, OUT, quiet setting)
        #
        quietAudioStreamingInterface = uac2.AudioStreamingInterfaceDescriptorEmitter()
        quietAudioStreamingInterface.bInterfaceNumber  = 1
        quietAudioStreamingInterface.bAlternateSetting = 0
        c.add_subordinate_descriptor(quietAudioStreamingInterface)

        # we need the default alternate setting to be stereo
        # out for windows to automatically recognize
        # and use this audio interface
        self.create_output_streaming_interface(c, no_channels=2, alt_setting_nr=1, max_packet_size=max_packet_size)
        if no_channels > 2:
            self.create_output_streaming_interface(c, no_channels=no_channels, alt_setting_nr=2, max_packet_size=max_packet_size)


    def create_input_streaming_interface(self, c, *, no_channels: int, alt_setting_nr: int, channel_config: int=0, max_packet_size: int):
        # Interface Descriptor (Streaming, IN, active setting)
        activeAudioStreamingInterface = uac2.AudioStreamingInterfaceDescriptorEmitter()
        activeAudioStreamingInterface.bInterfaceNumber  = 2
        activeAudioStreamingInterface.bAlternateSetting = alt_setting_nr
        activeAudioStreamingInterface.bNumEndpoints     = 1
        c.add_subordinate_descriptor(activeAudioStreamingInterface)

        # AudioStreaming Interface Descriptor (General)
        audioStreamingInterface                 = uac2.ClassSpecificAudioStreamingInterfaceDescriptorEmitter()
        audioStreamingInterface.bTerminalLink   = 5
        audioStreamingInterface.bFormatType     = uac2.FormatTypes.FORMAT_TYPE_I
        audioStreamingInterface.bmFormats       = uac2.TypeIFormats.PCM
        audioStreamingInterface.bNrChannels     = no_channels
        audioStreamingInterface.bmChannelConfig = channel_config
        c.add_subordinate_descriptor(audioStreamingInterface)

        # AudioStreaming Interface Descriptor (Type I)
        typeIStreamingInterface  = uac2.TypeIFormatTypeDescriptorEmitter()
        typeIStreamingInterface.bSubslotSize   = 4
        typeIStreamingInterface.bBitResolution = 24 # we use all 24 bits
        c.add_subordinate_descriptor(typeIStreamingInterface)

        # Endpoint Descriptor (Audio out)
        audioOutEndpoint = standard.EndpointDescriptorEmitter()
        audioOutEndpoint.bEndpointAddress     = USBDirection.IN.to_endpoint_address(2) # EP 2 IN
        audioOutEndpoint.bmAttributes         = USBTransferType.ISOCHRONOUS  | \
                                                (USBSynchronizationType.ASYNC << 2) | \
                                                (USBUsageType.DATA << 4)
        audioOutEndpoint.wMaxPacketSize = max_packet_size
        audioOutEndpoint.bInterval      = 1
        c.add_subordinate_descriptor(audioOutEndpoint)

        # AudioControl Endpoint Descriptor
        audioControlEndpoint = uac2.ClassSpecificAudioStreamingIsochronousAudioDataEndpointDescriptorEmitter()
        c.add_subordinate_descriptor(audioControlEndpoint)


    def create_input_channels_descriptor(self, c, no_channels: int, max_packet_size: int):
        #
        # Interface Descriptor (Streaming, IN, quiet setting)
        #
        quietAudioStreamingInterface = uac2.AudioStreamingInterfaceDescriptorEmitter()
        quietAudioStreamingInterface.bInterfaceNumber  = 2
        quietAudioStreamingInterface.bAlternateSetting = 0
        c.add_subordinate_descriptor(quietAudioStreamingInterface)

        # Windows wants a stereo pair as default setting, so let's have it
        self.create_input_streaming_interface(c, no_channels=2, alt_setting_nr=1, channel_config=0x3, max_packet_size=max_packet_size)
        if no_channels > 2:
            self.create_input_streaming_interface(c, no_channels=no_channels, alt_setting_nr=2, channel_config=0x0, max_packet_size=max_packet_size)


    def create_midi_interface_descriptor(self):
        midi_interface = midi1.StandardMidiStreamingInterfaceDescriptorEmitter()
        midi_interface.bInterfaceNumber = 3
        midi_interface.bNumEndpoints    = 2

        midi_streaming_interface = midi1.ClassSpecificMidiStreamingInterfaceDescriptorEmitter()

        outToHostJack = midi1.MidiOutJackDescriptorEmitter()
        outToHostJack.bJackID = 1
        outToHostJack.bJackType = midi1.MidiStreamingJackTypes.EMBEDDED
        outToHostJack.add_source(2)
        midi_streaming_interface.add_subordinate_descriptor(outToHostJack)

        inToDeviceJack = midi1.MidiInJackDescriptorEmitter()
        inToDeviceJack.bJackID = 2
        inToDeviceJack.bJackType = midi1.MidiStreamingJackTypes.EXTERNAL
        midi_streaming_interface.add_subordinate_descriptor(inToDeviceJack)

        inFromHostJack = midi1.MidiInJackDescriptorEmitter()
        inFromHostJack.bJackID = 3
        inFromHostJack.bJackType = midi1.MidiStreamingJackTypes.EMBEDDED
        midi_streaming_interface.add_subordinate_descriptor(inFromHostJack)

        outFromDeviceJack = midi1.MidiOutJackDescriptorEmitter()
        outFromDeviceJack.bJackID = 4
        outFromDeviceJack.bJackType = midi1.MidiStreamingJackTypes.EXTERNAL
        outFromDeviceJack.add_source(3)
        midi_streaming_interface.add_subordinate_descriptor(outFromDeviceJack)

        outEndpoint = midi1.StandardMidiStreamingBulkDataEndpointDescriptorEmitter()
        outEndpoint.bEndpointAddress = USBDirection.OUT.to_endpoint_address(3)
        outEndpoint.wMaxPacketSize = self.MAX_PACKET_SIZE_MIDI
        midi_streaming_interface.add_subordinate_descriptor(outEndpoint)

        outMidiEndpoint = midi1.ClassSpecificMidiStreamingBulkDataEndpointDescriptorEmitter()
        outMidiEndpoint.add_associated_jack(3)
        midi_streaming_interface.add_subordinate_descriptor(outMidiEndpoint)

        inEndpoint = midi1.StandardMidiStreamingBulkDataEndpointDescriptorEmitter()
        inEndpoint.bEndpointAddress = USBDirection.IN.to_endpoint_address(3)
        inEndpoint.wMaxPacketSize = self.MAX_PACKET_SIZE_MIDI
        midi_streaming_interface.add_subordinate_descriptor(inEndpoint)

        inMidiEndpoint = midi1.ClassSpecificMidiStreamingBulkDataEndpointDescriptorEmitter()
        inMidiEndpoint.add_associated_jack(1)
        midi_streaming_interface.add_subordinate_descriptor(inMidiEndpoint)

        return (midi_interface, midi_streaming_interface)
