// SPDX-License-Identifier: LGPL-3.0-or-later
// See Notices.txt for copyright information

use cpu::{
    main_memory_and_io::{
        AddressRange, MemoryInterface, MemoryInterfaceConfig, MemoryOperationErrorKind,
        MemoryOperationFinish, MemoryOperationFinishKind, MemoryOperationKind,
        MemoryOperationStart, main_memory_and_io, simple_uart::SIMPLE_UART_TRANSMIT_OFFSET,
    },
    next_pc::FETCH_BLOCK_ID_WIDTH,
};
use fayalite::{prelude::*, sim::vcd::VcdWriterDecls, util::RcWriter};
use std::num::{NonZeroU64, NonZeroUsize, Wrapping};

const CLOCK_FREQUENCY: f64 = 2_000_000.0;

const fn half_period(frequency: f64) -> SimDuration {
    SimDuration::from_attos((SimDuration::from_secs(1).as_attos() as f64 / frequency / 2.0) as u128)
}

const CLOCK_HALF_PERIOD: SimDuration = half_period(CLOCK_FREQUENCY);

fn clock_input_properties() -> PhantomConst<peripherals::ClockInputProperties> {
    peripherals::ClockInput::new(CLOCK_FREQUENCY).properties
}

const SRAM_START: u64 = 0x1000;
const SRAM_SIZE: NonZeroU64 = NonZeroU64::new(0x30).unwrap();
const UART_START: u64 = 0x100000;
const UART_BAUD_RATE: f64 = 115200.0;
const UART_RECEIVER_QUEUE_SIZE: NonZeroUsize = NonZeroUsize::new(16).unwrap();
const SRAM_INITIAL_CONTENTS: &[u8; SRAM_SIZE.get() as usize] =
    b"This is a main memory and I/O test.             ";
const SRAM_WRITING_CONTENTS: &[u8; SRAM_SIZE.get() as usize] =
    b"Testing. Hello World. abcdefghijklmnopqrstuvwxyz";

const CONFIG: MemoryInterfaceConfig =
    MemoryInterfaceConfig::new(3, 8, FETCH_BLOCK_ID_WIDTH, AddressRange::Full);
const BUS_WIDTH_IN_BYTES: usize = CONFIG.bus_width_in_bytes();

fn config() -> PhantomConst<MemoryInterfaceConfig> {
    PhantomConst::new_sized(CONFIG)
}

#[hdl_module]
fn test_harness() {
    let config = config();

    #[hdl]
    let cd: ClockDomain = m.input();
    #[hdl]
    let memory_interface: MemoryInterface<PhantomConst<MemoryInterfaceConfig>> =
        m.input(MemoryInterface[config]);

    #[hdl]
    let dut = instance(main_memory_and_io(
        config,
        clock_input_properties(),
        AddressRange::Limited {
            start: Wrapping(SRAM_START),
            size: SRAM_SIZE,
        },
        SRAM_INITIAL_CONTENTS,
        UART_START,
        UART_BAUD_RATE,
        UART_RECEIVER_QUEUE_SIZE,
    ));
    connect(dut.cd, cd);
    connect(dut.memory_interface, memory_interface);
    connect(dut.uart.rx, dut.uart.tx); // loop back for testing
}

struct MainMemoryAndIoTester {
    sim: Simulation<test_harness>,
    next_op_id: u64,
}

