mod memory;
pub use memory::BrainfuckTape;

use std::{fmt, ops::Not};

#[derive(Debug, PartialEq, Clone, Copy)]
pub enum BfOpKind {
    InvalidOp,
    IncPtr,
    DecPtr,
    IncData,
    DecData,
    ReadStdin,
    WriteStdout,
    LoopSetToZero,
    LoopMovePtr,
    LoopMoveData,
    JumpIfDataZero,
    JumpIfDataNotZero,
}

impl BfOpKind {
    fn from_character(instruction: u8) -> BfOpKind {
        match instruction as char {
            '>' => BfOpKind::IncPtr,
            '<' => BfOpKind::DecPtr,
            '+' => BfOpKind::IncData,
            '-' => BfOpKind::DecData,
            ',' => BfOpKind::ReadStdin,
            '.' => BfOpKind::WriteStdout,
            '[' => BfOpKind::JumpIfDataZero,
            ']' => BfOpKind::JumpIfDataNotZero,
            _ => BfOpKind::InvalidOp,
        }
    }

    pub fn to_character(&self) -> char {
        match self {
            BfOpKind::IncPtr => '>',
            BfOpKind::DecPtr => '<',
            BfOpKind::IncData => '+',
            BfOpKind::DecData => '-',
            BfOpKind::ReadStdin => ',',
            BfOpKind::WriteStdout => '.',
            BfOpKind::JumpIfDataZero => '[',
            BfOpKind::JumpIfDataNotZero => ']',
            _ => '?',
        }
    }
}

impl Not for BfOpKind {
    type Output = BfOpKind;

    fn not(self) -> Self {
        match self {
            BfOpKind::IncData => BfOpKind::DecData,
            BfOpKind::DecData => BfOpKind::IncData,
            BfOpKind::DecPtr => BfOpKind::IncPtr,
            BfOpKind::IncPtr => BfOpKind::DecPtr,
            kind => kind,
        }
    }
}

pub struct BfOp {
    kind: BfOpKind,
    argument: isize,
}

impl fmt::Debug for BfOp {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(format!("{} {:?}", self.argument, self.kind).as_str())
    }
}

impl PartialEq for BfOp {
    fn eq(&self, other: &Self) -> bool {
        self.kind == other.kind
    }
}

impl BfOp {
    pub fn to_string(&self) -> String {
        let c = self.kind.to_character();

        match c {
            '?' | '[' | ']' => String::from(c),
            other => String::from(other).repeat(self.argument as usize)
        }
    }
}

fn optimize_loop(ops: &Vec<BfOp>, loop_start: isize) -> Vec<BfOp> {
    let mut new_ops: Vec<BfOp> = Vec::new();

    if ops.len() as isize - loop_start == 2 {
        let repeated_op = &ops[loop_start as usize + 1];
        if repeated_op.kind == BfOpKind::IncData || repeated_op.kind == BfOpKind::DecData {
            new_ops.push(BfOp {
                kind: BfOpKind::LoopSetToZero,
                argument: 0,
            });
        } else if repeated_op.kind == BfOpKind::IncPtr {
            new_ops.push(BfOp {
                kind: BfOpKind::LoopMovePtr,
                argument: repeated_op.argument,
            });
        } else if repeated_op.kind == BfOpKind::DecPtr {
            new_ops.push(BfOp {
                kind: BfOpKind::LoopMovePtr,
                argument: -repeated_op.argument,
            });
        } else if ops.len() as isize - loop_start == 5 {
            if ops[loop_start as usize + 1].kind == BfOpKind::DecData
                && ops[loop_start as usize + 3].kind == BfOpKind::IncData
                && ops[loop_start as usize + 1].argument == 1
                && ops[loop_start as usize + 3].argument == 1
            {
                if ops[loop_start as usize + 2].kind == BfOpKind::IncPtr
                    && ops[loop_start as usize + 4].kind == BfOpKind::DecPtr
                    && ops[loop_start as usize + 2].argument
                        == ops[loop_start as usize + 4].argument
                {
                    new_ops.push(BfOp {
                        kind: BfOpKind::LoopMoveData,
                        argument: ops[loop_start as usize + 2].argument,
                    });
                } else if ops[loop_start as usize + 2].kind == BfOpKind::DecPtr
                    && ops[loop_start as usize + 4].kind == BfOpKind::IncPtr
                    && ops[loop_start as usize + 2].argument
                        == ops[loop_start as usize + 4].argument
                {
                    new_ops.push(BfOp {
                        kind: BfOpKind::LoopMoveData,
                        argument: -ops[loop_start as usize + 2].argument,
                    });
                }
            }
        }
    }

    new_ops
}

