#![no_std]

use core::cmp::max;

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

mod types;
pub use types::*;

use types::{DPRAM_SIZE, ENDPOINTS_ALLOC_LAYOUT, ONE_MS_16_MGHZ};

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. //
        free(|_cs| {
            if ep_addr == Some(EndpointAddress::from_parts(0, UsbDirection::In)) {
                Ok(ep_addr.unwrap())
            } else {
                let address = match ep_addr {
                    // If current endpoint doesn't allocated, assign ep_addr to variable. //
                    Some(addr) if !self.ep_table[addr.index()].is_allocated => addr,

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

                        EndpointAddress::from_parts(index, ep_dir)
                    }
                };

                // Select endpoint info by address index. //

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

                // Get power of two number of endpoint size. //

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

                // Endpoint allocation marker. //

                if DPRAM_SIZE - self.dpram_already_used < ep_size {
                    Err(UsbError::EndpointMemoryOverflow)
                } else {
                    // Set endpoint parameters. //

                    target_endpoint.set_dir(ep_dir);
                    target_endpoint.set_type(ep_type);
                    target_endpoint.set_size(ep_size)?;

                    // Add used dpram memory. //

                    target_endpoint.is_allocated = true;
                    self.dpram_already_used += ep_size;

                    Ok(address)
                }
            }
        })
    }

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

            // Enable USB pads regulators. //

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

            // PLL configuration //

            pll.pllcsr.write(|w| w.pindiv().set_bit());

            pll.pllfrq
                .write(|w| w.pdiv().mhz96().plltm().factor_15().pllusb().set_bit());

            // Enable PLL //

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

            // Check PLL lock //

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

            // Enable USB interface. //

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

            // Unfreeze clock. //

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

            // Endpoint configuration //

            self.allocated_endpoints().for_each(|(i, _ep)| {
                self.configure_endpoint(cs, i).unwrap();
            });

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

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

            // Interrupts. //

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

    fn force_reset(&self) -> UsbResult<()> {
        free(|cs| {
            let usbcon = &self.usb.borrow(cs).usbcon;
            usbcon.modify(|_, w| w.usbe().set_bit());
        });

        delay_cycles(ONE_MS_16_MGHZ);

        free(|cs| {
            let usbcon = &self.usb.borrow(cs).usbcon;
            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_set(),
            Err(_) => false,
        })
    }

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

            let (udint, udien, usbint) = (usb.udint.read(), usb.udien.read(), usb.usbint.read());

            if usbint.vbusti().bit_is_set() {
                usb.usbint
                    .write(|w| unsafe { w.bits(0x01) }.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| unsafe { w.bits(0x7d) }.sofi().clear_bit());
            }

            if usb.usbcon.read().frzclk().bit_is_clear() {
                let (mut ep_out, mut ep_in_complete, mut ep_setup) = (0u16, 0u16, 0u16);
                let pending_ins = self.pending_ins.borrow(cs);

                for (ep_index, _ep) in self.allocated_endpoints() {
                    if self.select_endpoint(cs, ep_index).is_err() {
                        break;
                    } else {
                        let ueintx = usb.ueintx.read();

                        if ueintx.rxouti().bit_is_set() {
                            ep_out |= 1 << ep_index;
                        }
                        if ueintx.rxstpi().bit_is_set() {
                            ep_setup |= 1 << ep_index;
                        }

                        if pending_ins.get() & (1 << ep_index) != 0 && ueintx.txini().bit_is_set() {
                            ep_in_complete |= 1 << ep_index;
                            pending_ins.set(pending_ins.get() & !(1 << ep_index));
                        }
                    }
                }
                if ep_out | ep_in_complete | ep_setup != 0 {
                    return PollResult::Data {
                        ep_out,
                        ep_in_complete,
                        ep_setup,
                    };
                }
            }
            PollResult::None
        })
    }

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

            if let Err(error) = self.select_endpoint(cs, ep_addr.index()) {
                Err(error)
            } else {
                let ep = &self.ep_table[ep_addr.index()];

                if ep.ep_type == 0 {
                    let ueintx = usb.ueintx.read();

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

                    let buf_size = self.get_size(cs);
                    if buf.len() < buf_size {
                        return Err(UsbError::BufferOverflow);
                    }

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

                    usb.ueintx.write(|w| {
                        unsafe { w.bits(0xdf) }
                            .rxouti()
                            .clear_bit()
                            .rxstpi()
                            .clear_bit()
                    });

                    Ok(buf_size)
                } else {
                    if usb.ueintx.read().rxouti().bit_is_clear() {
                        return Err(UsbError::WouldBlock);
                    }
                    usb.ueintx
                        .write(|w| unsafe { w.bits(0xdf) }.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| unsafe { w.bits(0xdf) }.fifocon().clear_bit());
                    Ok(bytes_read)
                }
            }
        })
    }

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

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

            self.allocated_endpoints().for_each(|(i, _)| {
                self.configure_endpoint(cs, i).unwrap();
            });

            // Clear resume informations. //

            usb.udint.write(|w| {
                unsafe { w.bits(0x7d) }
                    .wakeupi()
                    .clear_bit()
                    .suspi()
                    .clear_bit()
            });

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

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

            // PLL enable //

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

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

            // Resuming //

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

            usb.udint.write(|w| {
                unsafe { w.bits(0x7d) }
                    .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, pll) = (self.usb.borrow(cs), self.pll.borrow(cs));

            usb.udint.write(|w| {
                unsafe { w.bits(0x7d) }
                    .wakeupi()
                    .clear_bit()
                    .suspi()
                    .clear_bit()
            });

            // Suspend. //

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

            // Freeze clock. //

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

            // Disable PLL. //

            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);

            if let Err(error) = self.select_endpoint(cs, ep_addr.index()) {
                Err(error)
            } else {
                let ep = &self.ep_table[ep_addr.index()];

                // Endpoint type confitions //

                if ep.ep_type == 0 {
                    if usb.ueintx.read().txini().bit_is_clear() {
                        return Err(UsbError::WouldBlock);
                    }

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

                    for &byte in buf {
                        usb.uedatx.write(|w| w.bits(byte));
                    }

                    usb.ueintx
                        .write(|w| unsafe { w.bits(0xdf) }.txini().clear_bit());
                } else {
                    if usb.ueintx.read().txini().bit_is_clear() {
                        return Err(UsbError::WouldBlock);
                    }
                    usb.ueintx.write(|w| {
                        unsafe { w.bits(0xdf) }
                            .txini()
                            .clear_bit()
                            .rxouti()
                            .clear_bit()
                    });

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

                    usb.ueintx.write(|w| {
                        unsafe { w.bits(0xdf) }
                            .rxouti()
                            .clear_bit()
                            .fifocon()
                            .clear_bit()
                    });
                }

                let pending_ins = self.pending_ins.borrow(cs);
                pending_ins.set(pending_ins.get() | 1 << ep_addr.index());

                Ok(buf.len())
            }
        })
    }
}