pub mod addr;
pub mod ctrl_pipe;
pub mod ext_reg;
pub mod pck_size;
pub mod status_bk;
pub mod status_pipe;

use addr::Addr;
use ctrl_pipe::CtrlPipe;
use ext_reg::ExtReg;
use pck_size::PckSize;
use status_bk::StatusBk;
use status_pipe::StatusPipe;

use super::device::{Endpoint, TransferType};

use super::usbproto::*;

use atsamd_hal::target_device::usb::{
    self,
    host::{BINTERVAL, PCFG, PINTFLAG, PSTATUS, PSTATUSCLR, PSTATUSSET},
};
use core::convert::TryInto;
use log::{trace, warn};

// Maximum time to wait for a control request with data to finish. cf
// §9.2.6.1 of USB 2.0.
const USB_TIMEOUT: usize = 5 * 1024; // 5 Seconds

// samd21 only supports 8 pipes.
const MAX_PIPES: usize = 8;

// How many times to retry a transaction that has transient errors.
const NAK_LIMIT: usize = 15;

// For use in control requests that do not require a data stage.
pub(crate) const NO_DATA_STAGE: Option<&mut ()> = None;

#[derive(Copy, Clone, Debug, PartialEq)]
pub(crate) enum PipeErr {
    ShortPacket,
    InvalidPipe,
    InvalidToken,
    Stall,
    TransferFail,
    PipeErr,
    Flow,
    HWTimeout,
    DataToggle,
    SWTimeout,
    Other(&'static str),
}
impl From<&'static str> for PipeErr {
    fn from(v: &'static str) -> Self {
        Self::Other(v)
    }
}

pub(crate) struct PipeTable {
    tbl: [PipeDesc; MAX_PIPES],
}

impl PipeTable {
    pub(crate) fn new() -> Self {
        let tbl = {
            let mut tbl: [core::mem::MaybeUninit<PipeDesc>; MAX_PIPES] =
                unsafe { core::mem::MaybeUninit::uninit().assume_init() };

            for e in &mut tbl[..] {
                unsafe { core::ptr::write(e.as_mut_ptr(), PipeDesc::new()) }
            }

            unsafe { core::mem::transmute(tbl) }
        };
        Self { tbl: tbl }
    }

    pub(crate) fn pipe_for<'a, 'b>(
        &'a mut self,
        host: &'b mut usb::HOST,
        endpoint: &Endpoint,
    ) -> Pipe<'a, 'b> {
        // Just use two pipes for now. 0 is always for control
        // endpoints, 1 for everything else.
        //
        // TODO: cache in-use pipes and return them without init if
        // possible.
        let i = if endpoint.num == 0 { 0 } else { 1 };

        let pregs = PipeRegs::from(host, i);
        let pdesc = &mut self.tbl[i];

        pregs.cfg.write(|w| {
            let ptype = PType::from(endpoint.transfer_type) as u8;
            unsafe { w.ptype().bits(ptype) }
        });
        trace!(
            "setting paddr of pipe {} to {}:{}",
            i,
            endpoint.addr,
            endpoint.num
        );
        pdesc.bank0.ctrl_pipe.write(|w| {
            w.pdaddr().set_addr(endpoint.addr);
            w.pepnum().set_epnum(endpoint.num)
        });
        Pipe {
            num: i,
            regs: pregs,
            desc: pdesc,
        }
    }
}

// TODO: hide regs/desc fields. Needed right now for init_pipe0.
pub(crate) struct Pipe<'a, 'b> {
    num: usize,

