use kaptan_mavlink::{
    MavlinkFrame,
    encoder::MavlinkEncoder,
    generated::kaptan::{FILE_DATA_DATA, MavMessage},
};
use kaptan_tx::{FrameClass, PendingFrame};
use usb_device::bus::PollResult;

#[path = "../src/f722/usb.rs"]
mod usb;
use usb::test_bus as usb_bus;
use usb_bus::with_link;
#[path = "../src/f722/serial.rs"]
mod serial;
use serial::{IrqTransmit, UartIo};

#[derive(Default)]
struct MockUart {
    bytes: Vec<u8>,
    ready: bool,
    complete: bool,
    ready_irq: bool,
    complete_irq: bool,
}
impl UartIo for MockUart {
    fn tx_ready(&mut self) -> bool {
        self.ready
    }
    fn tx_complete(&mut self) -> bool {
        self.complete
    }
    fn write_byte(&mut self, byte: u8) {
        self.bytes.push(byte);
        self.complete = false;
    }
    fn interrupts(&mut self, ready: bool, complete: bool) {
        self.ready_irq = ready;
        self.complete_irq = complete;
    }
}
fn flush_uart(tx: &mut IrqTransmit, io: &mut MockUart) {
    io.ready = true;
    for _ in 0..2048 {
        if !io.ready_irq && !io.complete_irq {
            return;
        }
        let previous = io.bytes.len();
        io.complete = true;
        tx.on_interrupt(io);
        assert!(
            io.bytes.len() <= previous + 1,
            "IRQ work is bounded to one byte"
        );
    }
    panic!("UART did not drain its bounded queue");
}

#[test]
fn moving_irq_queue_and_waiting_for_tc_preserves_exact_frame_bytes() {
    let mut io = MockUart {
        ready: true,
        ..MockUart::default()
    };
    let mut tx = IrqTransmit::new();
    let mut encoder = MavlinkEncoder::new(1, 1);
    let frame = encoder.encode(&kaptan_mavlink::generated::kaptan::HEARTBEAT_DATA::default());
    assert!(tx.send(
        &mut io,
        PendingFrame {
            class: FrameClass::Vital,
            frame
        }
    ));
    assert!(!tx.on_interrupt(&mut io));
    let mut moved = tx;
    for _ in 1..frame.raw_bytes().len() {
        assert!(!moved.on_interrupt(&mut io));
    }
    assert_eq!(io.bytes, frame.raw_bytes());
    assert!(!io.ready_irq && io.complete_irq);
    assert!(!moved.on_interrupt(&mut io));
    io.complete = true;
    assert!(moved.on_interrupt(&mut io));
    assert!(!io.ready_irq && !io.complete_irq);
    assert!(!moved.on_interrupt(&mut io));
}

#[test]
fn usb_and_two_uarts_have_independent_sequences_queues_and_completion() {
    use kaptan_data_channel::{DataChannel, DataChannelRates, TelemetrySamples};
    use kaptan_time::Timestamp;
    with_link(|host, usb| {
        host.configure(usb, true);
        host.dtr(usb, true);
        let rates = DataChannelRates {
            imu_hz: 20,
            heartbeat_hz: 2,
            range_hz: None,
            optical_flow_hz: None,
        };
        let mut channels = core::array::from_fn::<_, 3, _>(|_| {
            DataChannel::<100>::new(1, 1, rates, 100_000).unwrap()
        });
        let mut uart1 = IrqTransmit::new();
        let mut uart2 = IrqTransmit::new();
        let mut slow = MockUart::default();
        let mut fast = MockUart::default();
        let samples = TelemetrySamples::default();
        let mut usb_bytes = Vec::new();
        let mut expected = Vec::new();
        let mut dropped = false;
        for tick in 0..400 {
            let now = Timestamp::from_microseconds(tick * 10_000);
            let frames = channels
                .each_mut()
                .map(|channel| channel.on_service_tick(&samples, now));
            if let [Some(a), Some(b), Some(c)] = frames {
                assert_eq!(a.frame.raw_bytes(), b.frame.raw_bytes());
                assert_eq!(b.frame.raw_bytes(), c.frame.raw_bytes());
                expected.extend_from_slice(c.frame.raw_bytes());
                assert!(usb.send(a));
                dropped |= !uart1.send(&mut slow, b);
                assert!(uart2.send(&mut fast, c));
            }
            flush_uart(&mut uart2, &mut fast);
            usb.poll();
            usb_bytes.extend(host.collect_packets(usb).into_iter().flatten());
        }
        assert!(
            dropped,
            "blocked UART must exhaust only its own bounded queue"
        );
        assert!(slow.bytes.is_empty());
        assert_eq!(usb_bytes, expected);
        assert_eq!(fast.bytes, expected);
        let fast_before = fast.bytes.clone();
        flush_uart(&mut uart1, &mut slow);
        assert!(!slow.bytes.is_empty());
        assert_eq!(
            fast.bytes, fast_before,
            "UART1 completions cannot advance UART2"
        );
    });
}

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.events_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);
    });
}

#[test]
fn bus_fault_closes_link_and_discards_pending_data() {
    for fail in [
        usb_bus::Host::fail_poll,
        usb_bus::Host::fail_read,
        usb_bus::Host::fail_write,
    ] {
        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());
            fail(host);
            link.with_rx_interrupt_gate(
                |link| assert!(!link.poll()),
                |_| panic!("interrupt register access after terminal fault"),
            );
            assert!(!link.is_open());
            assert!(!link.send(pending(frame)));
            assert!(link.take_received().is_none());
            assert!(!link.poll());
            assert!(!host.nvic_enabled());
        });
    }
}