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

use alloc::vec::Vec;
use alloc::vec;
use alloc::str::FromStr;
use alloc::fmt::{Display, Formatter, Error as E};

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


#[repr(u8)]
#[derive(Debug, Clone, PartialEq)]
pub enum DeserializerControlCode {
    SciNumber   = 0x00,
    FltNumber   = 0x01,
    FraNumber   = 0x02,
    SymInf      = 0x03,
    SymNan      = 0x04,
    SymNegInf   = 0x05,
    SymNegNan   = 0x06,
    BoolFalse   = 0x07,
    BoolTrue    = 0x08,
    Char        = 0x09,
    String      = 0x0A,
    ByteVec     = 0x0B,
    Vector      = 0x0C,
    EmptyCons   = 0x0D,
    LeftCons    = 0x0E,
    RightCons   = 0x0F,
    FullCons    = 0x10,
    DataChunk   = 0x11,
    CodeChunk   = 0x12,
}

#[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
    Bool   = 0xf9, // immutable access only
    Char   = 0xfa, // 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<Datum>, 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),
            _ if val == Address::Bool   as u8 => Ok(Address::Bool),
            _ if val == Address::Char   as u8 => Ok(Address::Char),
            _ => 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,
            Address::Bool  => 1,
            Address::Char  => 1,
            _ => 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 FromStr for Operand {
    type Err = &'static str;
    fn from_str(v: &str) -> Result<Self, Self::Err> {
        match v {
            "$expr"  => Ok(Operand(Address::Expr,  0)),
            "$oper1" => Ok(Operand(Address::Oper1, 0)),
            "$oper2" => Ok(Operand(Address::Oper2, 0)),
            "$oper3" => Ok(Operand(Address::Oper3, 0)),
            "$oper4" => Ok(Operand(Address::Oper4, 0)),
            "true"   => Ok(Operand(Address::Bool,  1)),
            "false"  => Ok(Operand(Address::Bool,  0)),
            a if a.len() > 1 &&
                 a.chars().nth(0).unwrap() == '%' &&
                 a[1..].parse::<usize>().is_ok() =>
                    Ok(Operand(Address::Stack, a[1..].parse::<usize>().unwrap())),
            a if a.len() > 1 &&
                 a.chars().nth(0).unwrap() == '@' &&
                 a[1..].parse::<usize>().is_ok() =>
                    Ok(Operand(Address::Instr, a[1..].parse::<usize>().unwrap())),
            a if a.len() == 3 &&
                 a.chars().nth(0).unwrap() == '\'' &&
                 a.chars().nth(2).unwrap() == '\'' =>
                    Ok(Operand(Address::Char, a.chars().nth(1).unwrap() as usize)),
            a if a.parse::<usize>().is_ok() =>
                Ok(Operand(Address::Numer, a.parse::<usize>().unwrap())),
            _ => Err("invalid operand")
        }
    }
}

impl Display for Operand {
    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), E> {
        match self.0 {
            Address::Expr  => write!(f, "$expr"),
            Address::Oper1 => write!(f, "$oper1"),
            Address::Oper2 => write!(f, "$oper2"),
            Address::Oper3 => write!(f, "$oper3"),
            Address::Oper4 => write!(f, "$oper4"),
            Address::Bool  =>
                write!(f, "{}", if self.1 > 0 { "true" } else { "false" }),
            Address::Stack => write!(f, "%")
                .and_then(|_| write!(f, "{}", self.1)),
            Address::Instr => write!(f, "@")
                .and_then(|_| write!(f, "{}", self.1)),
            Address::Numer => write!(f, "{}", self.1),
            Address::Char  => write!(f, "'{}'", self.1 as u8 as char),
        }
    }
}

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 FromStr for Instruction {
    type Err = &'static str;
    fn from_str(v: &str) -> Result<Self, Self::Err> {
        let toks: Vec<&str> = v
            .trim()
            .split(' ')
            .filter(|x| x.len() > 0)
            .collect();
        if toks.len() < 1 {
            return Err("empty string");
        }

        let oper = Operation::from_str(toks[0])?;
        let mut args = vec![];
        if toks.len() == 1 && oper.num_args()? == 0 {
            return Ok(Instruction(oper, args));
        }
        for i in toks[1..].iter() {
            args.push(Operand::from_str(i.trim_matches(','))?);
        }

        if oper.num_args()? as usize != args.len() {
            return Err("instruction has incorrect number of operands");
        }

        Ok(Instruction(oper, args))
    }
}

impl Display for Instruction {
    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), E> {
        write!(f, "{} \t", self.0)?;
        if self.1.len() > 0 {
            write!(f, "{}", self.1[0])?;
        }
        if self.1.len() > 1 {
            for i in self.1[1..].iter() {
                write!(f, ", {}", i)?;
            }
        }
        Ok(())
    }
}

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 data: Vec<Datum> = vec![];
        let mut prog: Vec<Instruction> = vec![];

        let mut parse_data = || -> Result<usize, Self::Error> {
            let mut cur = 0;
            if value[cur] != DeserializerControlCode::DataChunk as u8 {
                return Ok(cur);
            }
            cur += 1;
            while value[cur] != DeserializerControlCode::CodeChunk as u8 {
                let datum: Datum = value[cur..].try_into()?;
                cur += datum.byte_length();
                data.push(datum);
            }
            Ok(cur)
        };

        let mut parse_code = |cur: usize| -> Result<(), Self::Error> {
            let mut cur = cur;
            if value[cur] != DeserializerControlCode::CodeChunk as u8 {
                return Err("no code chunk detected in program");
            }
            cur += 1;
            while cur < value.len() {

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

        parse_code(parse_data()?)?;
        Ok(Program(data, prog))
    }
}

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

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

