use std::collections::HashMap;
use std::vec::Vec;

const PAGE_BITS: usize = 12;
const PAGE_SIZE: usize = 2 << PAGE_BITS;

const TABLE_BITS: i64 = 9;
const TABLE_SIZE: i64 = 2 << TABLE_BITS;

const SV39_PAGE_BITS: i64 = 9 + 9 + 9;

// Actively used bits in our limited implementation:
const BIT_V: i64 = 1;
const BIT_R: i64 = 2;
const BIT_W: i64 = 4;
const BIT_X: i64 = 8;

// (Not used) - user mode accessible
const BIT_U: i64 = 16;

// (Not used) - global mapping
const BIT_G: i64 = 32;

// (Not presently used) accessed -- set if we've touched this since last time A cleared
const BIT_A: i64 = 64;

// (Not presently used) dirty -- written since last time D cleared
// Note each core/hart has its own dirty state for local JIT
const BIT_D: i64 = 128;

// (Not used) reserved for supervisor
const BITS_RSW_LO: i64 = 256;
const BITS_RSW_HI: i64 = 512;
const BITS_RSW: i64 = BITS_RSW_LO | BITS_RSW_HI;


#[repr(C)]
#[derive(Clone, Copy)]
struct PageTableEntry {
    flags: i64
}

impl PageTableEntry {
    fn new(flags: i64) -> Self {
        return Self {
            flags
        }
    }

    fn as_i64(&self) -> i64 {
        return self.flags;
    }

    fn is_pointer(&self) -> bool {
        return self.flags & (BIT_V | BIT_R | BIT_W | BIT_X) == BIT_V;
    }

    fn is_valid(&self) -> bool {
        return self.flags & BIT_V == BIT_V;
    }

    fn is_readable(&self) -> bool {
        return self.flags & (BIT_V | BIT_R) == (BIT_V | BIT_R);
    }

    fn is_writable(&self) -> bool {
        return self.flags & (BIT_V | BIT_R | BIT_W) == (BIT_V | BIT_R | BIT_W);
    }

    fn is_executable(&self) -> bool {
        return self.flags & (BIT_V | BIT_R | BIT_X) == (BIT_V | BIT_R | BIT_X);
    }

    fn get_physical_page(&self) -> i64 {
        return (self.flags << 10) >> 20;
    }

    fn get_physical(&self) -> i64 {
        return self.get_physical_page() << PAGE_BITS;
    }

}

type ExecutorFunc = fn(i64, &mut CoreState, &mut MachineState) -> i64;

#[repr(C)]
struct MachineState {
    memory: Vec<u8>,
    pages: Vec<PageTableEntry>,
}

/**
 * Note that physical memory accessors can traverse page boundaries;
 * we lay out linear memory from 0 to +4 gigabytes and will allocate
 * as many page table entries as are needed to cover RAM. These eat
 * up an extra 1 byte per 4 KiB of address space used (1 MiB per 4 GiB),
 * initially allocating enough for all physical memory allocated.
 *
 * This will be relatively space-inefficient for sparse address spaces
 * in the range of several gigabytes and more but requires only one
 * table lookup per load/store, or two if spanning pages.
 *
 * All threads in the process's machine will have access to the same
 * page tables, even if running on different threads.
 */
impl MachineState {
    fn new_with_state(memory: Vec<u8>, pages: Vec<PageTableEntry>) -> Self {
        if ((memory.len() >> PAGE_BITS) << PAGE_BITS) != memory.len() {
            panic!("memory size must be a multiple of 4096 bytes");
        }
        if ((memory.len() >> PAGE_BITS)) != pages.len() {
            panic!("page data is wrong length for memory size");
        }
        return Self {
            memory,
            pages
        };
    }

    fn new(memory_size: usize) -> Self {
        let memory = vec![0u8; memory_size];
        let pages = vec![PageTableEntry::new(0); memory_size >> PAGE_BITS];
        return Self::new_with_state(memory, pages);
    }

    fn lb_physical(&self, address: i64) -> i64 {
        if (address as u64) < (self.memory.len() as u64) {
            return self.memory[address as usize] as i8 as i64;
        }
        panic!("@fixme: trap on out of bounds physical read");
    }

    fn lbu_physical(&self, address: i64) -> i64 {
        if (address as u64) < (self.memory.len() as u64) {
            return self.memory[address as usize] as u8 as i64;
        }
        panic!("@fixme: trap on out of bounds physical read");
    }

