/* 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::ops::DerefMut;
use core::array;

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;

#[derive(Clone)]
pub struct VM {
    // execution environment
    pub stack:  StackStack<Gc<Datum>>,
    pub symtab: QuickMap<Operand>,
    pub prog:   Program,
    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 From<Program> for VM {
    fn from(value: Program) -> Self {
        VM{
            stack: StackStack::new(),
            symtab: QuickMap::new(),
            prog: value,
            traps: vec![],
            expr: Datum::None.into(),
            oper: array::from_fn(|_| Datum::None.into()),
            retn: 0,
            ictr: 0,
            errr: Datum::None.into(),
            running: true,
            err_state: false
        }
    }
}

impl VM {
    pub fn new_with_opts(
        program: Program,
        traps: Option<Vec<Arc<dyn Fn(&mut VM)>>>,
        stack: Option<StackStack<Gc<Datum>>>,
        syms:  Option<QuickMap<Operand>>
    ) -> VM {
        Into::<VM>::into(program)
            .with_stack(stack)
            .with_symbols(syms)
            .with_traps(traps)
            .to_owned() // not efficient, but we are not executing
    }

    pub fn with_stack(
        &mut self,
        maybe_stack: Option<StackStack<Gc<Datum>>>,
    ) -> &mut VM {
        if let Some(stack) = maybe_stack {
            self.stack = stack;
        }

        self
    }

    pub fn with_symbols(
        &mut self,
        maybe_symbols: Option<QuickMap<Operand>>,
    ) -> &mut VM {
        if let Some(symbols) = maybe_symbols {
            self.symtab = symbols;
        }

        self
    }

    pub fn with_traps(
        &mut self,
        maybe_traps: Option<Vec<Arc<dyn Fn(&mut VM)>>>,
    ) -> &mut VM {
        if let Some(traps) = maybe_traps {
            self.traps = traps;
        }

        self
    }

    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(vec![]).into(),

            i::MKBVEC => self.expr = Datum::ByteVector(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]).clone() {
                    Datum::Vector(ref v) =>
                        self.expr  = (*v[idx].clone()).clone().into(),
                    Datum::ByteVector(ref bv) =>
                        self.expr = Datum::Char(bv[idx]).into(),
                    Datum::String(ref s) =>
                        self.expr = Datum::Char(s[idx]).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) =>
                    self.expr = Datum::Number(Number::Fra(Fraction(
                                v.len() as isize, 1))).into(),
                Datum::ByteVector(ref bv) =>
                    self.expr = Datum::Number(Number::Fra(Fraction(
                                bv.len() as isize, 1))).into(),
                Datum::String(ref s) =>
                    self.expr = Datum::Number(Number::Fra(Fraction(
                                s.len() as isize, 1))).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]).clone().deref_mut() {
                    Datum::Vector(v) =>
                        self.expr  = Datum::Vector(v[st..ed].to_vec()).into(),
                    Datum::ByteVector(bv) =>
                        self.expr = Datum::ByteVector(bv[st..ed].to_vec()).into(),
                    Datum::String(s) =>
                        self.expr = Datum::String(s[st..ed].to_vec()).into(),
                    Datum::Cons(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]).clone().deref_mut() {
                    Datum::Vector(v) => {
                        v.insert(idx, access!(&instr.1[1]).clone());
                    },
                    Datum::ByteVector(bv) => {
                        let Datum::Char(b) = **access!(&instr.1[1]) else {
                            e!("INSER instruction can only insert a byte into a bytevector");
                        };
                        bv.insert(idx, b);
                    },
                    Datum::String(st) => {
                        let Datum::Char(b) = **access!(&instr.1[1]) else {
                            e!("INSER instruction can only insert a char into a string");
                        };
                        st.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());
                }
            },

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

                let Datum::String(ref right) = **access!(&instr.1[1]) else {
                    e!("illegal argument to CONCAT instruction");
                };

                let (left, right) = unsafe { (str::from_utf8_unchecked(left).to_owned(),
                                             str::from_utf8_unchecked(right)) };

                self.expr = Datum::String((left + right).as_bytes().to_vec()).into();
            },

            i::S_APPEND => {
                let mut left = access!(&instr.1[0]).clone();
                let Datum::String(left) = left.deref_mut() else {
                    e!("illegal argument to S_APPEND instruction");
                };

                let Datum::Char(right) = **access!(&instr.1[1]) else {
                    e!("illegal argument to S_APPEND instruction");
                };

                left.push(right);
            },

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

#[cfg(test)]
mod tests {
    use core::array;
    use super::*;
    use crate::instr;

    const ISA_TESTS: [Option<Vec<(VM, TestResult)>>; instr::TOTAL_INSTRUCTIONS] = [
        // TRAP
        None,
        // BIND
        None,
        // UNBIND
        None,
        // BOUND
        None,
        // PUSH
        None,
        // POP
        None,
        // ENTER
        None,
        // EXIT
        None,
        // LOAD
        None,
        // DUPL
        None,
        // CLEAR
        None,
        // NOP
        None,
        // HALT
        None,
        // PANIC
        None,
        // JMP
        None,
        // JMPIF
        None,
        // EQ
        None,
        // LT
        None,
        // GT
        None,
        // LTE
        None,
        // GTE
        None,
        // BOOL_NOT
        None,
        // BOOL_AND
        None,
        // BOOL_OR
        None,
        // BYTE_AND
        None,
        // BYTE_OR
        None,
        // XOR
        None,
        // BYTE_NOT
        None,
        // ADD
        None,
        // SUB
        None,
        // MUL
        None,
        // FDIV
        None,
        // IDIV
        None,
        // POW
        None,
        // MODULO
        None,
        // REM
        None,
        // INC
        None,
        // DEC
        None,
        // CTON
        None,
        // NTOC
        None,
        // NTOI
        None,
        // NTOE
        None,
        // CONST
        None,
        // MKVEC
        None,
        // MKBVEC
        None,
        // INDEX
        None,
        // LENGTH
        None,
        // SUBSL
        None,
        // INSER
        None,
        // CONS
        None,
        // CAR
        None,
        // CDR
        None,
        // CONCAT
        None,
        // S_APPEND
        None,
    ];

    fn has_stack(state: &VM, idx: usize, target: Gc<Datum>) -> bool {
        *(state.stack[idx]) == *target
    }

    fn has_sym(state: &VM, sym: &str, operand: Option<Operand>) -> bool {
        (operand.is_some() == state.symtab.contains_key(sym)) &&
            (operand.is_none() || state.symtab.get(&sym.to_owned()) == 
                operand.as_ref())
    }

    #[derive(Clone)]
    struct TestResult {
        expr: Option<Gc<Datum>>,
        stack: Vec<(usize, Gc<Datum>)>,
        syms: Vec<(&'static str, Option<Operand>)>,
        errr: Option<&'static str>
    }

    impl TestResult {
        fn chk_expr(&self, state: &VM) -> bool {
            self.expr.is_none() || *(state.expr) == *(self.expr.clone().unwrap())
        }

        fn chk_err(&self, state: &VM) -> bool {
            if let Datum::String(ref msg) = *state.errr {
                let msg = unsafe { str::from_utf8_unchecked(msg) };
                msg == self.errr.unwrap()
            } else if let Datum::None = *state.errr {
                self.errr.is_none()
            } else {
                false
            }
        }

        fn chk_stack(&self, state: &VM) -> bool {
            for i in &self.stack {
                if !has_stack(&state, i.0, i.1.clone()) {
                    return false
                }
            }
            true
        }

        fn chk_syms(&self, state: &VM) -> bool {
            for i in &self.syms {
                if !has_sym(&state, i.0, i.1.clone()) {
                    return false
                }
            }
            true
        }

        fn test_passes(&self, state: VM) -> bool {
            (self.expr.is_none()   || self.chk_expr(&state))  &&
            (self.stack.is_empty() || self.chk_stack(&state)) &&
            (self.syms.is_empty()  || self.chk_syms(&state))  &&
            (self.errr.is_none()   || self.chk_err(&state))
        }
    }

    #[test]
    fn run_isa_tests() {
        for i in &ISA_TESTS.to_vec() {
            let Some(test_case) = i else {
                assert!(false); // dont let untested instructions happen
                return
            };

            for i in test_case {
                let (mut vm, result) = (i.0.clone(), i.1.clone());
                vm.run_program();
                assert!(result.test_passes(vm));
            }
        }
    }
}