/* 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 crate::instr::Operation;

use alloc::vec::Vec;
use alloc::vec;

use core::ops::Index;
use core::mem::transmute;

#[repr(u8)]
#[derive(Debug, Clone, PartialEq)]
pub enum Address {
    Stack = 0xf0, // immutable access only
    Instr = 0xf1, // immutable access only
    Expr  = 0xf2, // mutable access allowed
    Oper1 = 0xf3, // mutable access allowed
    Oper2 = 0xf4, // mutable access allowed
    Oper3 = 0xf5, // mutable access allowed
    Oper4 = 0xf6, // mutable access allowed
    Numer = 0xf8, // immutable access only
}

#[derive(Debug, Clone, PartialEq)]
pub struct Operand(pub Address, pub usize);

#[derive(Debug, Clone, PartialEq)]
pub struct Instruction(pub Operation, pub Vec<Operand>);

#[derive(Debug, Clone, PartialEq)]
pub struct Program(pub Vec<Instruction>);

impl Into<u8> for Address {
    fn into(self) -> u8 {
        unsafe { transmute::<Address, u8>(self) }
    }
}

impl TryFrom<u8> for Address {
    type Error = &'static str;
    fn try_from(val: u8) -> Result<Self, Self::Error> {
        match val {
            _ if val == Address::Stack as u8 => Ok(Address::Stack),
            _ if val == Address::Instr as u8 => Ok(Address::Instr),
            _ if val == Address::Expr  as u8 => Ok(Address::Expr),
            _ if val == Address::Oper1 as u8 => Ok(Address::Oper1),
            _ if val == Address::Oper2 as u8 => Ok(Address::Oper2),
            _ if val == Address::Oper3 as u8 => Ok(Address::Oper3),
            _ if val == Address::Oper4 as u8 => Ok(Address::Oper4),
            _ if val == Address::Numer as u8 => Ok(Address::Numer),
            _ => Err("illegal addressing mode")
        }
    }
}

impl Address {
    fn operand_size(&self) -> u8 {
        match self {
            Address::Stack => (usize::BITS / 8) as u8,
            Address::Instr => (usize::BITS / 8) as u8,
            Address::Numer => (usize::BITS / 8) as u8,
            _ => 0,
        }
    }
}

impl TryFrom<&[u8]> for Operand {
    type Error = &'static str;
    fn try_from(value: &[u8]) -> Result<Self, Self::Error> {
        let addr_mode: Address = value[0].try_into()?;
        let operand_size = addr_mode.operand_size();
        if value.len() < (operand_size + 1).into() {
            return Err("truncated address data")
        }

        let mut operand_bytes: [u8; 8] = [0, 0, 0, 0, 0, 0, 0, 0];
        for (&src, dest) in value[1..(1+operand_size) as usize]
            .iter()
            .zip(operand_bytes.iter_mut()) {
                *dest = src;
        }

        Ok(Operand(addr_mode, usize::from_ne_bytes(operand_bytes)))
    }
}

impl Into<Vec<u8>> for Operand {
    fn into(self) -> Vec<u8> {
        let mut res = vec![];
        res.push(self.0.clone() as u8);
        res.append(&mut self.1.to_ne_bytes()[..self.0.operand_size() as usize].to_vec());
        res
    }
}

impl Operand {
    fn byte_length(&self) -> u8 {
        1 + self.0.operand_size()
    }
}

impl TryFrom<&[u8]> for Instruction {
    type Error = &'static str;
    fn try_from(value: &[u8]) -> Result<Self, Self::Error> {
        let operation: Operation = value[0].try_into()?;
        let mut operands: Vec<Operand> = vec![];

        let mut cur = 1;
        for _ in 0..operation.num_args()? {
            if cur >= value.len() {
                return Err("operand data truncated")
            }
            let operand: Operand = value[cur..].try_into()?;
            cur += operand.byte_length() as usize;
            operands.push(operand);
        }

        Ok(Instruction(operation, operands))
    }
}

impl Into<Vec<u8>> for Instruction {
    fn into(self) -> Vec<u8> {
        let mut res = vec![];
        res.push(self.0.0);
        for op in self.1 {
            res.append(&mut op.into())
        }
        res
    }
}

impl Instruction {
    fn byte_length(&self) -> u8 {
        self.1.iter()
            .fold(0, |total, oper|
                total + oper.byte_length()) + 1
    }
}

impl TryFrom<&[u8]> for Program {
    type Error = &'static str;
    fn try_from(value: &[u8]) -> Result<Self, Self::Error> {
        let mut prog: Vec<Instruction> = vec![];
        let mut cur = 0;

        while cur < value.len() {
            let instruction: Instruction = value[cur..].try_into()?;
            cur += instruction.byte_length() as usize;
            prog.push(instruction);
        }

        Ok(Program(prog))
    }
}

impl Into<Vec<u8>> for Program {
    fn into(self) -> Vec<u8> {
        let mut res: Vec<u8> = vec![];
        for instr in self.0 {
            res.append(&mut instr.into())
        }
        res
    }
}

impl<'a> Index<usize> for Program {
    type Output = Instruction;
    fn index(&self, index: usize) -> &Instruction {
        self.0.get(index).expect("access to out of bounds instruction in vm")
    }
}


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

    #[test]
    fn test_operand_parse() {
        let bad_addressing =
            TryInto::<Operand>::try_into(&[0x13, 0x39][..]);
        assert_eq!(bad_addressing, Err("illegal addressing mode"));

        let truncated_address =
            TryInto::<Operand>::try_into(&[0xf1][..]);
        assert_eq!(truncated_address, Err("truncated address data"));

        let usize_case =
            TryInto::<Operand>::try_into(&[Address::Stack.into(),
                0x23, 0x23, 0x23, 0x23, 0x23, 0x23, 0x23, 0x23][..]);
        assert!(usize_case.is_ok());
        assert_eq!(usize_case.unwrap().0, Address::Stack);

        let register_operand = Operand(Address::Expr, 0);
        let operand_byte_arr =
            TryInto::<Vec<u8>>::try_into(register_operand.clone());
        assert!(operand_byte_arr.is_ok());
        let br = operand_byte_arr.unwrap();
        let operand_bytes = br.as_slice();
        assert_eq!(operand_bytes, &[0xf2][..]);
        let operand_conv =
            TryInto::<Operand>::try_into(operand_bytes);
        assert!(operand_conv.is_ok());
        assert_eq!(register_operand, operand_conv.unwrap());
    }

    #[test]
    fn test_instruction_parse() {
        let illegal_instruction =
            TryInto::<Instruction>::try_into(&[0x88][..]);
        assert_eq!(illegal_instruction, Err("illegal instruction"));

        let bad_operand =
            TryInto::<Instruction>::try_into(&[instr::TRAP.0, 0xf1][..]);
        assert_eq!(bad_operand, Err("truncated address data"));

        let need_more_opers =
            TryInto::<Instruction>::try_into(&[instr::TRAP.0][..]);
        assert_eq!(need_more_opers, Err("operand data truncated"));

        let no_operands =
            TryInto::<Instruction>::try_into(&[instr::POP.0][..]);
        assert!(no_operands.is_ok());
        let nop = no_operands.unwrap();
        assert_eq!(nop.0, instr::POP);
        let nop_bytes =
            TryInto::<Vec<u8>>::try_into(nop);
        assert!(nop_bytes.is_ok());
        assert_eq!(nop_bytes.unwrap(), vec![instr::POP.0]);

        let one_operand =
            TryInto::<Instruction>::try_into(&[instr::TRAP.0, 0xf3][..]);
        assert!(one_operand.is_ok());
        let oe_oper = one_operand.unwrap();
        assert_eq!(oe_oper.0, instr::TRAP);
        assert_eq!(oe_oper.1.len(), 1);
        assert_eq!(oe_oper.1[0], Operand(Address::Oper1, 0));
        let oe_bytes =
            TryInto::<Vec<u8>>::try_into(oe_oper);
        assert!(oe_bytes.is_ok());
        assert_eq!(oe_bytes.unwrap(), vec![instr::TRAP.0, 0xf3]);

        let two_operands =
            TryInto::<Instruction>::try_into(&[instr::LINK.0, 0xf3, 0xf4][..]);
        assert!(two_operands.is_ok());
        let two_oper = two_operands.unwrap();
        assert_eq!(two_oper.0, instr::LINK);
        assert_eq!(two_oper.1.len(), 2);
        let two_bytes =
            TryInto::<Vec<u8>>::try_into(two_oper.clone());
        assert!(two_bytes.is_ok());
        assert_eq!(two_bytes.unwrap(), vec![instr::LINK.0, 0xf3, 0xf4]);
        assert_eq!(two_oper.1[0], Operand(Address::Oper1, 0));
        assert_eq!(two_oper.1[1], Operand(Address::Oper2, 0));
   }

    #[test]
    fn test_program_parse() {
        let bytes1 = [instr::LINK.0, 0xf3, 0xf4];
        let out1 = vec![Instruction(instr::LINK,
            vec![Operand(Address::Oper1, 0), Operand(Address::Oper2, 0)])];
        let res1 =
            TryInto::<Program>::try_into(&bytes1[..]);
        assert!(res1.is_ok());
        assert_eq!(res1.unwrap().0, out1);

        let bytes2 = [
            instr::LINK.0, 0xf3, 0xf4,
            instr::CLEAR.0, 0xf0, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00
        ];
        let out2 = vec![
            Instruction(instr::LINK, vec![
                Operand(Address::Oper1, 0),
                Operand(Address::Oper2, 0)
            ]),

            Instruction(instr::CLEAR, vec![
                Operand(Address::Stack, 1)
            ])
        ];
        let res2 =
            TryInto::<Program>::try_into(&bytes2[..]);
        assert!(res2.is_ok());
        assert_eq!(res2.unwrap().0, out2);
    }
}