    pub(crate) regs: PipeRegs<'b>,
    pub(crate) desc: &'a mut PipeDesc,
}
impl Pipe<'_, '_> {
    pub(crate) fn control_transfer<T>(
        &mut self,
        ep: &mut Endpoint,
        bm_request_type: RequestType,
        b_request: RequestCode,
        w_value: WValue,
        w_index: u16,
        buf: Option<&mut T>,
        millis: &dyn Fn() -> usize,
    ) -> Result<(), PipeErr> {
        let len = match buf {
            None => 0,
            _ => core::mem::size_of::<T>(),
        };

        /*
         * Setup stage.
         */
        let mut setup_packet = SetupPacket {
            bm_request_type: bm_request_type,
            b_request: b_request,
            w_value: w_value,
            w_index: w_index,
            w_length: len as u16,
        };
        self.send(
            ep,
            PToken::Setup,
            &DataBuf::from(&mut setup_packet),
            NAK_LIMIT,
            millis,
        )?;

        /*
         * Data stage.
         */
        if let Some(b) = buf {
            // TODO: data stage, has up to 5,000ms (in 500ms
            // per-packet chunks) to complete. cf §9.2.6.4 of USB 2.0.
            match bm_request_type.direction()? {
                RequestDirection::DeviceToHost => {
                    self.in_transfer(ep, b, NAK_LIMIT, millis)?;
                }

                RequestDirection::HostToDevice => {
                    self.out_transfer(ep, b, NAK_LIMIT, millis)?;
                }
            }
        }

        /*
         * Status stage.
         */
        // TODO: status stage has up to 50ms to complete. cf §9.2.6.4
        // of USB 2.0.
        self.desc.bank0.pcksize.write(|w| {
            unsafe { w.byte_count().bits(0) };
            unsafe { w.multi_packet_size().bits(0) }
        });

        let token = match bm_request_type.direction()? {
            RequestDirection::DeviceToHost => PToken::Out,
            RequestDirection::HostToDevice => PToken::In,
        };

        trace!("dispatching status stage");
        self.dispatch_retries(ep, token, NAK_LIMIT, millis)?;

        Ok(())
    }

    fn send(
        &mut self,
        ep: &mut Endpoint,
        token: PToken,
        buf: &DataBuf,
        nak_limit: usize,
        millis: &dyn Fn() -> usize,
    ) -> Result<(), PipeErr> {
        trace!("p{}: sending {:?}", self.num, buf);

        self.desc
            .bank0
            .addr
            .write(|w| unsafe { w.addr().bits(buf.ptr as u32) });
        // configure packet size PCKSIZE.SIZE
        self.desc.bank0.pcksize.write(|w| {
            unsafe { w.byte_count().bits(buf.len as u16) };
            unsafe { w.multi_packet_size().bits(0) }
        });

        self.dispatch_retries(ep, token, nak_limit, millis)
    }

    pub(crate) fn in_transfer<T>(
        &mut self,
        ep: &mut Endpoint,
        buf: &mut T,
        nak_limit: usize,
        millis: &dyn Fn() -> usize,
    ) -> Result<usize, PipeErr> {
        // TODO: pull this from pipe descriptor for this addr/ep.
        let packet_size = 8;

        let db: DataBuf = buf.into();

        trace!("p{}: Should IN for {}b.", self.num, db.len);
        self.desc.bank0.pcksize.write(|w| {
            unsafe { w.byte_count().bits(db.len as u16) };
            unsafe { w.multi_packet_size().bits(0) }
        });

        // Read until we get a short packet (indicating that there's
        // nothing left for us in this transaction) or the buffer is
        // full.
        //
        // TODO: It is sometimes valid to get a short packet when
        // variable length data is desired by the driver. cf §5.3.2 of
        // USB 2.0.
        let mut bytes_received = 0;
        while bytes_received < db.len {
            // Move the buffer pointer forward as we get data.
            self.desc
                .bank0
                .addr
                .write(|w| unsafe { w.addr().bits(db.ptr as u32 + bytes_received as u32) });
            self.regs.statusclr.write(|w| w.bk0rdy().set_bit());
            trace!("--- !!! regs-pre-dispatch !!! ---");
            self.log_regs();

            self.dispatch_retries(ep, PToken::In, nak_limit, millis)?;
            let recvd = self.desc.bank0.pcksize.read().byte_count().bits() as usize;
            bytes_received += recvd;
            trace!("!! read {} of {}", bytes_received, db.len);
            if recvd < packet_size {
                break;
            }

            // Don't allow writing past the buffer.
            assert!(bytes_received <= db.len);
        }

        self.regs.statusset.write(|w| w.pfreeze().set_bit());
        if bytes_received < db.len {
            self.log_regs();
            // TODO: honestly, this is probably a panic condition,
            // since whatever's in DataBuf.ptr is totally
            // invalid. Alternately, this function should be declared
            // `unsafe`. OR! Make the function take a mutable slice of
            // u8, and leave it up to the caller to figure out how to
            // deal with short packets.
            Err(PipeErr::ShortPacket)
        } else {
            Ok(bytes_received)
        }
    }

