/* 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::Datum;

use core::cell::RefCell;

use alloc::vec;
use alloc::rc::Rc;
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<Datum>,
    pub symtab: QuickMap<Operand>,
    pub prog:   Program,
    pub fds:    Vec<u64>,
    pub traps:  Vec<Arc<dyn Fn(&mut VM)>>,

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

    // control flow registers
    pub retn: usize,
    pub ictr: usize,
    pub errr: 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());
                    return;
                }
            }
        }

        macro_rules! deref {
            ( $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!("attempt to dereference constant numeric data"),
                    Address::Instr => e!("bad access to instruction data"),
                }
            }
        }

        macro_rules! deref_mut {
            ( $oper:expr ) => {
                match $oper.0 {
                    Address::Expr  => &mut self.expr,
                    Address::Oper1 => &mut self.oper[0],
                    Address::Oper2 => &mut self.oper[1],
                    Address::Oper3 => &mut self.oper[2],
                    Address::Oper4 => &mut self.oper[3],
                    Address::Instr => e!("bad mutable access to instruction data"),
                    // Stack, Numer
                    _ => e!("mutable access to immutable data"),
                }
            }
        }

        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 deref!(&instr.1[0]){
                        Datum::$in_type(l) => l,
                        _ => e!("illegal argument to instruction"),
                    } $oper *match deref!(&instr.1[1]){
                        Datum::$in_type(l) => l,
                        _ => e!("illegal argument to instruction"),
                    })
            }
        }

        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(tag) = deref!(&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(tag) = deref!(&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(tag) = deref!(&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(deref!(&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 => *deref_mut!(&instr.1[1]) = deref!(&instr.1[0]).clone(),
            i::CLEAR => *deref_mut!(&instr.1[0]) = Datum::None,

            // control flow ops
            i::NOP => (),
            i::HALT => self.running = false,
            i::PANIC => {
                self.running = false;
                self.err_state = false;
                self.errr = deref!(&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(*deref!(&instr.1[0]) == *deref!(&instr.1[1])),
            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
                });
            },

            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
                });
            },

            // 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(l) = deref!(&instr.1[0]) else {
                    e!("illegal argument to IDIV instruction");
                };

                let Datum::Number(r) = deref!(&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)));
            },

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

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

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

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

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

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

            i::NTOC => {
                let src = deref_mut!(&instr.1[0]);
                if let Datum::Number(snum) = src {
                    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");
                    }
                    *src = Datum::Char(n.0.trunc() as u64 as u8);

                } else {
                    e!("illegal argument to NTOC instruction");
                }
            },

            i::MKVEC  => self.expr = Datum::Vector(RefCell::from(vec![])),
            i::MKBVEC => self.expr = Datum::ByteVector(RefCell::from(vec![])),
            i::INDEX  => {
                let Datum::Number(idx) = deref!(&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 deref!(&instr.1[0]) {
                    Datum::Vector(v) => {
                        let a = (*v.borrow()[idx].clone()).clone();
                        self.expr = a;
                    },
                    Datum::ByteVector(bv) => {
                        let a = Datum::Char(bv.borrow()[idx]);
                        self.expr = a;
                    },
                    Datum::List(l) => self.expr = l[idx].clone(),
                    _ => e!("illegal argument to INDEX instruction")
                };
            },

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

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

                let Datum::Number(ed) = deref!(&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 deref!(&instr.1[0]) {
                    Datum::Vector(v) => {
                        let a = Datum::Vector(RefCell::from(v.borrow()[st..ed].to_vec()));
                        self.expr = a;
                    },
                    Datum::ByteVector(bv) => {
                        let a = Datum::ByteVector(RefCell::from(bv.borrow()[st..ed].to_vec()));
                        self.expr = a;
                    },
                    Datum::List(a) =>
                        self.expr = Datum::List(Rc::new(
                                (**a).subsl(st as isize, ed as isize))),
                    _ => e!("illegal argument to SUBSL instruction")
                };
            }

            i::INSER  => {
                let Datum::Number(idx) = deref!(&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 deref!(&instr.1[0]) {
                    Datum::Vector(v) => {
                        v.borrow_mut().insert(idx, deref!(&instr.1[1]).clone().into());
                    },
                    Datum::ByteVector(bv) => {
                        let Datum::Char(b) = deref!(&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::List(arg) = deref!(&instr.1[0]) else {
                    e!("illegal argument to CAR instruction");
                };

                self.expr = (*arg.0).clone();
            },

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

                self.expr = (*arg.1).clone();
            },

            i::CONS  => {
                /* CONS BEHAVIOR
                 * L Datum is not list means create a new standard form list
                 * L Datum is list then append the second element to the first
                 */
            },

            // in order to maintain a language agnostic VM these must be traps
            //i::PARSE => todo!("implement AST API"),
            //i::EVAL  => todo!("implement AST API"),

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