impl MainMemoryAndIoTester {
    fn reset(&mut self) {
        let sim = &mut self.sim;
        sim.write(
            sim.io().memory_interface.start.data,
            sim.io().ty().memory_interface.start.data.HdlNone(),
        );
        sim.write(sim.io().memory_interface.finish.ready, false);
        sim.write_clock(sim.io().cd.clk, false);
        sim.write_reset(sim.io().cd.rst, true);
        sim.advance_time(CLOCK_HALF_PERIOD);
        sim.write_clock(sim.io().cd.clk, true);
        sim.advance_time(CLOCK_HALF_PERIOD);
        sim.write_clock(sim.io().cd.clk, false);
        sim.write_reset(sim.io().cd.rst, false);
    }
    fn wait(&mut self, cycles: u64) {
        let sim = &mut self.sim;
        sim.write(
            sim.io().memory_interface.start.data,
            sim.io().ty().memory_interface.start.data.HdlNone(),
        );
        sim.write(sim.io().memory_interface.finish.ready, false);
        for _ in 0..cycles {
            sim.advance_time(CLOCK_HALF_PERIOD);
            sim.write_clock(sim.io().cd.clk, true);
            sim.advance_time(CLOCK_HALF_PERIOD);
            sim.write_clock(sim.io().cd.clk, false);
        }
    }
    #[track_caller]
    #[hdl]
    fn try_rw_op(
        &mut self,
        addr: u64,
        write: Option<[u8; BUS_WIDTH_IN_BYTES]>,
        rw_mask: [bool; BUS_WIDTH_IN_BYTES],
        max_wait: u64,
    ) -> Result<[u8; BUS_WIDTH_IN_BYTES], SimValue<MemoryOperationErrorKind>> {
        let config = config();
        let sim = &mut self.sim;
        let op_id = self.next_op_id.cast_to(UInt[FETCH_BLOCK_ID_WIDTH]);
        self.next_op_id += 1;
        sim.write(
            sim.io().memory_interface.start.data,
            #[hdl(sim)]
            (sim.io().ty().memory_interface.start.data).HdlSome(match write {
                Some(write) =>
                {
                    #[hdl(sim)]
                    MemoryOperationStart::<_> {
                        kind: #[hdl(sim)]
                        MemoryOperationKind.Write(),
                        addr,
                        write_data: &write[..],
                        rw_mask: &rw_mask[..],
                        op_id,
                        config,
                    }
                }
                None =>
                {
                    #[hdl(sim)]
                    MemoryOperationStart::<_> {
                        kind: #[hdl(sim)]
                        MemoryOperationKind.Read(),
                        addr,
                        write_data: &[0u8; BUS_WIDTH_IN_BYTES][..],
                        rw_mask: &rw_mask[..],
                        op_id,
                        config,
                    }
                }
            }),
        );
        sim.write(sim.io().memory_interface.finish.ready, true);
        let mut cleared_start = false;
        for _ in 0..max_wait {
            sim.advance_time(CLOCK_HALF_PERIOD);
            let start_ready = sim.read_bool(sim.io().memory_interface.start.ready);
            let finish_data = sim.read(sim.io().memory_interface.finish.data);
            sim.write_clock(sim.io().cd.clk, true);
            sim.advance_time(CLOCK_HALF_PERIOD);
            sim.write_clock(sim.io().cd.clk, false);
            if start_ready && !cleared_start {
                cleared_start = true;
                sim.write(
                    sim.io().memory_interface.start.data,
                    sim.io().ty().memory_interface.start.data.HdlNone(),
                );
            }
            #[hdl(sim)]
            if let HdlSome(finish) = finish_data {
                #[hdl(sim)]
                let MemoryOperationFinish::<_> {
                    kind,
                    read_data,
                    config: _,
                } = finish;
                #[hdl(sim)]
                match kind {
                    MemoryOperationFinishKind::Success(_) => {
                        return Ok(std::array::from_fn(|i| read_data[i].as_int()));
                    }
                    MemoryOperationFinishKind::Error(e) => {
                        return Err(e);
                    }
                };
            }
        }
        panic!("waited too many cycles");
    }
    #[track_caller]
    fn try_read<const N: usize>(
        &mut self,
        addr: u64,
        max_wait: u64,
    ) -> Result<[u8; N], SimValue<MemoryOperationErrorKind>> {
        const { assert!(N <= BUS_WIDTH_IN_BYTES) }
        let bus_addr = addr - addr % BUS_WIDTH_IN_BYTES as u64;
        let mut rw_mask = [false; BUS_WIDTH_IN_BYTES];
        let start = (addr % BUS_WIDTH_IN_BYTES as u64) as usize;
        rw_mask[start..start + N].fill(true);
        self.try_rw_op(bus_addr, None, rw_mask, max_wait)
            .map(|v| std::array::from_fn(|i| v[i + start]))
    }
    #[track_caller]
    fn read<const N: usize>(&mut self, addr: u64, max_wait: u64) -> [u8; N] {
        self.try_read(addr, max_wait).expect("read returned error")
    }
    #[track_caller]
    fn try_write<const N: usize>(
        &mut self,
        addr: u64,
        data: [u8; N],
        max_wait: u64,
    ) -> Result<(), SimValue<MemoryOperationErrorKind>> {
        const { assert!(N <= BUS_WIDTH_IN_BYTES) }
        let bus_addr = addr - addr % BUS_WIDTH_IN_BYTES as u64;
        let mut rw_mask = [false; BUS_WIDTH_IN_BYTES];
        let mut write_data = [0u8; BUS_WIDTH_IN_BYTES];
        let start = (addr % BUS_WIDTH_IN_BYTES as u64) as usize;
        rw_mask[start..start + N].fill(true);
        write_data[start..start + N].copy_from_slice(&data);
        self.try_rw_op(bus_addr, Some(write_data), rw_mask, max_wait)?;
        Ok(())
    }
    #[track_caller]
    fn write<const N: usize>(&mut self, addr: u64, data: [u8; N], max_wait: u64) {
        self.try_write(addr, data, max_wait)
            .expect("read returned error")
    }
}