    pub(crate) fn out_transfer<T>(
        &mut self,
        ep: &mut Endpoint,
        buf: &mut T,
        nak_limit: usize,
        millis: &dyn Fn() -> usize,
    ) -> Result<usize, PipeErr> {
        let db: DataBuf = buf.into();

        trace!("p{}: Should OUT for {}b.", self.num, db.len);
        self.desc.bank0.pcksize.write(|w| {
            unsafe { w.byte_count().bits(db.len as u16) };
            unsafe { w.multi_packet_size().bits(0) }
        });

        let mut bytes_sent = 0;
        while bytes_sent < db.len {
            self.desc
                .bank0
                .addr
                .write(|w| unsafe { w.addr().bits(db.ptr as u32 + bytes_sent as u32) });
            self.dispatch_retries(ep, PToken::Out, nak_limit, millis)?;

            let sent = self.desc.bank0.pcksize.read().byte_count().bits() as usize;
            bytes_sent += sent;
            trace!("!! wrote {} of {}", bytes_sent, db.len);
        }

        Ok(bytes_sent)
    }

    pub(crate) fn dtgl(&mut self, ep: &mut Endpoint, token: PToken) {
        // TODO: this makes no sense to me, and docs are unclear. If
        // the status bit is set, set it again? if it's clear then
        // clear it? This is what I get for having to work from
        // Arduino sources.
        trace!(
            "~~~ tok: {:?}, dtgl: {}, i: {}, o: {}",
            token,
            self.regs.status.read().dtgl().bit(),
            ep.in_toggle,
            ep.out_toggle,
        );
        if self.regs.status.read().dtgl().bit_is_set() {
            self.dtgl_clear();
        } else {
            self.dtgl_set();
        }
        if token == PToken::In {
            ep.in_toggle = self.regs.status.read().dtgl().bit_is_set()
        } else {
            ep.out_toggle = self.regs.status.read().dtgl().bit_is_set()
        }
    }

    pub(crate) fn dtgl_set(&mut self) {
        self.regs.statusset.write(|w| w.dtgl().set_bit());
    }

    pub(crate) fn dtgl_clear(&mut self) {
        self.regs.statusclr.write(|w| unsafe {
            // No function for this. FIXME: need to patch the SVD for
            // PSTATUSCLR.DTGL at bit0. No? This is in the SVD, but
            // not the rust output.
            w.bits(1)
        });
    }

    // This is the only function that calls `millis`. If we can make
    // this just take the current timestamp, we can make this
    // non-blocking.
    fn dispatch_retries(
        &mut self,
        ep: &mut Endpoint,
        token: PToken,
        retries: usize,
        millis: &dyn Fn() -> usize,
    ) -> Result<(), PipeErr> {
        self.dispatch_packet(ep, token);

        let until = millis() + USB_TIMEOUT;
        let mut last_err = PipeErr::SWTimeout;
        let mut naks = 0;
        while millis() < until {
            let res = self.dispatch_result(token);
            match res {
                Ok(true) => {
                    if token == PToken::In {
                        ep.in_toggle = self.regs.status.read().dtgl().bit_is_set();
                    } else if token == PToken::Out {
                        ep.out_toggle = self.regs.status.read().dtgl().bit_is_set();
                    }
                    return Ok(());
                }
                Ok(false) => continue,

                Err(e) => {
                    last_err = e;
                    match last_err {
                        PipeErr::DataToggle => self.dtgl(ep, token),
                        PipeErr::Stall => break,
                        _ => {
                            naks += 1;
                            if naks > retries {
                                break;
                            } else {
                                self.dispatch_packet(ep, token);
                            }
                        }
                    }
                }
            }
        }

        Err(last_err)
    }

