/// USB Device Class for I2C transactions.
//
// It's not very good, and the API is weird. Someone with more USB device design experience could
// easily come up with something better.
//
// Control OUT transactions are used to perform I2C transfers to/from an internal buffer.
// Bulk IN/OUT transactions are used to transfer contents of the buffer to the host. It has not
// been optimized for speed or pipelining.
//
// To perform an I2C read:
//   1) Control OUT: ReadI2C(Address: 0xAA, Length: N)
//      (0xAA is the device address, N is the amount of bytes to read. Cannot be larger than
//      BUFFER_SIZE).
//      This performs an I2C read of N bytes into the inne buffer of the device, starting at
//      address 0.
//   2) Control IN: GetStatus()
//      The host ensures that the transaction was either ACK or NACK by getting one byte of status
//      from the device.
//   3) Control OUT: ReadBuffer(Address: X, Length: N)
//      (X is the address within the buffer, N is the amount of bytes to transfer to the host. N
//      cannot be larger than PACKET_SIZE).
//   4) Bulk IN: Read PACKET_SIZE bytes.
// Steps 3/4 can be skipped for scanning (the device won't mind the inner buffer not being read).
//
// To perform an I2C write:
//   1) Control OUT: SetWritePointer(Addrss: X)
//   2) Bulk OUT: Write at most PACKET_SIZE bytes.
// Repeat steps 1/2 to fill buffer with an I2c transaction.
//   3) Control OUT: WriteI2C(Address: 0x00, Length: N)
//      (0xAA is the device address, N is the amount of bytes to write. Cannot be larger than
//      BUFFER_SIZE).
//   4) Control IN: GetStatus()
//      The host ensures that the transaction was either ACK or NACK by getting one byte of status
//      from the device.

use embedded_hal::digital::v2::OutputPin;
use usb_device::class_prelude::*;
use nb::Error as NbError;
use stm32f1xx_hal::{
    gpio::{gpiob::*, Alternate, OpenDrain},
    i2c::{BlockingI2c, Error as I2CError},
    pac::I2C1,
    prelude::*,
};

use num_derive::FromPrimitive;
use num_traits::FromPrimitive;

use crate::{hprint, hprintln};

// Size of buffer within class, in bytes. Dictates maximum I2C transaction size.
const BUFFER_SIZE: usize = 1024;
// Size of bulk packets.
const PACKET_SIZE: usize = 64;

// All IN/OUT references bellow conform to typical USB naming, where IN: from device to host; OUT:
// from host to device.

/// Request number passed within Control IN requests to the I2C interface (ie. 'gets' from device).
#[derive(FromPrimitive)]
#[repr(u8)]
enum ControlInRequest {
    /// Write the current status as a single byte in response.
    GetStatus = 1,
}

/// Request number passed within Control OUT requests to the I2C interface (ie. 'sets' from the
/// host).
#[derive(FromPrimitive)]
#[repr(u8)]
enum ControlOutRequest {
    /// Set LED on or off (value == 0 -> off; on otherwise).
    SetLED = 1,

    /// Perform I2C bus read of a given length from a given I2C address.
    /// I2C Address: lower 8 bits of value.
    /// Read Length: upper 8 bits of value.
    ReadI2C = 2,

    /// Schedule a BULK IN transaction on the USB bus with the contents of the inner buffer.
    /// Buffer start address: lower 8 bits of value
    /// Read Length: upper 8 bits of value.
    ReadBuffer = 3,

    /// Perform I2C bus write of a given length to a given I2C address.
    /// I2C Address: lower 8 bits of value.
    /// Read Length: upper 8 bits of value.
    WriteI2C = 4,

    /// Set inner buffer write pointer. Any subsequent BULK OUT will write to the buffer at that
    /// address (but will not auto advance the pointer).
    SetWritePointer = 5,
}

/// Status of the I2C class. Combines information about requested transactions and I2C bus
/// responses.
#[derive(Copy, Clone)]
#[repr(u8)]
enum Status {
    /// Last request okay.
    OK = 0,
    /// Last request contained an invalid argument.
    InvalidArgument = 1,
    /// Last request okay, resulted in a successful I2C transaction.
    Ack = 2,
    /// Last request okay, resulted in a NACKd I2C transaction.
    Nack = 3,
    /// Last request okay, resulted in a fully failed I2C transaction.
    BusError = 4,
}

