use std::{
    collections::VecDeque,
    sync::{Arc, Mutex},
};

use kaptan_mavlink::{
    MavlinkFrame,
    encoder::MavlinkEncoder,
    generated::kaptan::{FILE_DATA_DATA, MavMessage},
};
use kaptan_tx::{FrameClass, PendingFrame};
use usb_device::{
    UsbDirection, UsbError,
    bus::{PollResult, UsbBus, UsbBusAllocator},
    device::{UsbDeviceBuilder, UsbVidPid},
    endpoint::{EndpointAddress, EndpointType},
};
use usbd_serial::CdcAcmClass;

#[path = "../src/f722/usb.rs"]
mod usb;
use usb::UsbMavlink;

// One pending packet per endpoint, with host-controlled completion. Device and
// CDC control requests run through the real usb-device/usbd-serial code.
#[derive(Default)]
struct BusState {
    events: VecDeque<PollResult>,
    input: [Option<Vec<u8>>; 4],
    output: [Option<Vec<u8>>; 4],
    sizes: [[u16; 4]; 2],
    stalled: [[bool; 4]; 2],
    bulk_in: usize,
    bulk_out: usize,
}

#[derive(Clone, Default)]
struct Host(Arc<Mutex<BusState>>);

struct PacketBus(Arc<Mutex<BusState>>);

impl UsbBus for PacketBus {
    fn alloc_ep(
        &mut self,
        direction: UsbDirection,
        address: Option<EndpointAddress>,
        kind: EndpointType,
        size: u16,
        _: u8,
    ) -> Result<EndpointAddress, UsbError> {
        let mut state = self.0.lock().unwrap();
        let sizes = &mut state.sizes[usize::from(direction == UsbDirection::In)];
        let index = match address {
            Some(address) => address.index(),
            None => (1..sizes.len())
                .find(|index| sizes[*index] == 0)
                .ok_or(UsbError::EndpointOverflow)?,
        };
        if index >= sizes.len() || sizes[index] != 0 {
            return Err(UsbError::InvalidEndpoint);
        }
        sizes[index] = size;
        if kind == EndpointType::Bulk {
            match direction {
                UsbDirection::In => state.bulk_in = index,
                UsbDirection::Out => state.bulk_out = index,
            }
        }
        Ok(EndpointAddress::from_parts(index, direction))
    }

    fn enable(&mut self) {}

    fn reset(&self) {
        let mut state = self.0.lock().unwrap();
        state.input.fill(None);
        state.output.fill(None);
        state.events.clear();
    }

    fn set_device_address(&self, _: u8) {}

    fn write(&self, address: EndpointAddress, bytes: &[u8]) -> Result<usize, UsbError> {
        let mut state = self.0.lock().unwrap();
        if bytes.len() > usize::from(state.sizes[1][address.index()]) {
            return Err(UsbError::BufferOverflow);
        }
        let packet = &mut state.input[address.index()];
        if packet.is_some() {
            return Err(UsbError::WouldBlock);
        }
        *packet = Some(bytes.to_vec());
        Ok(bytes.len())
    }

    fn read(&self, address: EndpointAddress, bytes: &mut [u8]) -> Result<usize, UsbError> {
        let mut state = self.0.lock().unwrap();
        let packet = state.output[address.index()]
            .as_ref()
            .ok_or(UsbError::WouldBlock)?;
        if packet.len() > bytes.len() {
            return Err(UsbError::BufferOverflow);
        }
        let packet = state.output[address.index()].take().unwrap();
        bytes[..packet.len()].copy_from_slice(&packet);
        Ok(packet.len())
    }

    fn set_stalled(&self, address: EndpointAddress, stalled: bool) {
        self.0.lock().unwrap().stalled[usize::from(address.is_in())][address.index()] = stalled;
    }

    fn is_stalled(&self, address: EndpointAddress) -> bool {
        self.0.lock().unwrap().stalled[usize::from(address.is_in())][address.index()]
    }

    fn suspend(&self) {}
    fn resume(&self) {}

    fn poll(&self) -> PollResult {
        self.0
            .lock()
            .unwrap()
            .events
            .pop_front()
            .unwrap_or(PollResult::None)
    }
}

impl Host {
    fn event(&self, event: PollResult) {
        self.0.lock().unwrap().events.push_back(event);
    }

    fn control(&self, link: &mut UsbMavlink<'_, PacketBus>, request: [u8; 8]) {
        self.0.lock().unwrap().output[0] = Some(request.to_vec());
        self.event(PollResult::Data {
            ep_out: 0,
            ep_in_complete: 0,
            ep_setup: 1,
        });
        link.poll();
        assert_eq!(self.take_input(0), Some(Vec::new()));
        link.poll();
    }

    fn configure(&self, link: &mut UsbMavlink<'_, PacketBus>, configured: bool) {
        self.control(link, [0, 9, u8::from(configured), 0, 0, 0, 0, 0]);
    }

    fn dtr(&self, link: &mut UsbMavlink<'_, PacketBus>, open: bool) {
        self.control(link, [0x21, 0x22, u8::from(open), 0, 0, 0, 0, 0]);
    }

