/* Mycelium Scheme
 * Copyright (C) 2025 Ava Affine
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */


use organelle::{Fraction, Number, Numeric};

use crate::hmap::QuickMap;
use crate::stackstack::StackStack;
use crate::instr as i;
use crate::util::{Operand, Program, Address};
use crate::heap::{Gc, Datum, Cons};

use core::cell::RefCell;
use core::ops::DerefMut;

use alloc::vec;
use alloc::vec::Vec;
use alloc::sync::Arc;
use alloc::borrow::ToOwned;

use num::pow::Pow;


const NUM_OPERAND_REGISTERS: usize = 4;

pub struct VM {
    // execution environment
    pub stack:  StackStack<Gc<Datum>>,
    pub symtab: QuickMap<Operand>,
    pub prog:   Program,
    pub fds:    Vec<u64>,
    pub traps:  Vec<Arc<dyn Fn(&mut VM)>>,

    // data registers
    pub expr:  Gc<Datum>,
    pub oper: [Gc<Datum>; NUM_OPERAND_REGISTERS],

    // control flow registers
    pub retn: usize,
    pub ictr: usize,
    pub errr: Gc<Datum>,

    // state
    pub running: bool,
    pub err_state: bool,
}

impl VM {
    pub fn run_program(&mut self) {
        if self.prog.0.len() < 1 {
            self.running = false;
        }

        while self.ictr < self.prog.0.len() {
            if self.err_state || !self.running {
                return;
            }

            self.execute_instruction();
            self.ictr += 1;
        }

        self.running = false;
    }