    fn lh_physical(&self, address: i64) -> i64 {
        if (address as u64) < (self.memory.len() as u64 - 1) {
            return (
                ( self.memory[address as usize    ] as u16)       |
                ((self.memory[address as usize + 1] as u16) << 8)
            ) as i16 as i64;
        }
        panic!("@fixme: trap on out of bounds physical read");
    }

    fn lhu_physical(&self, address: i64) -> i64 {
        if (address as u64) < (self.memory.len() as u64 - 1) {
            return (
                ( self.memory[address as usize   ] as u16)       |
                ((self.memory[address as usize + 1] as u16) << 8)
            ) as i64;
        }
        panic!("@fixme: trap on out of bounds physical read");
    }

    fn lw_physical(&self, address: i64) -> i64 {
        if (address as u64) < (self.memory.len() as u64 - 3) {
            return (
                ( self.memory[address as usize    ] as u32)        |
                ((self.memory[address as usize + 1] as u32) <<  8) |
                ((self.memory[address as usize + 2] as u32) << 16) |
                ((self.memory[address as usize + 3] as u32) << 24)
            ) as i32 as i64;
        }
        panic!("@fixme: trap on out of bounds physical read");
    }

    fn lwu_physical(&self, address: i64) -> i64 {
        if (address as u64) < (self.memory.len() as u64 - 3) {
            return (
                ( self.memory[address as usize    ] as u32)        |
                ((self.memory[address as usize + 1] as u32) <<  8) |
                ((self.memory[address as usize + 2] as u32) << 16) |
                ((self.memory[address as usize + 3] as u32) << 24)
            ) as u32 as i64;
        }
        panic!("@fixme: trap on out of bounds physical read");
    }

    fn ld_physical(&self, address: i64) -> i64 {
        if (address as u64) < (self.memory.len() as u64 - 7) {
            return (
                ( self.memory[address as usize   ] as u64)        |
                ((self.memory[address as usize + 1] as u64) <<  8) |
                ((self.memory[address as usize + 2] as u64) << 16) |
                ((self.memory[address as usize + 3] as u64) << 24) |
                ((self.memory[address as usize + 4] as u64) << 32) |
                ((self.memory[address as usize + 5] as u64) << 40) |
                ((self.memory[address as usize + 6] as u64) << 48) |
                ((self.memory[address as usize + 7] as u64) << 56)
            ) as i64;
        }
        panic!("@fixme: trap on out of bounds physical read");
    }

    fn sb_physical(&mut self, address: i64, value: i64) {
        self.memory[address as usize] = value as u8;
    }

    fn sh_physical(&mut self, address: i64, value: i64) {
        self.memory[address as usize]     =  value       as u8;
        self.memory[address as usize + 1] = (value >> 8) as u8;
    }

    fn sw_physical(&mut self, address: i64, value: i64) {
        self.memory[address as usize]     =  value        as u8;
        self.memory[address as usize + 1] = (value >>  8) as u8;
        self.memory[address as usize + 2] = (value >> 16) as u8;
        self.memory[address as usize + 3] = (value >> 24) as u8;
    }

    fn sd_physical(&mut self, address: usize, value: i64) {
        self.memory[address as usize]     =  value        as u8;
        self.memory[address as usize + 1] = (value >>  8) as u8;
        self.memory[address as usize + 2] = (value >> 16) as u8;
        self.memory[address as usize + 3] = (value >> 24) as u8;
        self.memory[address as usize + 4] = (value >> 32) as u8;
        self.memory[address as usize + 5] = (value >> 40) as u8;
        self.memory[address as usize + 6] = (value >> 48) as u8;
        self.memory[address as usize + 7] = (value >> 56) as u8;
    }

    fn store_buffer_physical(&mut self, address: usize, bytes: &[u8]) {
        self.memory[address..address + bytes.len()].copy_from_slice(bytes);
    }
}

#[repr(C)]
struct CoreState {
    // Integer registers
    x: [i64; 32],
    satp: i64,

    // Do we need pc? we're passing it around as active state

    // Floating point registers
    // f32 values get... NaN-boxed into f64 values? wild
    // probably most efficiently treated like a union so can
    // do f32-native loads and stores and also update the top
    // bytes separately
    f: [f64; 32],

    // * fflags, accrued exceptions: bits 0-4
    //   * nx: bit 0
    //   * uf: bit 1
    //   * of: bit 2
    //   * dz: bit 3
    //   * nv: bit 4
    // * frm, rounding mode: bits 5-7
    // * reserved: bits 8-31
    fcsr: i32,