pub struct I2CClass<'a, B: UsbBus, LED> {
    interface: InterfaceNumber,
    /// Bulk IN endpoint for buffer transfers to host.
    ep_in: EndpointIn<'a, B>,
    /// Bulk OUT endpoint for buffer transfers from host.
    ep_out: EndpointOut<'a, B>,

    /// LED used for debugging.
    led: LED,

    /// The underlying I2C device.
    i2c_dev: BlockingI2c<I2C1, (PB6<Alternate<OpenDrain>>, PB7<Alternate<OpenDrain>>)>,

    /// Marker that is true when the host requested a BULK OUT via ReadBuffer.
    expect_bulk_out: bool,

    /// The underlying buffer and its write pointer.
    buffer: [u8; BUFFER_SIZE],
    write_pointer: usize,

    /// The device's main status byte, used by host to check whether operations were succesful.
    status: Status,
}

impl<B: UsbBus, LED: OutputPin> I2CClass<'_, B, LED> {
    pub fn new(
        alloc: &UsbBusAllocator<B>,
        led: LED,
        i2c_dev: BlockingI2c<I2C1, (PB6<Alternate<OpenDrain>>, PB7<Alternate<OpenDrain>>)>,
    ) -> I2CClass<'_, B, LED> {
        I2CClass {
            interface: alloc.interface(),
            ep_in: alloc.bulk(PACKET_SIZE as u16),
            ep_out: alloc.bulk(PACKET_SIZE as u16),
            led, i2c_dev,

            expect_bulk_out: false,

            buffer: [0; BUFFER_SIZE],
            write_pointer: 0usize,
            status: Status::OK,
        }
    }
}

impl<'a, B: UsbBus, LED: OutputPin> UsbClass<B> for I2CClass<'a, B, LED> {
    fn reset(&mut self) {
        self.expect_bulk_out = false;
        self.status = Status::OK;
    }

    fn control_in(&mut self, xfer: ControlIn<B>) {
        let req = xfer.request();

        if req.request_type != control::RequestType::Vendor
            || req.recipient != control::Recipient::Interface
            || req.index != u8::from(self.interface) as u16 {
            return
        }

        match FromPrimitive::from_u8(req.request) {
            /// Serve GetStatus: return this.status.
            Some(ControlInRequest::GetStatus) => {
                let status = self.status.clone() as u8;
                xfer.accept(|buf| {
                  buf[0] = status;
                  Ok(1usize)
                }).ok();
            },
            _ => {
                hprintln!("Unhandled control in on iface: {:?}", req).unwrap();
            },
        }
    }