pub fn parse_bytes(code: &Vec<u8>, optimize_loops: bool) -> Vec<BfOp> {
    let mut ptr = 0;
    let code_size = code.len();
    let mut ops = Vec::new();
    let mut open_bracket_stack: Vec<isize> = Vec::new();

    while ptr < code_size {
        let instruction = code[ptr];
        if instruction == '[' as u8 {
            open_bracket_stack.push(ops.len() as isize);
            ops.push(BfOp {
                kind: BfOpKind::JumpIfDataZero,
                argument: -1,
            });
            ptr += 1;
        } else if instruction == ']' as u8 {
            if !open_bracket_stack.is_empty() {
                let open_bracket_offset = open_bracket_stack[open_bracket_stack.len() - 1];
                open_bracket_stack.pop();

                let mut optimized_loop = if optimize_loops {
                    optimize_loop(&ops, open_bracket_offset)
                } else {
                    Vec::new()
                };

                if optimized_loop.is_empty() {
                    ops[open_bracket_offset as usize].argument = ops.len() as isize;
                    ops.push(BfOp {
                        kind: BfOpKind::JumpIfDataNotZero,
                        argument: open_bracket_offset,
                    });
                } else {
                    let ops_len: usize = ops.len();
                    ops.drain(open_bracket_offset as usize..ops_len);
                    ops.append(&mut optimized_loop);
                }
            }
            ptr += 1;
        } else {
            let mut kind = BfOpKind::from_character(instruction);
            let mut count: isize = 1;

            loop {
                ptr += 1;
                if ptr >= code_size {
                    break;
                }

                let next_kind = BfOpKind::from_character(code[ptr]);
                if next_kind == BfOpKind::InvalidOp {
                    continue;
                } else if next_kind == kind {
                    count += 1;
                } else if next_kind == !kind {
                    count -= 1;
                } else {
                    break;
                }
            }

            if count < 0 {
                count = count.abs();
                kind = !kind;
            }
            if count > 0 {
                match kind {
                    BfOpKind::InvalidOp => {}
                    BfOpKind::IncData | BfOpKind::DecData => {
                        ops.push(BfOp {
                            kind: kind,
                            argument: count % 256,
                        });
                    }
                    _ => ops.push(BfOp {
                        kind: kind,
                        argument: count,
                    }),
                }
            }
        }
    }

    ops
}

pub fn parse(s: String, optimize_loops: bool) -> Vec<BfOp> {
    parse_bytes(&s.into_bytes(), optimize_loops)
}

pub fn execute(
    ops: &Vec<BfOp>,
    mut input: impl FnMut() -> u8,
    mut output: impl FnMut(u8),
    should_stop: impl Fn(u128) -> bool,
) -> u128 {
    let mut ticks: u128 = 0;
    let mut code_ptr: usize = 0;
    let mut data = BrainfuckTape::new();

    while code_ptr < ops.len() && !should_stop(ticks) {
        let op = &ops[code_ptr];
        let kind = &op.kind;

        match kind {
            &BfOpKind::IncPtr => data.increase_pointer(op.argument),
            &BfOpKind::DecPtr => data.decrease_pointer(op.argument),
            &BfOpKind::IncData => data.increase_cell(op.argument as u8),
            &BfOpKind::DecData => data.decrease_cell(op.argument as u8),
            &BfOpKind::ReadStdin => {
                for _ in 0..op.argument {
                    data.set_cell(input());
                }
            }
            &BfOpKind::WriteStdout => {
                for _ in 0..op.argument {
                    output(data.get_cell());
                }
            }
            &BfOpKind::LoopSetToZero => data.set_cell(0),
            &BfOpKind::LoopMovePtr => {
                while data.get_cell() != 0 {
                    data.move_ptr(op.argument);
                }
            }
            &BfOpKind::LoopMoveData => data.move_data(op.argument),
            &BfOpKind::JumpIfDataZero => {
                if data.get_cell() == 0 && op.argument >= 0 {
                    code_ptr = op.argument as usize;
                }
            }
            &BfOpKind::JumpIfDataNotZero => {
                if data.get_cell() != 0 && op.argument >= 0 {
                    code_ptr = op.argument as usize;
                }
            }
            &BfOpKind::InvalidOp => {}
        }
        code_ptr += 1;
        ticks += 1;
    }

    ticks
}