#![no_std]

use core::cmp::max;

use avr_device::interrupt::free;
use usb_device::{
    bus::{PollResult, UsbBus},
    endpoint::{EndpointAddress, EndpointType},
    Result as UsbResult, UsbDirection, UsbError,
};

mod types;
use types::{UsbDevice, DPRAM_SIZE, ENDPOINTS_ALLOC_LAYOUT};

impl<const L: usize> UsbBus for UsbDevice<L> {
    fn alloc_ep(
        &mut self,
        ep_dir: UsbDirection,
        ep_addr: Option<EndpointAddress>,
        ep_type: EndpointType,
        max_packet_size: u16,
        _interval: u8,
    ) -> UsbResult<EndpointAddress> {
        // Handle first endpoint. //
        if ep_addr == Some(EndpointAddress::from_parts(0, UsbDirection::In)) {
            return Ok(ep_addr.unwrap());
        }

        let address = match ep_addr {
            // If current endpoint doesn't allocated, assign ep_addr to variable. //
            Some(ep_addr) if !self.ep_table[ep_addr.index()].is_allocated => ep_addr,

            // If ep_aadr not provided, or current endpoint is allocated, try to find next free endpoint, otherwise return UsbError. //
            None | Some(_) => {
                let endpoint = self
                    .ep_table
                    .iter()
                    .enumerate()
                    .skip(1)
                    .find(|(i, &ep)| {
                        !ep.is_allocated && max_packet_size <= ENDPOINTS_ALLOC_LAYOUT[*i]
                    })
                    .ok_or(UsbError::EndpointMemoryOverflow)?;

                EndpointAddress::from_parts(endpoint.0, ep_dir)
            }
        };

        // Select endpoint info by address index. //

        let target_endpoint = &mut self.ep_table[address.index()];

        // Endpoint allocation marker. //

        if DPRAM_SIZE - self.dpram_already_used <= max_packet_size || max_packet_size >= 512 {
            Err(UsbError::EndpointMemoryOverflow)
        } else {
            // Get power of two number of endpoint size. //

            let max_packet_size = max(8, max_packet_size.next_power_of_two());

            // Set endpoint parameters. //

            target_endpoint.set_size(max_packet_size);
            target_endpoint.set_dir(ep_dir);
            target_endpoint.set_type(ep_type);
            target_endpoint.is_allocated = true;

            // Add used dpram memory. //

            self.dpram_already_used += max_packet_size;

            Ok(address)
        }
    }

    fn enable(&mut self) {
        free(|cs| {
            let (pll, usb) = (self.pll.borrow(cs), self.usb.borrow(cs));

            // Enable USB pads regulators. //

            usb.uhwcon.modify(|_, w| w.uvrege().set_bit());

            // Enable USB interface. //

            usb.usbcon
                .modify(|_, w| w.usbe().set_bit().frzclk().set_bit());

            // Configuring PLL. //

            pll.pllfrq
                .modify(|_, w| w.pdiv().mhz96().plltm().factor_15().pllusb().set_bit());

            // Enable PLL. //

            pll.pllcsr
                .modify(|_, w| w.pindiv().set_bit().plle().set_bit());

            while pll.pllcsr.read().plock().bit_is_clear() {}

            // Unfreeze clock. //

            usb.usbcon
                .modify(|_, w| w.frzclk().clear_bit().otgpade().set_bit());

            // Interrupts. //

            usb.udien
                .modify(|_, w| w.eorste().set_bit().sofe().set_bit());

            // Set high speed and attach the USB. //

            usb.udcon
                .modify(|_, w| w.lsm().set_bit().detach().clear_bit());
        })
    }

    fn force_reset(&self) -> UsbResult<()> {
        free(|cs| {
            let usb = self.usb.borrow(cs);

            usb.usbcon.modify(|_, w| w.usbe().clear_bit());
            usb.usbcon.modify(|_, w| w.usbe().set_bit());

            Ok(())
        })
    }

    fn is_stalled(&self, ep_addr: EndpointAddress) -> bool {
        free(|cs| match self.select_endpoint(cs, ep_addr.index()) {
            Ok(_) => self.usb.borrow(cs).ueconx.read().stallrq().bit_is_clear(),
            Err(_) => false,
        })
    }

