use core::panic;

use crate::{glk::Glk, instructions::{Instruction, Operand}, opcodes::OpCode};

pub type MemAddress = u32;
pub type StackAddress = u32;
pub type FunctionPointer = MemAddress;

#[repr(u32)]
enum StubDestType {
    DoNotStore = 0,
    MainMemory = 1,
    LocalVariable = 2,
    PushOnStack = 3,
    ResumeCompressedString = 10,
    ResumeFunctionFromString = 11,
    ResumeInteger = 12,
    ResumeCString = 13,
    ResumeUnicodeString = 14,
}

enum ReadWriteDest {
    Stack,
    Memory,
}

pub struct HeapChunk {
    pub start_address: MemAddress,
    pub length: u32,
    pub contents: Option<Vec<u32>>
}

pub enum Output {
    Null,
    Filter(FunctionPointer),
    Glk(Glk, u32),
}

pub struct Interpreter {
    pub memory: Vec<u8>,
    pub stack: Vec<u8>,
    pub mem_writeable: MemAddress,
    pub mem_min: MemAddress,
    pub heap: Vec<HeapChunk>,

    pub reg_program: MemAddress,
    pub reg_stack: StackAddress,
    pub reg_frame: StackAddress,
    pub reg_strdec: MemAddress,

    pub output_mode: Output,
    tmp_dest_oper: Option<Operand>,
}

impl Interpreter {
    pub fn init(
        mut source_file: Vec<u8>,
        ram_start: u32,
        ext_start: u32,
        ram_end: u32,
        stack_size: u32,
        str_decode: u32,
    ) -> Self {
        // Trick to ensure everything after ext_start is cleared
        // even if the original file was longer
        source_file.resize(ext_start as usize, 0);
        source_file.resize(ext_start as usize, 0);

        let mut stack = Vec::new();
        stack.resize(stack_size as usize, 0);
        
        Self {
            memory: source_file,
            stack,
            mem_writeable: ram_start,
            mem_min: ram_end,
            heap: Vec::new(),

            reg_program: 0,
            reg_stack: 0,
            reg_frame: 0,
            reg_strdec: str_decode,

            output_mode: Output::Null,
            tmp_dest_oper: None,
        }
    }

    pub fn run(&mut self, entry_func: MemAddress) {
        self.reg_program = self.create_stack_frame(entry_func, Vec::new());

        loop {
            let pc = self.reg_program as usize;
            let mut instruction_slice = &self.memory[pc..pc+100];
            let instruction_read = Instruction::read(&mut instruction_slice);
            if let Err(unknown_opcode) = instruction_read {
                panic!("Unknown opcode 0x{unknown_opcode:03x} at 0x{pc:08x}")
            }
            let (instruction, bytes_read) = instruction_read.unwrap();
            self.reg_program += bytes_read;

            println!("Executing instruction {instruction:?}");

            let in_operands = instruction.in_operands.iter().map(|op| self.eval_source_operand(&op)).collect();
            let no_store = OpCode::should_gen_return_stub(instruction.opcode.clone());
            self.tmp_dest_oper = instruction.out_operands.get(0).cloned();

            let result_values = OpCode::call_opcode(self, instruction.opcode, in_operands);

            if !no_store {
                instruction.out_operands.iter().enumerate().for_each(|(index, op)| self.store_dest_operand(&op, result_values[index]));
            }
        }
    }

    pub fn call_function(&mut self, func_addr: FunctionPointer, args: Vec<u32>) {
        let new_program_counter = self.create_stack_frame(func_addr, args);
        println!("Jumping to function at 0x{new_program_counter:04x}");
        self.reg_program = new_program_counter;
    }

    pub fn create_return_stub(&mut self) {
        let dest_oper = self.tmp_dest_oper.clone();
        
        let (dest_type, dest_addr) = match dest_oper {
            None => (StubDestType::DoNotStore, 0),
            Some(Operand::Constant(_)) => (StubDestType::DoNotStore, 0),
            Some(Operand::PopStack) => (StubDestType::PushOnStack, 0),
            Some(Operand::FrameIndirect(offset)) => (StubDestType::LocalVariable, offset),
            Some(Operand::MemIndirect(offset)) => (StubDestType::MainMemory, offset),
            Some(Operand::MemRamIndirect(offset)) => (StubDestType::MainMemory, self.mem_writeable + offset),
        };

        self.push_stack(dest_type as u32);
        self.push_stack(dest_addr);
        self.push_stack(self.reg_program);
        self.push_stack(self.reg_frame);
    }