    fn control_out(&mut self, xfer: ControlOut<B>) {
        let req = xfer.request();

        if req.request_type != control::RequestType::Vendor
            || req.recipient != control::Recipient::Interface
            || req.index != u8::from(self.interface) as u16 {
            return
        }

        match FromPrimitive::from_u8(req.request) {
            // Serve SetLED.
            Some(ControlOutRequest::SetLED) => {
                let on: bool = req.value > 0;
                match on {
                    true => self.led.set_low(),
                    false => self.led.set_high(),
                }.ok();
                xfer.accept().ok();
            },

            // Serve ReadI2C: read len bytes from I2C addr into internal buffer.
            Some(ControlOutRequest::ReadI2C) => {
                let addr: u8 = (req.value & 0xff) as u8;
                let len: u8 = (req.value >> 8) as u8;
                if len as usize > BUFFER_SIZE || len < 1u8 {
                    self.status = Status::InvalidArgument;
                    xfer.accept().ok();
                    return
                }
                if addr > 127u8 {
                    self.status = Status::InvalidArgument;
                    xfer.accept().ok();
                    return
                }
                match self.i2c_dev.read(addr, &mut self.buffer[0usize..(len as usize)]) {
                    Ok(_) => {
                        self.status = Status::Ack;
                    },
                    Err(NbError::Other(I2CError::Acknowledge)) => {
                        self.status = Status::Nack;
                    },
                    Err(e) => {
                        hprintln!("When reading I2C (addr {}, {} bytes): {:?}", addr, len, e).ok();
                        self.status = Status::BusError;
                    },
                }
                xfer.accept().ok();
            },

            // Serve ReadBuffer: send BULK IN with slice of buffer.
            Some(ControlOutRequest::ReadBuffer) => {
                let addr: u8 = (req.value & 0xff) as u8;
                let len: u8 = (req.value >> 8) as u8;

                if len as usize > PACKET_SIZE || len < 1u8 {
                    self.status = Status::InvalidArgument;
                    xfer.accept().ok();
                    return
                }

                let start = addr as usize;
                let end = (addr + len) as usize;
                if end as usize > BUFFER_SIZE {
                    self.status = Status::InvalidArgument;
                    xfer.accept().ok();
                    return
                }

                hprintln!("READ BUFFER, addr: {}, len: {}", addr, len).ok();

                self.status = Status::OK;
                xfer.accept().ok();
                match self.ep_in.write(&self.buffer[start..end]) {
                    Ok(count) => {
                    },
                    Err(UsbError::WouldBlock) => {},
                    Err(err) => {
                        hprintln!("bulk write failed: {:?}", err).ok();
                    },
                }
            },

            // Serve WriteI2C: write len bytes to I2C bus at addr from internal buffer.
            Some(ControlOutRequest::WriteI2C) => {
                let addr: u8 = (req.value & 0xff) as u8;
                let len: u8 = (req.value >> 8) as u8;
                if len as usize > BUFFER_SIZE || len < 1u8 {
                    self.status = Status::InvalidArgument;
                    xfer.accept().ok();
                    return
                }
                if addr > 127u8 {
                    self.status = Status::InvalidArgument;
                    xfer.accept().ok();
                    return
                }

                hprintln!("WRITE I2C, addr: {}, len: {}", addr, len).ok();
                match self.i2c_dev.write(addr, &self.buffer[0usize..(len as usize)]) {
                    Ok(_) => {
                        self.status = Status::Ack;
                    },
                    Err(NbError::Other(I2CError::Acknowledge)) => {
                        self.status = Status::Nack;
                    },
                    Err(e) => {
                        hprintln!("When writing I2C (addr {}, {} bytes): {:?}", addr, len, e).ok();
                        self.status = Status::BusError;
                    },
                }
                xfer.accept().ok();
            },

            // Serve SetWritePointer: set start address at which bytes from a BULK OUT will be
            // written to. The write pointer does _not_ increment on every write, so will need to
            // be manually controler after every BULK transfer.
            Some(ControlOutRequest::SetWritePointer) => {
                let pointer = req.value;
                if (pointer as usize) >= BUFFER_SIZE {
                    self.status = Status::InvalidArgument;
                    xfer.accept().ok();
                    return
                }
                hprintln!("SET WRITE PTR, pointer: {}", pointer).ok();
                self.write_pointer = pointer as usize;
                self.status = Status::OK;
                xfer.accept().ok();
            },
            _ => {
                hprintln!("Unhandled control out on iface: {:?}", req).ok();
            },
        }
    }

    fn get_configuration_descriptors(
        &self,
        writer: &mut DescriptorWriter,
    ) -> usb_device::Result<()> {
        writer.interface(
            self.interface,
            0xff,
            21, 37,
        )?;
        writer.endpoint(&self.ep_in)?;
        writer.endpoint(&self.ep_out)?;

        Ok(())
    }

    fn poll(&mut self) {
        let mut temp_buf = [0; PACKET_SIZE];
        // Serve BULK OUT writes - copy bytes into internal buffer.
        match self.ep_out.read(&mut temp_buf) {
            Ok(count) => {
                if self.expect_bulk_out {
                    self.expect_bulk_out = false;
                } else {
                    panic!("unexpectedly read data from bulk out endpoint");
                }
                hprintln!("SET BUFFER: ptr {}, {} bytes", self.write_pointer, count).ok();
                for (i, c) in temp_buf.iter().enumerate() {
                    let ptr = self.write_pointer + i;
                    // Silently drop bytes that do not fit in buffer.
                    if ptr >= BUFFER_SIZE {
                        continue;
                    }
                    self.buffer[ptr] = c.clone();
                }
            },
            Err(UsbError::WouldBlock) => {},
            Err(err) => panic!("bulk read {:?}", err),
        }
    }

    fn endpoint_out(&mut self, addr: EndpointAddress) {
        if addr == self.ep_out.address() {
            self.expect_bulk_out = true;
        }
    }

    fn endpoint_in_complete(&mut self, addr: EndpointAddress) {
        if addr == self.ep_in.address() {
            // TODO(q3k): should we be doing something here?
        }
    }
}