#![feature(allocator_api, box_as_ptr)]

unsafe extern "C" {
    fn heap_alloc(size: usize, align: usize) -> *mut u8;
    fn heap_dealloc(ptr: *mut u8, size: usize, align: usize);

    fn init_source(source: *mut Source, fd: i32) -> *mut Source;
    fn getc() -> SourceIterResult;
    fn peekc() -> SourceIterResult;

    fn eval(expr: Object, env: Object) -> Object;
    fn init_env() -> Object;

    fn parse_next_token() -> Object;

    #[link_name = "ifile"]
    static mut IFILE: Source;
    #[link_name = "nil"]
    static NIL: ();
    #[link_name = "env"]
    static GENV: ();
}

use std::ptr::NonNull;

#[repr(C)]
#[derive(Copy, Clone)]
struct Object(*mut ());

impl Object {
    const BYTE: u8 = 0;
    const NUM: u8 = 1;
    const PRIM: u8 = 2;
    const CONS: u8 = 3;
    const CLOS: u8 = 4;
    const ATOM: u8 = 5;
    const ARR: u8 = 6;

    const NUM_MAGIC: u32 = 0x5555;

    fn is_nil(&self) -> bool {
        self.0 as *const () == &raw const NIL
    }
    fn tag(&self) -> u8 {
        self.0.addr() as u8 & 0b111
    }
    fn bits(&self) -> usize {
        self.0.addr()
    }
    fn ptr(&self) -> *mut () {
        self.0.map_addr(|addr| addr & !0b111)
    }
    fn is_g_env(&self) -> bool {
        self.0 as *const () == unsafe { *ENV_INIT }.0 as *const ()
    }
}

#[derive(Copy, Clone)]
struct DebugObjectExplicitGenv(Object);

impl std::fmt::Debug for DebugObjectExplicitGenv {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        #[repr(C)]
        struct Cons {
            refc: u64,
            car: Object,
            cdr: Object,
        }

        if self.0.is_g_env() {
            let mut proxy = unsafe {
                Cons {
                    refc: 1,
                    car: self.0.ptr().byte_add(8).cast::<Object>().read(),
                    cdr: self.0.ptr().byte_add(16).cast::<Object>().read(),
                }
            };
            let obj =
                Object((&raw const proxy as *mut ()).map_addr(|addr| addr | Object::CONS as usize));
            write!(f, "{:?}", obj)
        } else {
            return write!(f, "{:?}", self.0);
        }
    }
}

impl std::fmt::Debug for Object {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        if self.is_nil() {
            return write!(f, "nil");
        }

        if self.is_g_env() {
            return write!(f, "#<g_env>");
        }

        fn fmt_byte(byte: *const (), f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            let c = unsafe { byte.cast::<u8>().read() };
            if c.is_ascii_graphic() || c == b' ' {
                write!(f, "'{}'", c as char)
            } else {
                write!(f, "'\\x{:02x}'", c)
            }
        }

        fn fmt_num(num: *const (), f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            let num_ptr = unsafe { num.cast::<i64>().read() };
            write!(f, "{}", num_ptr)
        }

        fn fmt_prim(prim: *const (), f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            let prim_ptr = unsafe { prim.cast::<*const ()>().read() };
            write!(f, "#<prim {:?}>", prim_ptr)
        }

        fn fmt_cons(cons: *const (), f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            let (car, cdr) = unsafe { cons.cast::<(Object, Object)>().read() };

            if !cdr.is_nil() && cdr.tag() == Object::CONS {
                // If the cdr is a cons cell, we can print it as a list
                write!(f, "({car:?}")?;
                let mut cdr = cdr;
                while !cdr.is_nil() && cdr.tag() == Object::CONS {
                    let (next_car, next_cdr) =
                        unsafe { cdr.ptr().byte_add(8).cast::<(Object, Object)>().read() };
                    write!(f, " {next_car:?}")?;
                    cdr = next_cdr;
                }
                if cdr.is_nil() {
                    write!(f, ")")
                } else {
                    write!(f, " {cdr:?})")
                }
            } else {
                write!(f, "({car:?} . {cdr:?})")
            }
        }

        fn fmt_clos(clos: *const (), f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            let (car, env) = unsafe { clos.cast::<(Object, Object)>().read() };
            assert_eq!(
                car.tag(),
                Object::CONS,
                "Expected a pair for lambda cdr, got tag {}",
                car.tag()
            );
            let (params, body) = unsafe { car.ptr().cast::<(Object, Object)>().read() };
            write!(f, "λ ")?;
            if env.0 as *const () == &raw const NIL {
                write!(f, " #<g_env> ++ ")
            } else {
                write!(f, " {env:?} ++ ")
            }
        }

        fn fmt_atom(atom: *const (), f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            let (ptr, len) = unsafe {
                (
                    atom.cast::<*const u8>().read(),
                    atom.byte_add(8).cast::<usize>().read(),
                )
            };

            let slice = unsafe { core::slice::from_raw_parts(ptr, len) };
            let string = core::str::from_utf8(slice).unwrap_or("<invalid utf-8>");
            write!(f, "'{string}")
        }