    fn poll(&self) -> PollResult {
        free(|cs| {
            let usb = self.usb.borrow(cs);

            let (usbint, udint, udien) = (usb.usbint.read(), usb.udint.read(), usb.udien.read());
            if usbint.vbusti().bit_is_set() {
                usb.usbint.write(|w| w.vbusti().clear_bit());
                if usb.usbsta.read().vbus().bit_is_set() {
                    return PollResult::Resume;
                } else {
                    return PollResult::Suspend;
                }
            }
            if udint.suspi().bit_is_set() && udien.suspe().bit_is_set() {
                return PollResult::Suspend;
            }
            if udint.wakeupi().bit_is_set() && udien.wakeupe().bit_is_set() {
                return PollResult::Resume;
            }
            if udint.eorsti().bit_is_set() {
                return PollResult::Reset;
            }
            if udint.sofi().bit_is_set() {
                usb.udint.write(|w| w.sofi().clear_bit());
            }

            // Can only query endpoints while clock is running
            // (e.g. not in suspend state)
            if usb.usbcon.read().frzclk().bit_is_clear() {
                let (mut ep_out, mut ep_setup, mut ep_in_complete) = (0u8, 0u8, 0u8);

                for (index, _ep) in self
                    .ep_table
                    .iter()
                    .enumerate()
                    .filter(|(_i, e)| e.is_allocated)
                {
                    if self.select_endpoint(cs, index).is_err() {
                        // Endpoint selection has stopped working...
                        break;
                    }

                    let ueintx = usb.ueintx.read();
                    if ueintx.rxouti().bit_is_set() {
                        ep_out |= 1 << index;
                    }
                    if ueintx.rxstpi().bit_is_set() {
                        ep_setup |= 1 << index;
                    }
                    if ueintx.txini().bit_is_set() {
                        ep_in_complete |= 1 << index;
                    }
                }
                if ep_out | ep_setup | ep_in_complete != 0 {
                    return PollResult::Data {
                        ep_out: ep_out as u16,
                        ep_in_complete: ep_in_complete as u16,
                        ep_setup: ep_setup as u16,
                    };
                }
            }

            PollResult::None
        })
    }

    fn read(&self, ep_addr: EndpointAddress, buf: &mut [u8]) -> UsbResult<usize> {
        free(|cs| {
            let usb = self.usb.borrow(cs);

            match self.select_endpoint(cs, ep_addr.index()) {
                Ok(()) => {
                    let target_endpoint = self.ep_table[ep_addr.index()];

                    let ueintx = usb.ueintx.read();

                    if ueintx.rxouti().bit_is_clear() {
                        return Err(UsbError::WouldBlock);
                    }

                    if target_endpoint.ep_type == 0 {
                        let bytes_count_to_read: usize = (usb.uebchx.read().bits() as usize) << 8
                            | (usb.uebclx.read().bits() as usize);

                        if bytes_count_to_read > buf.len() {
                            return Err(UsbError::BufferOverflow);
                        }

                        for slot in &mut buf[..bytes_count_to_read] {
                            *slot = usb.uedatx.read().bits();
                        }

                        usb.ueintx
                            .write(|w| w.rxouti().clear_bit().rxstpi().clear_bit());

                        Ok(bytes_count_to_read)
                    } else {
                        usb.ueintx.write(|w| w.rxouti().clear_bit());

                        let mut bytes_read = 0;
                        for slot in buf {
                            if usb.ueintx.read().rwal().bit_is_clear() {
                                break;
                            }
                            *slot = usb.uedatx.read().bits();
                            bytes_read += 1;
                        }

                        if usb.ueintx.read().rwal().bit_is_set() {
                            return Err(UsbError::BufferOverflow);
                        }

                        usb.ueintx.write(|w| w.fifocon().clear_bit());
                        Ok(bytes_read)
                    }
                }
                Err(err) => Err(err),
            }
        })
    }