    // When we add AOT and JIT, compiled functions will be
    // referenced in this lookup from PC to function pointer.
    // Because function references are linked separately on
    // each thread in WebAssembly, this has to live in each
    // core's state separately.
    executors: HashMap<i64, ExecutorFunc>,

    // Local dirty flags for JIT pages.
    // When we get a fence.i instructrion, look for all dirty
    // pages and invalidate any functions including them
    // Takes up to 1 byte per 4 KiB (1 MiB per 4 GiB) per thread.
    // Could be made more compact if only 1 bit is needed.
    dirty: Vec<u8>
}

impl CoreState {
    fn new(machine: &MachineState) -> Self {
        return Self {
            x: [
                0, 0, 0, 0, 0, 0, 0, 0,
                0, 0, 0, 0, 0, 0, 0, 0,
                0, 0, 0, 0, 0, 0, 0, 0,
                0, 0, 0, 0, 0, 0, 0, 0
            ],
            satp: 0,
            f: [
                0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
                0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
                0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
                0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0
            ],
            fcsr: 0,
            executors: HashMap::new(),
            dirty: vec![0u8; machine.memory.len()]
        }
    }

    fn trap(&mut self, machine: &mut MachineState, pc: i64) -> i64 {
        panic!("@fixme implement traps");
    }

    fn dereference(&mut self, machine: &mut MachineState, address: i64) -> (PageTableEntry, usize) {
        let high = ((address as u64) >> 39) as usize;
        let vpn2 = (((address as u64) << (64 - 39)) >> (64 - 9)) as usize;
        let vpn1 = (((address as u64) << (64 - 39 + 9)) >> (64 - 9 - 9)) as usize;
        let vpn0 = (((address as u64) << (64 - 39 + 9 + 9)) >> (64 - 9 - 9)) as usize;

        if high == 0 {
            let ppn = self.satp << (64 - 44) >> (64 - 44);
            let pte2 = PageTableEntry::new(
                machine.ld_physical(ppn + (vpn2 << 3))
            );
            if pte2.is_pointer() {
                let pte1 = PageTableEntry::new(
                    machine.ld_physical(pte2.get_physical() + (vpn1 << 3))
                );
                if pte1.is_pointer() {
                    let pte0 = PageTableEntry::new(machine.ld_physical(pte1.get_physical() + (vpn0 << 3)));
                    if !pte0.is_pointer() {
                        // 4 KiB pages
                        let shift = 64 - 12;
                        return (pte0, (raw << shift >> shift) as usize);
                    }
                    panic!("@fixme: trap on pointer node when we're out of levels");
                }
                // 2 MiB megapages
                let shift = 64 - 12 - 9;
                return (pte1, (raw << shift >> shift) as usize);
            }
            // 1 GiB gigapages
            let shift = 64 - 12 - 9 - 9;
            return (pte2, (raw << shift >> shift) as usize);
        }
        panic!("@fixme: handle trap for kernel address space or invalid SV39 address");
    }

    fn lb(&mut self, machine: &mut MachineState, pc: i64, rd: usize, rs1: usize, imm: i32) -> i64 {
        //...
        panic!("@fixme");
    }

    fn lbu(&self, machine: &MachineState, address: i64) -> i64 {
        let raw_page = (address as u64) >> PAGE_BITS;
        if raw_page < machine.pages.len() as u64 {
            let addr = address as usize;
            let page = raw_page as usize;
            let entry = machine.get_page_table_entry(page);
            if entry.is_readable() {
                return machine.lbu_physical(addr);
            }
        }
        panic!("@fixme: trap on read failure on page");
    }

    fn lh(&self, machine: &MachineState, address: i64) -> i64 {
        let raw_page = (address as u64) >> PAGE_BITS;
        if raw_page < machine.pages.len() as u64 {
            let addr = address as usize;
            let page = raw_page as usize;
            let entry = machine.get_page_table_entry(page);
            if entry.is_readable() {
                let offset = addr & (PAGE_SIZE - 1);
                if offset < (PAGE_SIZE - 2) || machine.get_page_table_entry(page + 1).is_readable() {
                    return machine.lh_physical(addr);
                } else {
                    panic!("@fixme: trap on read failure on page + 1");
                }
            }
        }
        panic!("@fixme: trap on read failure on page");
    }