    fn write_packet(&self, bytes: &[u8]) {
        assert!(bytes.len() <= 64);
        let mut state = self.0.lock().unwrap();
        let endpoint = state.bulk_out;
        assert!(state.output[endpoint].is_none());
        state.output[endpoint] = Some(bytes.to_vec());
        state.events.push_back(PollResult::Data {
            ep_out: 1 << endpoint,
            ep_in_complete: 0,
            ep_setup: 0,
        });
    }

    fn take_input(&self, endpoint: usize) -> Option<Vec<u8>> {
        let mut state = self.0.lock().unwrap();
        let packet = state.input[endpoint].take()?;
        state.events.push_back(PollResult::Data {
            ep_out: 0,
            ep_in_complete: 1 << endpoint,
            ep_setup: 0,
        });
        Some(packet)
    }

    fn collect_packets(&self, link: &mut UsbMavlink<'_, PacketBus>) -> Vec<Vec<u8>> {
        let endpoint = self.0.lock().unwrap().bulk_in;
        let mut packets = Vec::new();
        for _ in 0..16 {
            let Some(packet) = self.take_input(endpoint) else {
                return packets;
            };
            packets.push(packet);
            link.poll();
        }
        panic!("transmission did not finish");
    }
}

fn with_link(test: impl FnOnce(&Host, &mut UsbMavlink<'_, PacketBus>)) {
    let host = Host::default();
    let allocator = UsbBusAllocator::new(PacketBus(host.0.clone()));
    let class = CdcAcmClass::new(&allocator, 64);
    let device = UsbDeviceBuilder::new(&allocator, UsbVidPid(0x1209, 1))
        .max_packet_size_0(64)
        .unwrap()
        .build();
    let mut link = UsbMavlink::new(device, class);
    host.event(PollResult::Reset);
    link.poll();
    host.configure(&mut link, true);
    host.dtr(&mut link, true);
    assert!(link.is_open());
    test(&host, &mut link);
}

fn pending(frame: MavlinkFrame) -> PendingFrame<MavlinkFrame> {
    PendingFrame {
        class: FrameClass::Critical,
        frame,
    }
}

fn file_frame(length: usize) -> MavlinkFrame {
    // FILE_DATA has nine scalar bytes before its trailing data array. Setting
    // the final data byte controls MAVLink 2's trailing-zero truncation.
    let mut data = FILE_DATA_DATA::DEFAULT;
    data.data[..length - 12 - 9].fill(0x5a);
    let frame = MavlinkEncoder::new(1, 1).encode(&data);
    assert_eq!(frame.raw_bytes().len(), length);
    frame
}

#[test]
fn receives_split_frame_followed_by_coalesced_frames() {
    with_link(|host, link| {
        // HEARTBEAT ID 51027, source 255:1, sequence 0, empty payload,
        // CRC-16/MCRF4XX with CRC_EXTRA 161. Integrity admission precedes payload
        // validation, so these 12-byte frames must also survive USB batching.
        let mut minimum = MavlinkFrame::new();
        minimum.as_mut_slice()[..12]
            .copy_from_slice(&[0xfd, 0, 0, 0, 0, 255, 1, 0x53, 0xc7, 0, 0x2e, 0xa8]);
        assert!(minimum.has_valid_crc::<MavMessage>());
        let mut frames = [minimum; 6];
        frames[0] = MavlinkEncoder::new(255, 1).heartbeat();
        let split = frames[0].raw_bytes().len() - 1;
        host.write_packet(&frames[0].raw_bytes()[..split]);
        assert!(link.poll());
        assert!(link.take_received().is_none());

        // One carried-over CRC byte and five complete frames: six completions
        // in 61 bytes, exceeding the former five-frame receive queue.
        let mut packet = frames[0].raw_bytes()[split..].to_vec();
        for frame in &frames[1..] {
            packet.extend_from_slice(frame.raw_bytes());
        }
        host.write_packet(&packet);
        assert!(link.poll());
        for frame in &frames {
            assert_eq!(link.take_received().unwrap().raw(), frame);
        }
        assert!(link.take_received().is_none());
    });
}

#[test]
fn drains_endpoint_backlog_after_its_poll_event_was_consumed() {
    with_link(|host, link| {
        let mut encoder = MavlinkEncoder::new(255, 1);
        let first = encoder.heartbeat();
        let second = encoder.heartbeat();
        host.write_packet(first.raw_bytes());
        assert!(link.poll());
        host.write_packet(second.raw_bytes());
        assert!(link.poll());
        assert_eq!(link.take_received().unwrap().raw(), &first);
        assert!(link.take_received().is_none());

        assert!(host.0.lock().unwrap().events.is_empty());
        assert!(link.poll());
        assert_eq!(link.take_received().unwrap().raw(), &second);
    });
}

#[test]
fn endpoint_backpressure_neither_repeats_nor_interleaves_frame_bytes() {
    with_link(|host, link| {
        let frame = file_frame(130);
        let next = MavlinkEncoder::new(1, 1).heartbeat();
        assert!(link.send(pending(frame)));
        assert!(link.send(pending(next)));
        for _ in 0..4 {
            link.poll();
        }
        let packets = host.collect_packets(link);
        assert_eq!(
            packets.iter().map(Vec::len).collect::<Vec<_>>(),
            [64, 64, 2, 13]
        );
        assert_eq!(
            packets.concat(),
            [frame.raw_bytes(), next.raw_bytes()].concat()
        );
    });
}

#[test]
fn full_size_final_packet_is_followed_by_one_zlp() {
    for length in [64, 128] {
        with_link(|host, link| {
            let frame = file_frame(length);
            let next = MavlinkEncoder::new(1, 1).heartbeat();
            assert!(link.send(pending(frame)));
            assert!(link.send(pending(next)));
            let packets = host.collect_packets(link);
            let mut lengths = vec![64; length / 64];
            lengths.extend([0, 13]);
            assert_eq!(packets.iter().map(Vec::len).collect::<Vec<_>>(), lengths);
            assert_eq!(
                packets.concat(),
                [frame.raw_bytes(), next.raw_bytes()].concat()
            );
        });
    }
}

#[test]
fn bounded_admission_preserves_reserved_frame_and_replaces_only_pending_frames() {
    with_link(|host, link| {
        let mut encoder = MavlinkEncoder::new(1, 1);
        let first = encoder.heartbeat();
        let reserved = encoder.heartbeat();
        assert!(link.send(pending(first)));
        // The endpoint still owns the first packet; this start must retry.
        assert!(link.send(pending(reserved)));
        for _ in 0..4 {
            assert!(link.send(PendingFrame {
                class: FrameClass::BestEffort,
                frame: encoder.heartbeat(),
            }));
        }
        let critical = std::array::from_fn::<_, 4, _>(|_| encoder.heartbeat());
        for frame in critical {
            assert!(link.send(pending(frame)));
        }
        assert!(!link.send(pending(encoder.heartbeat())));
        let vital = encoder.heartbeat();
        assert!(link.send(PendingFrame {
            class: FrameClass::Vital,
            frame: vital,
        }));

        let packets = host.collect_packets(link);
        assert_eq!(packets.len(), 6);
        assert_eq!(packets[0], first.raw_bytes());
        assert_eq!(packets[1], reserved.raw_bytes());
        assert_eq!(packets[2], vital.raw_bytes());
        let mut pending = critical.map(|frame| frame.raw_bytes().to_vec()).to_vec();
        for packet in &packets[3..] {
            let index = pending.iter().position(|frame| frame == packet).unwrap();
            pending.remove(index);
        }
    });
}

#[test]
fn dtr_suspend_and_deconfiguration_preserve_unfinished_transmission() {
    enum Pause {
        Dtr,
        Suspend,
        Deconfigure,
    }
    for pause in [Pause::Dtr, Pause::Suspend, Pause::Deconfigure] {
        with_link(|host, link| {
            let frame = file_frame(130);
            assert!(link.send(pending(frame)));
            match pause {
                Pause::Dtr => host.dtr(link, false),
                Pause::Suspend => {
                    host.event(PollResult::Suspend);
                    link.poll();
                }
                Pause::Deconfigure => host.configure(link, false),
            }
            assert!(!link.is_open());
            assert!(!link.send(pending(MavlinkEncoder::new(1, 1).heartbeat())));
            match pause {
                Pause::Dtr => host.dtr(link, true),
                Pause::Suspend => {
                    host.event(PollResult::Resume);
                    link.poll();
                }
                Pause::Deconfigure => host.configure(link, true),
            }
            assert!(link.is_open());
            assert_eq!(host.collect_packets(link).concat(), frame.raw_bytes());
        });
    }
}

#[test]
fn dtr_close_discards_partial_receive_frame() {
    with_link(|host, link| {
        let mut encoder = MavlinkEncoder::new(255, 1);
        let interrupted = encoder.heartbeat();
        let next = encoder.heartbeat();
        host.write_packet(&interrupted.raw_bytes()[..8]);
        link.poll();
        assert!(link.take_received().is_none());
        host.dtr(link, false);
        host.dtr(link, true);
        host.write_packet(next.raw_bytes());
        link.poll();
        assert_eq!(link.take_received().unwrap().raw(), &next);
        assert!(link.take_received().is_none());
    });
}

#[test]
fn bus_reset_discards_transmit_and_receive_state() {
    with_link(|host, link| {
        let frame = file_frame(130);
        assert!(link.send(pending(frame)));
        host.write_packet(&frame.raw_bytes()[..8]);
        link.poll();
        assert!(link.take_received().is_none());
        host.event(PollResult::Reset);
        link.poll();
        assert!(!link.is_open());
        host.configure(link, true);
        host.dtr(link, true);
        assert!(host.collect_packets(link).is_empty());

        let next = MavlinkEncoder::new(1, 1).heartbeat();
        assert!(link.send(pending(next)));
        assert_eq!(host.collect_packets(link).concat(), next.raw_bytes());
        host.write_packet(next.raw_bytes());
        link.poll();
        assert_eq!(link.take_received().unwrap().raw(), &next);
    });
}