    #[inline(always)]
    fn execute_instruction(&mut self) {
        let instr = &self.prog.0[self.ictr].clone();

        macro_rules! e {
            ( $err:expr ) => {
                {
                    self.running = false;
                    self.err_state = true;
                    self.errr = Datum::String($err.as_bytes().to_vec()).into();
                    return;
                }
            }
        }

        // get or set according to addressing mode
        macro_rules! access {
            ( $oper:expr ) => {
                match $oper.0 {
                    Address::Expr  => &self.expr,
                    Address::Oper1 => &self.oper[0],
                    Address::Oper2 => &self.oper[1],
                    Address::Oper3 => &self.oper[2],
                    Address::Oper4 => &self.oper[3],
                    Address::Stack => &self.stack[$oper.1],
                    Address::Numer => e!("cannot access constant numeric data"),
                    Address::Instr => e!("bad access to instruction data"),
                }
            };

            ( $data:expr, $target:expr ) => {
                match $data.0 {
                    Address::Expr => self.expr = $target,
                    Address::Oper1 => self.oper[0] = $target,
                    Address::Oper2 => self.oper[1] = $target,
                    Address::Oper3 => self.oper[2] = $target,
                    Address::Oper4 => self.oper[3] = $target,
                    _ => e!("attempted mutation of immutable address"),
                }
            }
        }

        macro_rules! do_jmp {
            ( $idx:expr ) => {
                let Operand(Address::Instr, target) = instr.1[$idx] else {
                    e!("illegal argument to jump");
                };

                if target >= self.prog.0.len() {
                    e!("out of bounds jump caught");
                }

               self.ictr = target;
            }
        }

        macro_rules! lr_oper {
            ( $in_type:ident, $oper:tt, $out_type:ident ) => {
                self.expr = Datum::$out_type(match **access!(&instr.1[0]){
                        Datum::$in_type(l) => l,
                        _ => e!("illegal argument to instruction"),
                    } $oper match **access!(&instr.1[1]){
                        Datum::$in_type(l) => l,
                        _ => e!("illegal argument to instruction"),
                    }).into()
            }
        }

        match instr.0 {
            i::TRAP => {
                let Operand(Address::Numer, idx) = instr.1[0] else {
                    e!("illegal argument to TRAP instruction");
                };

                if idx >= self.traps.len() {
                    e!("access to out of bounds trap!")
                }

                self.traps[idx].clone()(self)
            },

            // symtable ops
            i::BIND => {
                let Datum::String(ref tag) = **access!(&instr.1[0]) else {
                    e!("illegal argument to BIND instruction");
                };
                let tag = unsafe { str::from_utf8_unchecked(tag).to_owned() };
                self.symtab.insert(tag, instr.1[1].clone());
            },

            i::UNBIND => {
                let Datum::String(ref tag) = **access!(&instr.1[0]) else {
                    e!("illegal argument to UNBIND instruction");
                };
                let tag = unsafe { str::from_utf8_unchecked(tag) };
                self.symtab.remove(tag);
            },

            i::BOUND => {
                let Datum::String(ref tag) = **access!(&instr.1[0]) else {
                    e!("illegal argument to BOUND instruction");
                };
                let tag = unsafe { str::from_utf8_unchecked(tag) };
                self.symtab.contains_key(tag);
            },

            // stack ops
            i::PUSH => self.stack.push_current_stack(
                access!(&instr.1[0]).clone()),
            i::POP => _ = self.stack.pop_current_stack(),
            i::ENTER => self.stack.add_stack(),
            i::EXIT => self.stack.destroy_top_stack(),

            // movement ops
            i::LOAD => access!(&instr.1[1], access!(&instr.1[0]).clone()),
            i::DUPL => access!(&instr.1[1], access!(&instr.1[0]).deep_copy()),
            i::CLEAR => access!(&instr.1[0], Datum::None.into()),

            // control flow ops
            i::NOP => (),
            i::HALT => self.running = false,
            i::PANIC => {
                self.running = false;
                self.err_state = false;
                self.errr = access!(&instr.1[0]).clone();
            },

            i::JMP => {
                do_jmp!(0);
            },

            i::JMPIF => {
                if let Datum::Bool(true) = *self.expr {
                    do_jmp!(0);
                }
            },

            // boolean ops
            i::EQ => self.expr =
                Datum::Bool(*access!(&instr.1[0]) == *access!(&instr.1[1])).into(),
            i::LT => lr_oper!(Number, <, Bool),
            i::GT => lr_oper!(Number, >, Bool),
            i::LTE => lr_oper!(Number, <=, Bool),
            i::GTE => lr_oper!(Number, >=, Bool),
            i::BOOL_NOT => {
                self.expr = Datum::Bool(!{
                    let Datum::Bool(a) = *self.expr else {
                        e!("illegal argument to BOOL_NOT instruction");
                    };
                    a
                }).into();
            },

            i::BOOL_AND => lr_oper!(Bool, &&, Bool),
            i::BOOL_OR => lr_oper!(Bool, ||, Bool),

            // char / byte ops
            i::BYTE_AND => lr_oper!(Char, &, Char),
            i::BYTE_OR => lr_oper!(Char, |, Char),
            i::XOR => lr_oper!(Char, ^, Char),
            i::BYTE_NOT => {
                self.expr = Datum::Char(!{
                    let Datum::Char(a) = *self.expr else {
                        e!("illegal argument to BYTE_NOT instruction");
                    };
                    a
                }).into();
            },

            // numeric ops
            i::ADD => lr_oper!(Number, +, Number),
            i::SUB => lr_oper!(Number, -, Number),
            i::MUL => lr_oper!(Number, *, Number),
            i::FDIV => lr_oper!(Number, /, Number),
            i::IDIV => {
                let Datum::Number(ref l) = **access!(&instr.1[0]) else {
                    e!("illegal argument to IDIV instruction");
                };

                let Datum::Number(ref r) = **access!(&instr.1[1]) else {
                    e!("illgal argument to IDIV instruction");
                };

                let Fraction(l, 1) = l.make_exact() else {
                    e!("integer division on non integer value");
                };

                let Fraction(r, 1) = r.make_exact() else {
                    e!("integer division on non integer value");
                };

                self.expr = Datum::Number(Number::Fra(Fraction(l / r, 1))).into();
            },

            i::POW => {
                let Datum::Number(ref l) = **access!(&instr.1[0]) else {
                    e!("illegal argument to POW instruction");
                };

                let Datum::Number(ref r) = **access!(&instr.1[1]) else {
                    e!("illgal argument to POW instruction");
                };

                self.expr = Datum::Number(l.clone().pow(r.clone())).into();
            },

            i::INC => access!(&instr.1[0], {
                if let Datum::Number(src) = **access!(&instr.1[0]) {
                    Datum::Number(src + Number::Fra(Fraction(1, 1))).into()
                } else {
                    e!("illegal argument to INC instruction");
                }
            }),

            i::DEC => access!(&instr.1[0], {
                if let Datum::Number(src) = **access!(&instr.1[0]) {
                    Datum::Number(src - Number::Fra(Fraction(1, 1))).into()
                } else {
                    e!("illegal argument to INC instruction");
                }
            }),

            // byte/char to and from number conversions
            i::CTON => access!(&instr.1[0], {
                if let Datum::Char(schr) = **access!(&instr.1[0]) {
                    Datum::Number(Number::Fra(Fraction(schr as isize, 1))).into()
                } else {
                    e!("illegal argument to INC instruction");
                }
            }),

            i::NTOC => access!(&instr.1[0], {
                if let Datum::Number(snum) = **access!(&instr.1[0]) {
                    let n = snum.make_inexact();
                    if !snum.is_exact() || n.0.fract() != 0.0 ||
                        n.0 > u8::MAX.into() || n.0 < 0.0 {
                            e!("input to NTOC cannot cleanly convert");
                        }
                    Datum::Char(n.0.trunc() as u64 as u8).into()
                } else {
                    e!("illegal argument to INC instruction");
                }
            }),

            i::NTOI => {
                let src = access!(&instr.1[0]);
                if let Datum::Number(snum) = **src {
                    access!(&instr.1[0],
                        Datum::Number(snum.make_inexact().into()).into())
                }
            },

            i::NTOE => {
                let src = access!(&instr.1[0]);
                if let Datum::Number(snum) = **src {
                    access!(&instr.1[0], Datum::Number(snum.make_inexact().into())
                                            .into())
                }
            },

            i::CONST => access!(&instr.1[0], {
                let Operand(Address::Numer, num) = instr.1[0] else {
                    e!("illegal argument to CONST instruction");
                };

                Datum::Number(Number::Fra(Fraction(num as isize, 1))).into()
            }),

            i::MKVEC  => self.expr = Datum::Vector(RefCell::from(vec![])).into(),
            i::MKBVEC => self.expr = Datum::ByteVector(RefCell::from(vec![])).into(),
            i::INDEX  => {
                let Datum::Number(ref idx) = **access!(&instr.1[1]) else {
                    e!("illegal argument to INDEX instruction");
                };
                let idx = idx.make_inexact();
                if !idx.is_exact() || idx.0.fract() != 0.0 {
                    e!("illegal argument to INDEX instruction");
                }
                let idx = idx.0.trunc() as usize;

                match **access!(&instr.1[0]) {
                    Datum::Vector(ref v) => {
                        let a = (*v.borrow()[idx].clone()).clone();
                        self.expr = a.into();
                    },
                    Datum::ByteVector(ref bv) => {
                        let a = Datum::Char(bv.borrow()[idx]);
                        self.expr = a.into();
                    },
                    Datum::Cons(ref l) => self.expr = l[idx].clone(),
                    _ => e!("illegal argument to INDEX instruction")
                };
            },

            i::LENGTH => match **access!(&instr.1[0]) {
                Datum::Vector(ref v) => {
                    let a = Datum::Number(Number::Fra(Fraction(
                                v.borrow().len() as isize, 1)));
                    self.expr = a.into();
                },
                Datum::ByteVector(ref bv) => {
                    let a = Datum::Number(Number::Fra(Fraction(
                                bv.borrow().len() as isize, 1)));
                    self.expr = a.into();
                },
                Datum::Cons(ref l) => self.expr =
                    Datum::Number(Number::Fra(Fraction(l.len() as isize, 1)))
                        .into(),
                _ => e!("illegal argument to LENGTH instruction"),
            },

            i::SUBSL  => {
                let Datum::Number(ref st) = **access!(&instr.1[1]) else {
                    e!("illegal argument to SUBSL instruction");
                };

                let Datum::Number(ref ed) = **access!(&instr.1[2]) else {
                    e!("illegal argument to SUBSL instruction");
                };

                if !st.is_exact() || !ed.is_exact() {
                    e!("illegal argument to SUBSL instruction");
                }

                let st = st.make_inexact();
                let ed = ed.make_inexact();

                if st.0.fract() != 0.0 || ed.0.fract() != 0.0 {
                    e!("SUBSL: FP precision error");
                }

                let st = st.0.trunc() as usize;
                let ed = ed.0.trunc() as usize;

                match **access!(&instr.1[0]) {
                    Datum::Vector(ref v) => {
                        let a = Datum::Vector(RefCell::from(v.borrow()[st..ed].to_vec()));
                        self.expr = a.into();
                    },
                    Datum::ByteVector(ref bv) => {
                        let a = Datum::ByteVector(RefCell::from(bv.borrow()[st..ed].to_vec()));
                        self.expr = a.into();
                    },

                    Datum::Cons(ref a) => self.expr =
                        Datum::Cons(a.subsl(st as isize, ed as isize)).into(),
                    _ => e!("illegal argument to SUBSL instruction")
                };
            }

            i::INSER  => {
                let Datum::Number(ref idx) = **access!(&instr.1[2]) else {
                    e!("illegal argument to INSER instruction");
                };

                let idx = idx.make_inexact();
                if !idx.is_exact() || idx.0.fract() != 0.0 {
                    e!("illegal argument to INSER instruction");
                }

                let idx = idx.0.trunc() as usize;

                match **access!(&instr.1[0]) {
                    Datum::Vector(ref v) => {
                        v.borrow_mut()
                            .insert(idx, access!(&instr.1[1])
                                .deep_copy());
                    },
                    Datum::ByteVector(ref bv) => {
                        let Datum::Char(b) = **access!(&instr.1[1]) else {
                            e!("INSER instruction can only insert a byte into a bytevector");
                        };
                        bv.borrow_mut().insert(idx, b);
                    },
                    _ => e!("illegal argument to INSER instruction")
                }
            },

            i::CAR   => {
                let Datum::Cons(ref arg) = **access!(&instr.1[0]) else {
                    e!("illegal argument to CAR instruction");
                };

                self.expr = arg.clone().0
                    .or(Some(Datum::None.into()))
                    .expect("CAR instruction option consistency");
            },

            i::CDR   => {
                let Datum::Cons(ref arg) = **access!(&instr.1[0]) else {
                    e!("illegal argument to CAR instruction");
                };

                self.expr = arg.clone().1
                    .or(Some(Datum::None.into()))
                    .expect("CDR instruction option consistency");
            },

            i::CONS  => {
                let mut l = access!(&instr.1[0]).clone();
                if let Datum::Cons(l) = l.deref_mut() {
                    l.append(access!(&instr.1[1]).clone());
                } else {
                    access!(&instr.1[0], Datum::Cons(Cons(
                        Some(l),
                        Some(access!(&instr.1[1]).clone())
                    )).into());
                }
            },

            _ => {
                e!("illegal instruction");
            },
        };
    }
}