    fn dispatch_packet(&mut self, ep: &mut Endpoint, token: PToken) {
        self.regs
            .cfg
            .modify(|_, w| unsafe { w.ptoken().bits(token as u8) });
        self.regs.intflag.modify(|_, w| w.trfail().set_bit());
        self.regs.intflag.modify(|_, w| w.perr().set_bit());

        match token {
            PToken::Setup => {
                self.regs.intflag.write(|w| w.txstp().set_bit());
                self.regs.statusset.write(|w| w.bk0rdy().set_bit());

                // Toggles should be 1 for host and function's
                // sequence at end of setup transaction. cf §8.6.1 of
                // USB 2.0.
                self.dtgl_clear();
                ep.in_toggle = true;
                ep.out_toggle = true;
            }
            PToken::In => {
                self.regs.statusclr.write(|w| w.bk0rdy().set_bit());
                if ep.in_toggle {
                    self.dtgl_set();
                } else {
                    self.dtgl_clear();
                }
            }
            PToken::Out => {
                self.regs.intflag.write(|w| w.trcpt0().set_bit());
                self.regs.statusset.write(|w| w.bk0rdy().set_bit());
                if ep.out_toggle {
                    self.dtgl_set();
                } else {
                    self.dtgl_clear();
                }
            }
            _ => {}
        }

        trace!("initial regs");
        self.log_regs();

        self.regs.statusclr.write(|w| w.pfreeze().set_bit());
    }

    fn dispatch_result(&mut self, token: PToken) -> Result<bool, PipeErr> {
        if self.is_transfer_complete(token)? {
            self.regs.statusset.write(|w| w.pfreeze().set_bit());
            Ok(true)
        } else if self.regs.intflag.read().trfail().bit_is_set() {
            self.regs.intflag.write(|w| w.trfail().set_bit());
            trace!("trfail");
            self.regs.statusset.write(|w| w.pfreeze().set_bit());
            self.log_regs();
            Err(PipeErr::TransferFail)
        } else if self.desc.bank0.status_bk.read().errorflow().bit_is_set() {
            trace!("errorflow");
            self.regs.statusset.write(|w| w.pfreeze().set_bit());
            self.log_regs();
            Err(PipeErr::Flow)
        } else if self.desc.bank0.status_pipe.read().touter().bit_is_set() {
            trace!("touter");
            self.regs.statusset.write(|w| w.pfreeze().set_bit());
            self.log_regs();
            Err(PipeErr::HWTimeout)
        } else if self.desc.bank0.status_pipe.read().dtgler().bit_is_set() {
            trace!("dtgler");
            self.regs.statusset.write(|w| w.pfreeze().set_bit());
            self.log_regs();
            Err(PipeErr::DataToggle)
        } else {
            // Nothing wrong, but not done yet.
            Ok(false)
        }
    }

    fn is_transfer_complete(&mut self, token: PToken) -> Result<bool, PipeErr> {
        match token {
            PToken::Setup => {
                if self.regs.intflag.read().txstp().bit_is_set() {
                    self.regs.intflag.write(|w| w.txstp().set_bit());
                    Ok(true)
                } else {
                    Ok(false)
                }
            }
            PToken::In => {
                if self.regs.intflag.read().trcpt0().bit_is_set() {
                    self.regs.intflag.write(|w| w.trcpt0().set_bit());
                    Ok(true)
                } else {
                    Ok(false)
                }
            }
            PToken::Out => {
                if self.regs.intflag.read().trcpt0().bit_is_set() {
                    self.regs.intflag.write(|w| w.trcpt0().set_bit());
                    Ok(true)
                } else {
                    Ok(false)
                }
            }
            _ => Err(PipeErr::InvalidToken),
        }
    }

    fn log_regs(&self) {
        // Pipe regs
        let cfg = self.regs.cfg.read().bits();
        let bin = self.regs.binterval.read().bits();
        let sts = self.regs.status.read().bits();
        let ifl = self.regs.intflag.read().bits();
        trace!(
            "p{}: cfg: {:x}, bin: {:x}, sts: {:x}, ifl: {:x}",
            self.num,
            cfg,
            bin,
            sts,
            ifl
        );

        // Pipe RAM regs
        let adr = self.desc.bank0.addr.read().bits();
        let pks = self.desc.bank0.pcksize.read().bits();
        let ext = self.desc.bank0.extreg.read().bits();
        let sbk = self.desc.bank0.status_bk.read().bits();
        let hcp = self.desc.bank0.ctrl_pipe.read().bits();
        let spi = self.desc.bank0.status_pipe.read().bits();
        trace!(
            "p{}: adr: {:x}, pks: {:x}, ext: {:x}, sbk: {:x}, hcp: {:x}, spi: {:x}",
            self.num,
            adr,
            pks,
            ext,
            sbk,
            hcp,
            spi
        );
    }
}