    // Returns program counter after parsing all the args
    pub fn create_stack_frame(&mut self, func: FunctionPointer, mut args: Vec<u32>) -> MemAddress {

        // Reading function signature
        let func_type = self.get_byte(func, ReadWriteDest::Memory);
        let args_in_locals = match func_type {
            0xC0 => false,
            0xC1 => true,
            _ => panic!("Tried to call object that is not a function: {}", func_type),
        };

        let mut local_pointer = func + 1;
        let mut format_of_locals = Vec::<(u8, u8)>::new();
        loop {
            let format = (
                self.get_byte(local_pointer, ReadWriteDest::Memory),
                self.get_byte(local_pointer + 1, ReadWriteDest::Memory),
            );

            local_pointer += 2;

            if format == (0,0) {
                break
            }

            format_of_locals.push(format);
        };

        let opcode_start = local_pointer;

        // Writing stack frame
        self.reg_frame = self.reg_stack;
        let mut header_len = 4 + 4 + format_of_locals.len() as StackAddress * 2;
        Interpreter::align_ptr(&mut header_len, 16);

        // Write locals pos
        self.write_word(self.reg_frame + 4, ReadWriteDest::Stack, header_len);

        // Write format of locals
        let mut local_format_ptr = self.reg_frame + 8;
        for (local_type, local_count) in format_of_locals.iter() {
            self.write_byte(local_format_ptr, ReadWriteDest::Stack, *local_type);
            self.write_byte(local_format_ptr + 1, ReadWriteDest::Stack, *local_count);
            local_format_ptr += 2;
        }

        // Write locals
        let mut local_pos = self.reg_frame+ header_len;
        for (local_type, local_count) in format_of_locals {
            Interpreter::align_ptr(&mut local_pos, local_type * 8);

            for _count in 0..local_count {
                let value = match args_in_locals {
                    true => if args.len() > 0 {args.remove(0)} else {0}
                    false => 0,
                };

                match local_type {
                    1 => self.write_byte(local_pos, ReadWriteDest::Stack, value as u8),
                    2 => self.write_twobyte(local_pos, ReadWriteDest::Stack, value as u16),
                    4 => self.write_word(local_pos, ReadWriteDest::Stack, value),
                    _ => panic!("Invalid local type {local_type}")
                }

                local_pos += local_type as u32;
            }
        }
        
        // Write frame len
        Interpreter::align_ptr(&mut local_pos, 32);
        self.write_word(self.reg_frame, ReadWriteDest::Stack, local_pos - self.reg_frame);
        self.reg_stack = local_pos;

        // Handle stack args
        if !args_in_locals {
            args.reverse();
            for arg in args {
                self.push_stack(arg);
            }
        }

        return opcode_start;
    }

    pub fn return_from_function(&mut self, value: u32) {
        // Discard frame
        self.reg_stack = self.reg_frame - 4;

        let frame_pointer = self.pop_stack();
        let program_counter = self.pop_stack();
        let dest_addr = self.pop_stack();
        let dest_type = self.pop_stack();
        let dest_type : StubDestType = unsafe {std::mem::transmute(dest_type)};

        self.reg_frame = frame_pointer;
        self.reg_program = program_counter;

        let locals_offset = self.get_word(self.reg_frame + 4, ReadWriteDest::Stack);

        match dest_type {
            StubDestType::DoNotStore => (),
            StubDestType::MainMemory => self.write_word(dest_addr, ReadWriteDest::Memory, value),
            StubDestType::LocalVariable => self.write_word(self.reg_frame + locals_offset, ReadWriteDest::Stack, value),
            StubDestType::PushOnStack => self.push_stack(value),
            StubDestType::ResumeCompressedString => todo!(),
            StubDestType::ResumeFunctionFromString => todo!(),
            StubDestType::ResumeInteger => todo!(),
            StubDestType::ResumeCString => todo!(),
            StubDestType::ResumeUnicodeString => todo!(),
        }
    }

    pub fn jump_with_offset(&mut self, offset: u32) {
        if offset == 1 || offset == 0 {
            return self.return_from_function(offset);
        }

        let offset = offset as i32;
        let destination = if offset < 0 {
            self.reg_program - -offset as u32
        } else {
            self.reg_program + offset as u32
        };

        println!("Jumping {offset} to 0x{destination:04x}");
        self.reg_program = destination;
    }

    pub fn push_stack(&mut self, value: u32) {
        self.write_word(self.reg_stack, ReadWriteDest::Stack, value);
        self.reg_stack += 4;
    }

    pub fn pop_stack(&mut self) -> u32 {
        self.reg_stack -= 4;
        self.get_word(self.reg_stack, ReadWriteDest::Stack)
    }