        fn fmt_arr(arr: *const (), f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            let len = unsafe { arr.cast::<u32>().read() } as usize;
            let cap = unsafe { arr.byte_add(4).cast::<u32>().read() } as usize;
            let data = unsafe { arr.byte_add(8).cast::<Object>().read() };

            let data_ptr = data.ptr();
            let data_tag = data.tag();

            match data_tag {
                Object::BYTE => {
                    let slice = unsafe { core::slice::from_raw_parts(data_ptr as *const u8, len) };
                    let string = core::str::from_utf8(slice).unwrap_or("<invalid utf-8>");
                    write!(f, "\"{string}\"")
                }
                _ => {
                    const SIZES: [usize; 7] = [1, 8, 8, 16, 16, 16, 16];
                    let _fn = match data_tag {
                        Object::BYTE => fmt_byte,
                        Object::NUM => fmt_num,
                        Object::PRIM => fmt_prim,
                        Object::CONS => fmt_cons,
                        Object::CLOS => fmt_clos,
                        Object::ATOM => fmt_atom,
                        Object::ARR => fmt_arr,
                        _ => unreachable!(),
                    };

                    write!(f, "[")?;
                    for i in 0..len {
                        let elem = unsafe { data_ptr.byte_add(i * SIZES[data_tag as usize]) };
                        _fn(elem, f)?;
                        if i + 1 < len {
                            write!(f, ", ")?;
                        }
                    }
                    write!(f, "]")
                }
            }
        }

        unsafe {
            match self.tag() {
                Object::BYTE => {
                    let byte = self.ptr().addr() >> 8 as u8;
                    fmt_byte(&raw const byte as _, f)
                }
                Object::NUM => {
                    let MASK: usize = 0x5555_0000_0000_0000 | Object::NUM as usize;

                    if self.bits() & MASK == MASK {
                        let num = (self.bits() >> 8) as i32 as i64;
                        fmt_num(&raw const num as _, f)
                    } else {
                        fmt_num(self.ptr().byte_add(8), f)
                    }
                }
                Object::PRIM => fmt_prim(self.ptr().byte_add(8), f),
                Object::CONS => fmt_cons(self.ptr().byte_add(8), f),
                Object::CLOS => fmt_clos(self.ptr().byte_add(8), f),
                Object::ATOM => fmt_atom(self.ptr().byte_add(8), f),
                Object::ARR => fmt_arr(self.ptr().byte_add(8), f),
                _ => write!(f, "#<unknown tag {}>", self.tag()),
            }
        }
    }
}

#[repr(C)]
struct Source {
    fd: i32,
    peeked: (u8, u8),
    buf: NonNull<u8>,
    buf_cur: usize,
    buf_end: usize,
}

#[repr(align(8))]
struct Bool(bool);

impl core::ops::Deref for Bool {
    type Target = bool;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

#[repr(C)]
struct SourceIterResult {
    c: u8,
    some: Bool,
}

struct HeapAlloc;

unsafe impl std::alloc::Allocator for HeapAlloc {
    fn allocate(
        &self,
        layout: std::alloc::Layout,
    ) -> Result<std::ptr::NonNull<[u8]>, std::alloc::AllocError> {
        let ptr = unsafe { heap_alloc(layout.size(), layout.align()) };
        if ptr.is_null() {
            Err(std::alloc::AllocError)
        } else {
            Ok(std::ptr::NonNull::slice_from_raw_parts(
                std::ptr::NonNull::new(ptr).unwrap(),
                layout.size(),
            ))
        }
    }

    unsafe fn deallocate(&self, ptr: std::ptr::NonNull<u8>, layout: std::alloc::Layout) {
        heap_dealloc(ptr.as_ptr(), layout.size(), layout.align());
    }
}

static mut ENV_INIT: std::cell::LazyCell<Object> =
    std::cell::LazyCell::new(|| unsafe { init_env() });

#[cfg(test)]
mod tests {
    use super::*;
    use std::fs::File;
    use std::mem::ManuallyDrop;
    use std::os::fd::AsRawFd;

    #[test]
    fn test_heap_alloc_dealloc() {
        let small = Box::new_in(42u8, HeapAlloc);
        let ptr_sized = Box::new_in(42usize, HeapAlloc);
        let big = Box::new_in([0u8; 48], HeapAlloc);
        let large = Box::new_in([0u8; 1024], HeapAlloc);
        let page_sized = Box::new_in([0u8; 4096], HeapAlloc);
        assert_eq!(*small, 42);
        assert_ne!(
            Box::as_ptr(&small) as *const (),
            Box::as_ptr(&ptr_sized) as *const ()
        );
        let small_ptr = Box::as_ptr(&small) as *const ();
        drop(small);
        let small2 = Box::new_in(43u8, HeapAlloc);
        assert_eq!(*small2, 43);
        assert_eq!(small_ptr, Box::as_ptr(&small2) as *const ());
    }