    fn lhu(&self, machine: &MachineState, address: i64) -> i64 {
        let raw_page = (address as u64) >> PAGE_BITS;
        if raw_page < machine.pages.len() as u64 {
            let addr = address as usize;
            let page = raw_page as usize;
            let entry = machine.get_page_table_entry(page);
            if entry.is_readable() {
                let offset = addr & (PAGE_SIZE - 1);
                if offset < (PAGE_SIZE - 2) || machine.get_page_table_entry(page + 1).is_readable() {
                    return machine.lhu_physical(addr);
                } else {
                    panic!("@fixme: trap on read failure on page + 1");
                }
            }
        }
        panic!("@fixme: trap on read failure on page");
    }

    fn lw(&self, machine: &MachineState, address: i64) -> i64 {
        let raw_page = (address as u64) >> PAGE_BITS;
        if raw_page < machine.pages.len() as u64 {
            let addr = address as usize;
            let page = raw_page as usize;
            let entry = machine.get_page_table_entry(page);
            if entry.is_readable() {
                let offset = addr & (PAGE_SIZE - 1);
                if offset < (PAGE_SIZE - 4) || machine.get_page_table_entry(page + 1).is_readable() {
                    return machine.lw_physical(addr);
                } else {
                    panic!("@fixme: trap on read failure on page + 1");
                }
            }
        }
        panic!("@fixme: trap on read failure on page");
    }

    fn lwu(&self, machine: &MachineState, address: i64) -> i64 {
        let raw_page = (address as u64) >> PAGE_BITS;
        if raw_page < machine.pages.len() as u64 {
            let addr = address as usize;
            let page = raw_page as usize;
            let entry = machine.get_page_table_entry(page);
            if entry.is_readable() {
                let offset = addr & (PAGE_SIZE - 1);
                if offset < (PAGE_SIZE - 4) || machine.get_page_table_entry(page + 1).is_readable() {
                    return machine.lwu_physical(addr);
                } else {
                    panic!("@fixme: trap on read failure on page + 1");
                }
            }
        }
        panic!("@fixme: trap on read failure on page");
    }

    fn ld(&self, machine: &MachineState, address: i64) -> i64 {
        let raw_page = (address as u64) >> PAGE_BITS;
        if raw_page < machine.pages.len() as u64 {
            let addr = address as usize;
            let page = raw_page as usize;
            let entry = machine.get_page_table_entry(page);
            if entry.is_readable() {
                let offset = addr & (PAGE_SIZE - 1);
                if offset < (PAGE_SIZE - 8) || machine.get_page_table_entry(page + 1).is_readable() {
                    return machine.ld_physical(addr);
                } else {
                    panic!("@fixme: trap on read failure on page + 1");
                }
            }
        }
        panic!("@fixme: trap on read failure on page");
    }
}


// R-type
// * opcode: bits 0-6
// * rd: bits 7-11
// * funct3: bits 12-14
// * rs1: bits 15-19
// * rs2: bits 20-24
// * func7: bits 25-31

// I-type
// * opcode: bits 0-6
// * rd: bits 7-11
// * funct3: bits 12-14
// * rs1: bits 15-19
// * imm[0:11]: bits 20-31

// S-type
// * opcode: bits 0-6
// * imm[0:4]: bits 7-11
// * funct3: bits 12-14
// * rs1: bits 15-19
// * rs2: bits 20-24
// * imm[5:11]: bits 25-31

// B-type
// * opcode: bits 0-6
// * imm[11]: bit 7
// * imm[1:4]: bits 8-11
// * funct3: bits 12-14
// * rs1: bits 15-19
// * rs2: bits 20-24
// * imm[5:10]: bits 25-30
// * imm[12]: bit 31

// U-type
// * opcode: bits 0-6
// * rd: bits 7-11
// * imm[12:31]

// J-type
// * opcode: bits 0-6
// * rd: bits 7-11
// * imm[12:19]: bits 12-19
// * imm[11]: bit 20
// * imm[1:10]: bits 21-30
// * imm[20]: bit 31

// sequential path is expected to be the common case on conditional branches
// this aligns nicely with encoding forward to the next unconditional branch
// but if possible unconditional jumps within a compilation unit may be doable

// memory ordering instructions -- how can any of this be done? :D

// FENCE.I could tell us to clear JIT caches :D



extern "C" fn interpreter(
    _machine: *mut MachineState,
    _state: *mut CoreState,
    pc: i64) -> i64 
 {
    println!("Hello, world!");
    return pc;
}

fn main() {
    let size = 8 * 1024 * 1024;
    let mut machine = MachineState::new(size);
    let mut core = CoreState::new(&machine);
    let pc = interpreter(&mut machine, &mut core, 0);
    println!("Ended with PC {}", pc);
}