    fn eval_source_operand(&mut self, operand: &Operand) -> u32 {
        match operand {
            Operand::Constant(value) => *value as u32,
            Operand::PopStack => self.pop_stack(),
            Operand::FrameIndirect(offset) => {
                let locals_pos = self.get_word(self.reg_frame + 4, ReadWriteDest::Stack);
                assert!(self.reg_frame + locals_pos + offset < self.reg_stack, "Frame offset out of frame");
                self.get_word(self.reg_frame + locals_pos + offset, ReadWriteDest::Stack)
            },
            Operand::MemIndirect(address) => {
                // TODO: Handle heap
                assert!(*address < self.memory.len() as u32 - 4, "Memory offset out of memory");
                self.get_word(*address, ReadWriteDest::Memory)
            },
            Operand::MemRamIndirect(offset) => {
                let address = self.mem_writeable + offset;
                // TODO: Handle heap
                assert!(address < self.memory.len() as u32 - 4, "Memory offset out of memory");
                self.get_word(address, ReadWriteDest::Memory)
            },
        }
    }

    fn store_dest_operand(&mut self, operand: &Operand, value: u32) {
        match operand {
            Operand::Constant(0) => (),
            Operand::Constant(constant_dst) => panic!("Cannot store to constant operand {constant_dst}"),
            Operand::PopStack => self.push_stack(value),
            Operand::FrameIndirect(offset) => {
                let locals_pos = self.get_word(self.reg_frame + 4, ReadWriteDest::Stack);
                assert!(self.reg_frame + locals_pos + offset < self.reg_stack, "Frame offset out of frame");
                self.write_word(self.reg_frame + locals_pos + offset, ReadWriteDest::Stack, value)
            },
            Operand::MemIndirect(address) => {
                assert!(*address >= self.mem_writeable, "Cannot write to ROM");
                // TODO: Handle heap
                assert!(*address < self.memory.len() as u32 - 4, "Memory offset out of memory");
                self.write_word(*address, ReadWriteDest::Memory, value)
            },
            Operand::MemRamIndirect(offset) => {
                let address = self.mem_writeable + offset;
                // TODO: Handle heap
                assert!(address < self.memory.len() as u32 - 4, "Memory offset out of memory");
                self.write_word(address, ReadWriteDest::Memory, value)
            },
        }
    }

    fn get_byte(&self, addr: u32, src: ReadWriteDest) -> u8 {
        let addr = addr as usize;
        let buffer = match src {
            ReadWriteDest::Stack => &self.stack,
            ReadWriteDest::Memory => &self.memory,
        };

        buffer[addr]
    }

    fn get_twobyte(&self, addr: u32, src: ReadWriteDest) -> u16 {
        let addr = addr as usize;
        let buffer = match src {
            ReadWriteDest::Stack => &self.stack,
            ReadWriteDest::Memory => &self.memory,
        };

        u16::from_be_bytes(buffer[addr..addr+2].try_into().unwrap())
    }

    fn get_word(&self, addr: u32, src: ReadWriteDest) -> u32 {
        let addr = addr as usize;
        let buffer = match src {
            ReadWriteDest::Stack => &self.stack,
            ReadWriteDest::Memory => &self.memory,
        };

        u32::from_be_bytes(buffer[addr..addr+4].try_into().unwrap())
    }

    fn write_byte(&mut self, addr: u32, dst: ReadWriteDest, value: u8) {
        let addr = addr as usize;
        let buffer = match dst {
            ReadWriteDest::Stack => &mut self.stack,
            ReadWriteDest::Memory => &mut self.memory,
        };

        buffer[addr] = value;
    }

    fn write_twobyte(&mut self, addr: u32, dst: ReadWriteDest, value: u16) {
        let addr = addr as usize;
        let buffer = match dst {
            ReadWriteDest::Stack => &mut self.stack,
            ReadWriteDest::Memory => &mut self.memory,
        };

        buffer[addr..addr+2].copy_from_slice(&value.to_be_bytes());
    }

    fn write_word(&mut self, addr: u32, dst: ReadWriteDest, value: u32) {
        let addr = addr as usize;
        let buffer = match dst {
            ReadWriteDest::Stack => &mut self.stack,
            ReadWriteDest::Memory => &mut self.memory,
        };

        buffer[addr..addr+4].copy_from_slice(&value.to_be_bytes());
    }

    fn align_ptr(ptr: &mut u32, align_size: u8) {
        let modulo = match align_size {
            8 => 1,
            16 => 2,
            32 => 4,
            _ => panic!("Invalid alignment: {align_size}"),
        };

        let bytes_off = *ptr % modulo;
        *ptr += bytes_off;
    }
}