    #[test]
    fn source() {
        let mut file = std::fs::File::open("lisp1.asm").unwrap();
        unsafe {
            init_source(&raw mut IFILE, file.as_raw_fd());
            let SourceIterResult { c, some } = getc();
            assert!(*some);
            assert_eq!(c, b'd');
            let SourceIterResult { c, some } = peekc();
            assert!(*some);
            assert_eq!(c, b'e');
            let SourceIterResult { c, some } = getc();
            assert!(*some);
            assert_eq!(c, b'e');
        }
    }

    #[test]
    fn parse_list() {
        let file = ManuallyDrop::new(File::open("tests/list.l").unwrap());
        unsafe {
            init_source(&raw mut IFILE, file.as_raw_fd());
            let obj = parse_next_token();
            println!("{:?}", obj);
        }
    }

    #[test]
    fn parse_quoted_list() {
        let file = ManuallyDrop::new(File::open("tests/quoted-list.l").unwrap());
        unsafe {
            init_source(&raw mut IFILE, file.as_raw_fd());
            let obj = parse_next_token();
            println!("{:?}", obj);
        }
    }

    #[test]
    fn parse_string() {
        let file = ManuallyDrop::new(File::open("tests/string.l").unwrap());
        unsafe {
            init_source(&raw mut IFILE, file.as_raw_fd());
            let obj = parse_next_token();
            println!("{:?}", obj);
        }
    }

    #[test]
    fn parse_atom() {
        let file = ManuallyDrop::new(File::open("tests/atom.l").unwrap());
        unsafe {
            init_source(&raw mut IFILE, file.as_raw_fd());
            let obj = parse_next_token();
            println!("{:?}", obj);
        }
    }

    #[test]
    fn parse_pair() {
        let file = ManuallyDrop::new(File::open("tests/pair.l").unwrap());
        unsafe {
            init_source(&raw mut IFILE, file.as_raw_fd());
            let obj = parse_next_token();
            println!("{:?}", obj);
        }
    }

    #[test]
    fn parse_big_number() {
        let file = ManuallyDrop::new(File::open("tests/big-number.l").unwrap());
        unsafe {
            init_source(&raw mut IFILE, file.as_raw_fd());
            let obj = parse_next_token();
            println!("{:?}", obj);
        }
    }

    #[test]
    fn global_env() {
        unsafe {
            let env = unsafe { *ENV_INIT };
            println!("{:?}", DebugObjectExplicitGenv(env));
        }
    }

    #[test]
    fn eval_add() {
        let file = ManuallyDrop::new(File::open("tests/add.l").unwrap());
        unsafe {
            init_source(&raw mut IFILE, file.as_raw_fd());
            let expr = parse_next_token();
            let env = unsafe { *ENV_INIT };
            let result = eval(expr, env);
            eprint!("done: ");
            eprintln!("{result:?}");
        }
    }

    #[test]
    fn eval_let() {
        let file = ManuallyDrop::new(File::open("tests/let.l").unwrap());
        unsafe {
            init_source(&raw mut IFILE, file.as_raw_fd());
            let expr = parse_next_token();
            let env = unsafe { *ENV_INIT };
            let result = eval(expr, env);
            eprint!("done: ");
            eprintln!("{result:?}");
        }
    }

    #[test]
    fn eval_letstar() {
        let file = ManuallyDrop::new(File::open("tests/letstar.l").unwrap());
        unsafe {
            init_source(&raw mut IFILE, file.as_raw_fd());
            let expr = parse_next_token();
            let env = unsafe { *ENV_INIT };
            let result = eval(expr, env);
            eprint!("done: ");
            eprintln!("{result:?}");
        }
    }

    #[test]
    fn eval_print() {
        let file = ManuallyDrop::new(File::open("tests/print.l").unwrap());
        unsafe {
            init_source(&raw mut IFILE, file.as_raw_fd());
            let expr = parse_next_token();
            let env = unsafe { *ENV_INIT };
            let result = eval(expr, env);
            eprint!("done: ");
            eprintln!("{result:?}");
        }
    }

    #[test]
    fn eval_if() {
        let file = ManuallyDrop::new(File::open("tests/if.l").unwrap());
        unsafe {
            init_source(&raw mut IFILE, file.as_raw_fd());
            let expr = parse_next_token();
            let env = unsafe { *ENV_INIT };
            let result = eval(expr, env);
            eprint!("done: ");
            eprintln!("{result:?}");
        }
    }

    #[test]
    fn eval_define() {
        let file = ManuallyDrop::new(File::open("tests/define.l").unwrap());
        unsafe {
            init_source(&raw mut IFILE, file.as_raw_fd());
            let expr = parse_next_token();
            let env = unsafe { *ENV_INIT };
            let result = eval(expr, env);
            eprint!("done: ");
            eprintln!("{result:?}");
        }
    }

    #[test]
    fn eval_typeof() {
        let file = ManuallyDrop::new(File::open("tests/typeof.l").unwrap());
        unsafe {
            init_source(&raw mut IFILE, file.as_raw_fd());
            let expr = parse_next_token();
            let env = unsafe { *ENV_INIT };
            let result = eval(expr, env);
            eprint!("done: ");
            eprintln!("{result:?}");
        }
    }
}