#[test]
#[hdl]
fn test_main_memory_and_io() {
    let _n = SourceLocation::normalize_files_for_tests();
    let mut sim = Simulation::new(test_harness());
    let writer = RcWriter::default();
    sim.add_trace_writer(VcdWriterDecls::new(writer.clone()));
    struct DumpVcdOnDrop {
        writer: Option<RcWriter>,
    }
    impl Drop for DumpVcdOnDrop {
        fn drop(&mut self) {
            if let Some(mut writer) = self.writer.take() {
                let vcd = String::from_utf8(writer.take()).unwrap();
                println!("####### VCD:\n{vcd}\n#######");
            }
        }
    }
    let mut writer = DumpVcdOnDrop {
        writer: Some(writer),
    };
    let mut tester = MainMemoryAndIoTester { sim, next_op_id: 0 };
    tester.reset();

    const UART_MESSAGE: &str = "Test.";
    for b in UART_MESSAGE.bytes() {
        tester.write(UART_START + SIMPLE_UART_TRANSMIT_OFFSET, [b], 200);
    }

    for (i, expected_chunk) in SRAM_INITIAL_CONTENTS
        .as_chunks::<BUS_WIDTH_IN_BYTES>()
        .0
        .iter()
        .enumerate()
    {
        let addr = (i as u64 * BUS_WIDTH_IN_BYTES as u64) + SRAM_START;
        let chunk = tester.read::<BUS_WIDTH_IN_BYTES>(addr, 10);
        assert_eq!(chunk, *expected_chunk);
    }

    for (i, chunk) in SRAM_WRITING_CONTENTS
        .as_chunks::<BUS_WIDTH_IN_BYTES>()
        .0
        .iter()
        .enumerate()
    {
        let addr = (i as u64 * BUS_WIDTH_IN_BYTES as u64) + SRAM_START;
        tester.write(addr, *chunk, 10);
    }

    for (i, expected_chunk) in SRAM_WRITING_CONTENTS
        .as_chunks::<BUS_WIDTH_IN_BYTES>()
        .0
        .iter()
        .enumerate()
    {
        let addr = (i as u64 * BUS_WIDTH_IN_BYTES as u64) + SRAM_START;
        let chunk = tester.read::<BUS_WIDTH_IN_BYTES>(addr, 10);
        assert_eq!(chunk, *expected_chunk);
    }

    tester.wait(400);

    for expected_byte in UART_MESSAGE.bytes() {
        let [byte] = tester.read(UART_START + SIMPLE_UART_TRANSMIT_OFFSET, 10);
        assert_eq!(byte, expected_byte);
    }

    let vcd = String::from_utf8(writer.writer.take().unwrap().take()).unwrap();
    println!("####### VCD:\n{vcd}\n#######");
    if vcd != include_str!("expected/main_memory_and_io.vcd") {
        panic!();
    }
}