impl Display for Program {
    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), E> {
        write!(f, "DATA:\n")?;
        for i in self.0.iter() {
            write!(f, "    {}\n", i)?;
        }

        write!(f, "CODE:\n")?;
        for i in self.1.iter() {
            write!(f, "    {}\n", i)?;
        }

        Ok(())
    }
}

impl FromStr for Program {
    type Err = &'static str;
    fn from_str(val: &str) -> Result<Self, Self::Err> {
        //let mut datum = vec![];
        let mut instrs = vec![];
        let lines: Vec<&str> = val
            .split('\n')
            .filter(|x| x.len() > 0)
            .collect();
        let mut cur = 0;
        let mut toggle = 0;

        while cur < lines.len() {
            if toggle == 1 {
                instrs.push(lines[cur].parse::<Instruction>()?);

            // TODO: toggle == 2 case for interpreting a DATA chunk

            } else {
                match lines[cur] {
                    "DATA:" => return Err("datum parser unimplemented"),
                    "CODE:" => toggle = 1,
                    _ => return Err("unknown section in document: "),
                }
            }

            cur += 1;
        }

        Ok(Program(vec![], instrs))
    }
}

impl TryFrom<u8> for DeserializerControlCode {
    type Error = &'static str;
    fn try_from(value: u8) -> Result<Self, Self::Error> {
        match value {
            0x00 => Ok(DeserializerControlCode::SciNumber),
            0x01 => Ok(DeserializerControlCode::FltNumber),
            0x02 => Ok(DeserializerControlCode::FraNumber),
            0x03 => Ok(DeserializerControlCode::SymInf),
            0x04 => Ok(DeserializerControlCode::SymNan),
            0x05 => Ok(DeserializerControlCode::SymNegInf),
            0x06 => Ok(DeserializerControlCode::SymNegNan),
            0x07 => Ok(DeserializerControlCode::BoolFalse),
            0x08 => Ok(DeserializerControlCode::BoolTrue),
            0x09 => Ok(DeserializerControlCode::Char),
            0x0A => Ok(DeserializerControlCode::String),
            0x0B => Ok(DeserializerControlCode::ByteVec),
            0x0C => Ok(DeserializerControlCode::Vector),
            0x0D => Ok(DeserializerControlCode::EmptyCons),
            0x0E => Ok(DeserializerControlCode::LeftCons),
            0x0F => Ok(DeserializerControlCode::RightCons),
            0x10 => Ok(DeserializerControlCode::FullCons),
            0x11 => Ok(DeserializerControlCode::DataChunk),
            0x12 => Ok(DeserializerControlCode::CodeChunk),
            _ => Err("invalid control code")
        }
    }
}

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

    #[test]
    fn test_operand_tofrom_str() {
        let cases = vec![
            ("$expr",  Operand(Address::Expr,  0)),
            ("$oper1", Operand(Address::Oper1, 0)),
            ("$oper2", Operand(Address::Oper2, 0)),
            ("$oper3", Operand(Address::Oper3, 0)),
            ("$oper4", Operand(Address::Oper4, 0)),
            ("true",   Operand(Address::Bool,  1)),
            ("false",  Operand(Address::Bool,  0)),
            ("%12",    Operand(Address::Stack, 12)),
            ("%1",     Operand(Address::Stack, 1)),
            ("@1",     Operand(Address::Instr, 1)),
            ("@12",    Operand(Address::Instr, 12)),
            ("1234",   Operand(Address::Numer, 1234)),
            ("'c'",    Operand(Address::Char,  'c' as usize)),
        ];

        for i in cases.iter() {
            let a = Operand::from_str(i.0).unwrap();
            assert_eq!(a, i.1);
            assert_eq!(i.0, a.to_string().as_str());
        }
    }

    #[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_instruction_tofrom_string() {
        let happy_cases = vec![
            "NOP",
            "    NOP",
            "NOP   ",
            "nop",
            "PUSH $expr",
            "CONST $expr 4",
            "jmp @3",
        ];

        let sad_cases = vec![
            "NOP 1",
            "push",
            "const 4",
        ];

        for i in happy_cases.iter() {
            assert!(i.parse::<Instruction>().is_ok());
        }

        for i in sad_cases.iter() {
            assert!(i.parse::<Instruction>().is_err());
        }
    }

    // TODO: test program from and to string

    #[test]
    fn test_program_parse() {
        let bytes1 = [
            DeserializerControlCode::DataChunk as u8,
            DeserializerControlCode::BoolTrue as u8,
            DeserializerControlCode::CodeChunk as u8,
            instr::LINK.0, 0xf3, 0xf4
        ];
        let out1a = vec![Datum::Bool(true)];
        let out1b = vec![Instruction(instr::LINK,
            vec![Operand(Address::Oper1, 0), Operand(Address::Oper2, 0)])];
        let res1 =
            TryInto::<Program>::try_into(&bytes1[..]);
        assert!(res1.is_ok());
        let res1 = res1.unwrap();
        assert_eq!(res1.0, out1a);
        assert_eq!(res1.1, out1b);

        let bytes2 = [
            DeserializerControlCode::CodeChunk as u8,
            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().1, out2);
    }

    #[test]
    fn test_serializer_control_code_consistency() {
        let mut input: u8 = 0x00;
        loop {
            if DeserializerControlCode::try_from(input)
                .and_then(|x| Ok(assert!(x as u8 == input)))
                    .is_err() {
                        break;
            }

            input += 1;
        }
    }
}