// TODO: merge into SVD for pipe cfg register.
#[derive(Copy, Clone, Debug, PartialEq)]
pub(crate) enum PToken {
    Setup = 0x0,
    In = 0x1,
    Out = 0x2,
    _Reserved = 0x3,
}

// TODO: merge into SVD for pipe cfg register.
#[derive(Copy, Clone, Debug, PartialEq)]
pub(crate) enum PType {
    Disabled = 0x0,
    Control = 0x1,
    ISO = 0x2,
    Bulk = 0x3,
    Interrupt = 0x4,
    Extended = 0x5,
    _Reserved0 = 0x06,
    _Reserved1 = 0x07,
}
impl From<TransferType> for PType {
    fn from(v: TransferType) -> Self {
        match v {
            TransferType::Control => Self::Control,
            TransferType::Isochronous => Self::ISO,
            TransferType::Bulk => Self::Bulk,
            TransferType::Interrupt => Self::Interrupt,
        }
    }
}

struct DataBuf<'a> {
    ptr: *mut u8,
    len: usize,
    _marker: core::marker::PhantomData<&'a ()>,
}
impl DataBuf<'_> {}

impl core::fmt::Debug for DataBuf<'_> {
    fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
        write!(f, "DataBuf {{ len: {}, ptr: [", self.len)?;
        for i in 0..self.len {
            write!(f, " {:x}", unsafe {
                *self.ptr.offset(i.try_into().unwrap())
            })?;
        }
        write!(f, " ] }}")
    }
}

impl<'a, T> From<&'a mut T> for DataBuf<'a> {
    fn from(v: &'a mut T) -> Self {
        Self {
            ptr: v as *mut T as *mut u8,
            len: core::mem::size_of::<T>(),
            _marker: core::marker::PhantomData,
        }
    }
}

pub(crate) struct PipeRegs<'a> {
    pub(crate) cfg: &'a mut PCFG,
    pub(crate) binterval: &'a mut BINTERVAL,
    pub(crate) statusclr: &'a mut PSTATUSCLR,
    pub(crate) statusset: &'a mut PSTATUSSET,
    pub(crate) status: &'a mut PSTATUS,
    pub(crate) intflag: &'a mut PINTFLAG,
}
impl<'a> PipeRegs<'a> {
    pub(crate) fn from(host: &'a mut usb::HOST, i: usize) -> PipeRegs {
        assert!(i < MAX_PIPES);
        match i {
            0 => Self {
                cfg: &mut host.pcfg0,
                binterval: &mut host.binterval0,
                statusclr: &mut host.pstatusclr0,
                statusset: &mut host.pstatusset0,
                status: &mut host.pstatus0,
                intflag: &mut host.pintflag0,
            },
            1 => Self {
                cfg: &mut host.pcfg1,
                binterval: &mut host.binterval1,
                statusclr: &mut host.pstatusclr1,
                statusset: &mut host.pstatusset1,
                status: &mut host.pstatus1,
                intflag: &mut host.pintflag1,
            },
            2 => Self {
                cfg: &mut host.pcfg2,
                binterval: &mut host.binterval2,
                statusclr: &mut host.pstatusclr2,
                statusset: &mut host.pstatusset2,
                status: &mut host.pstatus2,
                intflag: &mut host.pintflag2,
            },
            3 => Self {
                cfg: &mut host.pcfg3,
                binterval: &mut host.binterval3,
                statusclr: &mut host.pstatusclr3,
                statusset: &mut host.pstatusset3,
                status: &mut host.pstatus3,
                intflag: &mut host.pintflag3,
            },
            4 => Self {
                cfg: &mut host.pcfg4,
                binterval: &mut host.binterval4,
                statusclr: &mut host.pstatusclr4,
                statusset: &mut host.pstatusset4,
                status: &mut host.pstatus4,
                intflag: &mut host.pintflag4,
            },
            5 => Self {
                cfg: &mut host.pcfg5,
                binterval: &mut host.binterval5,
                statusclr: &mut host.pstatusclr5,
                statusset: &mut host.pstatusset5,
                status: &mut host.pstatus5,
                intflag: &mut host.pintflag5,
            },
            6 => Self {
                cfg: &mut host.pcfg6,
                binterval: &mut host.binterval6,
                statusclr: &mut host.pstatusclr6,
                statusset: &mut host.pstatusset6,
                status: &mut host.pstatus6,
                intflag: &mut host.pintflag6,
            },
            7 => Self {
                cfg: &mut host.pcfg7,
                binterval: &mut host.binterval7,
                statusclr: &mut host.pstatusclr7,
                statusset: &mut host.pstatusset7,
                status: &mut host.pstatus7,
                intflag: &mut host.pintflag7,
            },
            _ => unreachable!(),
        }
    }
}