    fn reset(&self) {
        free(|cs| {
            let usb = self.usb.borrow(cs);

            usb.udint.modify(|_, w| w.eorsti().clear_bit());

            // Disabling all endpoints before it reset //

            self.ep_table
                .iter()
                .filter(|&&ep| ep.is_allocated)
                .enumerate()
                .for_each(|(index, _ep)| {
                    if self.select_endpoint(cs, index).is_ok() {
                        usb.ueconx.modify(|_, w| w.epen().clear_bit());
                    }
                });

            // Reset endpoints //

            usb.uerst.modify(|_, w| unsafe { w.bits(u8::MAX >> 1) });

            // Clear resume informations. //

            usb.udint
                .modify(|_, w| w.wakeupi().clear_bit().suspi().clear_bit());
            usb.udien
                .modify(|_, w| w.wakeupe().clear_bit().suspe().set_bit());
        })
    }

    fn resume(&self) {
        free(|cs| {
            let usb = self.usb.borrow(cs);
            let pll = self.pll.borrow(cs);

            // Enable PLL and wait PLL lock. //

            pll.pllcsr.modify(|_, w| w.plle().set_bit());

            while pll.pllcsr.read().plock().bit_is_clear() {}

            // Unfreeze USB clock. //

            usb.usbcon.modify(|_, w| w.frzclk().clear_bit());

            // Clear resume informations. //

            usb.udint
                .modify(|_, w| w.wakeupi().clear_bit().suspi().clear_bit());

            usb.udien
                .modify(|_, w| w.wakeupe().clear_bit().suspe().set_bit());
        })
    }

    fn set_device_address(&self, addr: u8) {
        free(|cs| {
            let usb = self.usb.borrow(cs);

            // Set address. //

            usb.udaddr.modify(|_, w| w.uadd().bits(addr));

            // Note: ADDEN and UADD shall not be written at the same time.
            // (written in atmega32u4/16u4 docs)

            // Enable. //

            usb.udaddr.modify(|_, w| w.adden().set_bit());
        });
    }

    fn set_stalled(&self, ep_addr: EndpointAddress, stalled: bool) {
        free(|cs| {
            let usb = self.usb.borrow(cs);
            if self.select_endpoint(cs, ep_addr.index()).is_ok() {
                usb.ueconx
                    .modify(|_, w| w.stallrq().bit(stalled).stallrqc().bit(!stalled));
            }
        });
    }

    fn suspend(&self) {
        free(|cs| {
            let usb = self.usb.borrow(cs);
            let pll = self.pll.borrow(cs);

            usb.udint
                .modify(|_, w| w.suspi().clear_bit().wakeupi().clear_bit());
            usb.udien
                .modify(|_, w| w.suspe().clear_bit().wakeupe().clear_bit());

            usb.usbcon.modify(|_, w| w.frzclk().set_bit());

            pll.pllcsr.modify(|_, w| w.plle().clear_bit());
        })
    }

    fn write(&self, ep_addr: EndpointAddress, buf: &[u8]) -> UsbResult<usize> {
        free(|cs| {
            let usb = self.usb.borrow(cs);

            match self.select_endpoint(cs, ep_addr.index()) {
                Ok(()) => {
                    let target_endpoint = self.ep_table[ep_addr.index()];

                    let ueintx = usb.ueintx.read();

                    if ueintx.rxouti().bit_is_clear() {
                        return Err(UsbError::WouldBlock);
                    }

                    if target_endpoint.ep_type == 0 {
                        let bytes_count_to_read: usize = (usb.uebchx.read().bits() as usize) << 8
                            | (usb.uebclx.read().bits() as usize);

                        if bytes_count_to_read > buf.len() {
                            return Err(UsbError::BufferOverflow);
                        }

                        buf.iter()
                            .for_each(|&byte| usb.uedatx.write(|w| w.bits(byte)));

                        usb.ueintx
                            .write(|w| w.rxouti().clear_bit().rxstpi().clear_bit());

                        Ok(bytes_count_to_read)
                    } else {
                        usb.ueintx
                            .write(|w| w.txini().clear_bit().rxouti().clear_bit());

                        for &byte in buf {
                            if usb.ueintx.read().rwal().bit_is_set() {
                                return Err(UsbError::BufferOverflow);
                            } else {
                                usb.uedatx.write(|w| w.bits(byte));
                            }
                        }

                        usb.ueintx
                            .write(|w| w.fifocon().clear_bit().rxouti().clear_bit());
                        Ok(buf.len())
                    }
                }
                Err(err) => Err(err),
            }
        })
    }
}