// §32.8.7.1
pub(crate) struct PipeDesc {
    pub bank0: BankDesc,
    pub bank1: BankDesc,
}

// 2 banks: 32 bytes per pipe.
impl PipeDesc {
    pub fn new() -> Self {
        Self {
            bank0: BankDesc::new(),
            bank1: BankDesc::new(),
        }
    }
}

#[repr(C, packed)]
// 16 bytes per bank.
pub(crate) struct BankDesc {
    pub addr: Addr,
    pub pcksize: PckSize,
    pub extreg: ExtReg,
    pub status_bk: StatusBk,
    _reserved0: u8,
    pub ctrl_pipe: CtrlPipe,
    pub status_pipe: StatusPipe,
    _reserved1: u8,
}

impl BankDesc {
    fn new() -> Self {
        Self {
            addr: Addr::from(0),
            pcksize: PckSize::from(0),
            extreg: ExtReg::from(0),
            status_bk: StatusBk::from(0),
            _reserved0: 0,
            ctrl_pipe: CtrlPipe::from(0),
            status_pipe: StatusPipe::from(0),
            _reserved1: 0,
        }
    }
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn bank_desc_sizes() {
        assert_eq!(core::mem::size_of::<Addr>(), 4, "Addr register size.");
        assert_eq!(core::mem::size_of::<PckSize>(), 4, "PckSize register size.");
        assert_eq!(core::mem::size_of::<ExtReg>(), 2, "ExtReg register size.");
        assert_eq!(
            core::mem::size_of::<StatusBk>(),
            1,
            "StatusBk register size."
        );
        assert_eq!(
            core::mem::size_of::<CtrlPipe>(),
            2,
            "CtrlPipe register size."
        );
        assert_eq!(
            core::mem::size_of::<StatusPipe>(),
            1,
            "StatusPipe register size."
        );

        // addr at 0x00 for 4
        // pcksize at 0x04 for 4
        // extreg at 0x08 for 2
        // status_bk at 0x0a for 2
        // ctrl_pipe at 0x0c for 2
        // status_pipe at 0x0e for 1
        assert_eq!(
            core::mem::size_of::<BankDesc>(),
            16,
            "Bank descriptor size."
        );
    }

    #[test]
    fn bank_desc_offsets() {
        let bd = BankDesc::new();
        let base = &bd as *const _ as usize;

        assert_offset("Addr", &bd.addr, base, 0x00);
        assert_offset("PckSize", &bd.pcksize, base, 0x04);
        assert_offset("ExtReg", &bd.extreg, base, 0x08);
        assert_offset("StatusBk", &bd.status_bk, base, 0x0a);
        assert_offset("CtrlPipe", &bd.ctrl_pipe, base, 0x0c);
        assert_offset("StatusPipe", &bd.status_pipe, base, 0x0e);
    }

    #[test]
    fn pipe_desc_size() {
        assert_eq!(core::mem::size_of::<PipeDesc>(), 32);
    }

    #[test]
    fn pipe_desc_offsets() {
        let pd = PipeDesc::new();
        let base = &pd as *const _ as usize;

        assert_offset("Bank0", &pd.bank0, base, 0x00);
        assert_offset("Bank1", &pd.bank1, base, 0x10);
    }

    fn assert_offset<T>(name: &str, field: &T, base: usize, offset: usize) {
        let ptr = field as *const _ as usize;
        assert_eq!(ptr - base, offset, "{} register offset.", name);
    }
}