use derivative::Derivative;
use num::bigint::BigInt;
use num::traits::identities::Zero;
use rpds::RedBlackTreeMap;
use rpds::{list::List, vector::Vector};
use std::fmt;
use std::iter;
use std::iter::FromIterator;
use std::iter::IntoIterator;
use std::iter::Iterator;

pub type Intcode = BigInt;

#[derive(Clone, Derivative)]
#[derivative(Debug)]
pub struct IntcodeProgram {
    #[derivative(Debug(format_with = "std::fmt::Display::fmt"))]
    instruction_pointer: Intcode,
    #[derivative(Debug(format_with = "std::fmt::Display::fmt"))]
    relative_base: Intcode,
    pub error: Option<IntcodeProgramError>,
    pub halted: bool,
    pub awaiting_input: bool,
    #[derivative(Debug(format_with = "std::fmt::Display::fmt"))]
    memory: RedBlackTreeMap<Intcode, Intcode>,
    #[derivative(Debug(format_with = "std::fmt::Display::fmt"))]
    pub input: List<Intcode>,
    #[derivative(Debug(format_with = "std::fmt::Display::fmt"))]
    pub output: Vector<Intcode>,
}

impl FromIterator<Intcode> for IntcodeProgram {
    fn from_iter<I: IntoIterator<Item = Intcode>>(iter: I) -> Self {
        IntcodeProgram {
            instruction_pointer: Intcode::from(0),
            relative_base: Intcode::from(0),
            error: None,
            halted: false,
            awaiting_input: false,
            memory: iter
                .into_iter()
                .enumerate()
                .map(|(addr, val)| (Intcode::from(addr), val))
                .collect(),
            input: List::new(),
            output: Vector::new(),
        }
    }
}

pub fn intcode_to_bool(i: &Intcode) -> bool {
    *i != Intcode::from(0)
}

pub fn bool_to_intcode(i: bool) -> Intcode {
    if i {
        Intcode::from(1)
    } else {
        Intcode::from(0)
    }
}

impl IntcodeProgram {
    pub fn with_mem_0(&self, val: Intcode) -> IntcodeProgram {
        self.with_memory_set(0.into(), val)
    }

    pub fn with_noun_verb_input(&self, noun: Intcode, verb: Intcode) -> IntcodeProgram {
        self.with_memory_set(1.into(), noun)
            .with_memory_set(2.into(), verb)
    }

    pub fn with_input(&self, input: List<Intcode>) -> IntcodeProgram {
        IntcodeProgram {
            input,
            ..self.clone()
        }
    }

    pub fn with_additional_input(&self, input: List<Intcode>) -> IntcodeProgram {
        IntcodeProgram {
            input: self.input.iter().chain(input.iter()).cloned().collect(),
            ..self.clone()
        }
    }

    pub fn with_cleared_output(&self) -> IntcodeProgram {
        IntcodeProgram {
            output: Vector::new(),
            ..self.clone()
        }
    }

    pub fn execute(&self) -> Result<Vector<Intcode>, IntcodeProgramError> {
        self.run_to_termination().output_into_result()
    }

    pub fn execute_returning_memory_0(&self) -> Result<Intcode, IntcodeProgramError> {
        self.run_to_termination().memory_0_into_result()
    }

    pub fn run_to_termination(&self) -> IntcodeProgram {
        iter::successors(Some(self.clear_await_input()), |p| {
            Some(IntcodeProgram::next(&p))
        })
        .find(|p| p.halted)
        .unwrap() // successors doesn't terminate, so this will never be none.
    }

    pub fn run_to_termination_or_input(&self) -> IntcodeProgram {
        iter::successors(Some(self.clear_await_input()), |p| {
            Some(IntcodeProgram::next(&p))
        })
        .find(|p| p.halted || p.awaiting_input)
        .unwrap()
    }

    pub fn output_into_result(&self) -> Result<Vector<Intcode>, IntcodeProgramError> {
        match self.error {
            Some(ref error) => Err(error.clone()),
            None => Ok(self.output.clone()),
        }
    }

    fn with_instruction_pointer(&self, instruction_pointer: Intcode) -> IntcodeProgram {
        IntcodeProgram {
            instruction_pointer,
            ..self.clone()
        }
    }

    fn with_instruction_pointer_offset(&self, offset: usize) -> IntcodeProgram {
        IntcodeProgram {
            instruction_pointer: self.instruction_pointer.clone() + offset,
            ..self.clone()
        }
    }

    fn with_relative_base(&self, relative_base: Intcode) -> IntcodeProgram {
        IntcodeProgram {
            relative_base,
            ..self.clone()
        }
    }

    fn with_memory_set(&self, address: Intcode, value: Intcode) -> IntcodeProgram {
        IntcodeProgram {
            memory: self.memory.insert(address, value),
            ..self.clone()
        }
    }

    fn with_input_consumed(&self) -> IntcodeProgram {
        self.input
            .drop_first()
            .map(|input| IntcodeProgram {
                input,
                ..self.clone()
            })
            .unwrap_or(self.error(IntcodeProgramError::Unknown))
    }

    fn with_output(&self, print: Intcode) -> IntcodeProgram {
        IntcodeProgram {
            output: self.output.push_back(print),
            ..self.clone()
        }
    }

    fn memory_0_into_result(&self) -> Result<Intcode, IntcodeProgramError> {
        match self.error {
            Some(ref error) => Err(error.clone()),
            None => Ok(self
                .memory
                .get(&Intcode::from(0))
                .cloned()
                .unwrap_or(Intcode::from(0))),
        }
    }
    fn next(&self) -> IntcodeProgram {
        //dbg!(self);
        self.memory
            .get(&self.instruction_pointer)
            .map(|opcode| match opcode.to_radix_le(100).1[0] {
                1 => self.add(opcode),
                2 => self.multiply(opcode),
                3 => self.input(opcode),
                4 => self.output(opcode),
                5 => self.jump_if_true(opcode),
                6 => self.jump_if_false(opcode),
                7 => self.less_than(opcode),
                8 => self.equals(opcode),
                9 => self.set_relative_base(opcode),
                99 => self.halt(),
                unknown => self.error(IntcodeProgramError::InvalidOpcode(unknown.clone())),
            })
            .unwrap_or(self.error(IntcodeProgramError::Unknown))
    }

    fn add(&self, mode: &Intcode) -> IntcodeProgram {
        self.with_instruction_pointer_offset(4).with_memory_set(
            self.get_literal(3, mode),
            self.get(1, mode) + self.get(2, mode),
        )
    }

    fn multiply(&self, mode: &Intcode) -> IntcodeProgram {
        self.with_instruction_pointer_offset(4).with_memory_set(
            self.get_literal(3, mode),
            self.get(1, mode) * self.get(2, mode),
        )
    }

    fn input(&self, mode: &Intcode) -> IntcodeProgram {
        match self.input.first().cloned() {
            Some(input) => self
                .with_instruction_pointer_offset(2)
                .with_memory_set(self.get_literal(1, mode), input)
                .with_input_consumed(),
            None => self.await_input(),
        }
    }

    fn output(&self, mode: &Intcode) -> IntcodeProgram {
        self.with_instruction_pointer_offset(2)
            .with_output(self.get(1, mode))
    }

    fn jump_if_true(&self, mode: &Intcode) -> IntcodeProgram {
        if !self.get(1, mode).is_zero() {
            self.with_instruction_pointer(self.get(2, mode))
        } else {
            self.with_instruction_pointer_offset(3)
        }
    }
    fn jump_if_false(&self, mode: &Intcode) -> IntcodeProgram {
        if self.get(1, mode).is_zero() {
            self.with_instruction_pointer(self.get(2, mode))
        } else {
            self.with_instruction_pointer_offset(3)
        }
    }

    fn less_than(&self, mode: &Intcode) -> IntcodeProgram {
        self.with_instruction_pointer_offset(4).with_memory_set(
            self.get_literal(3, mode),
            if self.get(1, mode) < self.get(2, mode) {
                Intcode::from(1)
            } else {
                Intcode::from(0)
            },
        )
    }

    fn equals(&self, mode: &Intcode) -> IntcodeProgram {
        self.with_instruction_pointer_offset(4).with_memory_set(
            self.get_literal(3, mode),
            if self.get(1, mode) == self.get(2, mode) {
                Intcode::from(1)
            } else {
                Intcode::from(0)
            },
        )
    }

    fn set_relative_base(&self, mode: &Intcode) -> IntcodeProgram {
        self.with_instruction_pointer_offset(2)
            .with_relative_base(self.relative_base.clone() + self.get(1, mode))
    }

    fn halt(&self) -> IntcodeProgram {
        IntcodeProgram {
            halted: true,
            ..self.clone()
        }
    }

    fn await_input(&self) -> IntcodeProgram {
        IntcodeProgram {
            awaiting_input: true,
            ..self.clone()
        }
    }

    fn clear_await_input(&self) -> IntcodeProgram {
        IntcodeProgram {
            awaiting_input: false,
            ..self.clone()
        }
    }

    fn error(&self, error: IntcodeProgramError) -> IntcodeProgram {
        IntcodeProgram {
            halted: true,
            error: Some(error),
            ..self.clone()
        }
    }

    fn get(&self, pointer_offset: usize, mode: &Intcode) -> Intcode {
        match mode
            .to_radix_le(10)
            .1
            .get(pointer_offset + 1)
            .cloned()
            .unwrap_or(0)
        {
            0 => self.get_position(pointer_offset),
            1 => self.get_immediate(pointer_offset),
            2 => self.get_relative(pointer_offset),
            _ => Intcode::from(0),
        }
    }

    fn get_literal(&self, pointer_offset: usize, mode: &Intcode) -> Intcode {
        match mode
            .to_radix_le(10)
            .1
            .get(pointer_offset + 1)
            .cloned()
            .unwrap_or(0)
        {
            0 => self.get_immediate(pointer_offset),
            1 => self.get_immediate(pointer_offset),
            2 => self.get_immediate_relative(pointer_offset),
            _ => Intcode::from(0),
        }
    }

    fn get_immediate(&self, pointer_offset: usize) -> Intcode {
        self.memory
            .get(&(self.instruction_pointer.clone() + pointer_offset))
            .cloned()
            .unwrap_or(Intcode::from(0))
    }

    fn get_position(&self, pointer_offset: usize) -> Intcode {
        self.memory
            .get(&self.get_immediate(pointer_offset))
            .cloned()
            .unwrap_or(Intcode::from(0))
    }

    fn get_relative(&self, pointer_offset: usize) -> Intcode {
        self.memory
            .get(&(self.get_immediate(pointer_offset) + self.relative_base.clone()))
            .cloned()
            .unwrap_or(Intcode::from(0))
    }

    fn get_immediate_relative(&self, pointer_offset: usize) -> Intcode {
        self.get_immediate(pointer_offset) + self.relative_base.clone()
    }
}

#[derive(Debug, PartialEq, Clone)]
pub enum IntcodeProgramError {
    InvalidOpcode(u8),
    Unknown,
}

impl fmt::Display for IntcodeProgramError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        use IntcodeProgramError::*;

        match self {
            InvalidOpcode(i) => write!(f, "{} is not a valid opcode", i),
            Unknown => write!(f, "Unknown error"),
        }
    }
}
impl std::error::Error for IntcodeProgramError {}

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

    fn i32_vec_to_intcode_program(input: Vec<i32>) -> IntcodeProgram {
        input.into_iter().map(Intcode::from).collect()
    }

    fn i32_vec_to_intcode_memory(input: Vec<i32>) -> RedBlackTreeMap<Intcode, Intcode> {
        input
            .into_iter()
            .enumerate()
            .map(|(i, val)| (Intcode::from(i), Intcode::from(val)))
            .collect()
    }

    fn i32_vec_to_intcode_vec(input: Vec<i32>) -> Vector<Intcode> {
        input.into_iter().map(Intcode::from).collect()
    }

    fn test_example_program(
        before_execution: Vec<i32>,
        after_execution: Vec<i32>,
    ) -> IntcodeProgram {
        let program = i32_vec_to_intcode_program(before_execution).run_to_termination();

        assert_eq!(program.error, None);
        assert_eq!(program.memory, i32_vec_to_intcode_memory(after_execution));
        program
    }

    #[test]
    fn day_2_example_1() {
        test_example_program(vec![1, 0, 0, 0, 99], vec![2, 0, 0, 0, 99]);
    }

    #[test]
    fn day_2_example_2() {
        test_example_program(vec![2, 3, 0, 3, 99], vec![2, 3, 0, 6, 99]);
    }

    #[test]
    fn day_5_example_1() {
        test_example_program(vec![1002, 4, 3, 4, 33], vec![1002, 4, 3, 4, 99]);
    }

    #[test]
    fn day_9_example_1() {
        let quine = vec![
            109, 1, 204, -1, 1001, 100, 1, 100, 1008, 100, 16, 101, 1006, 101, 0, 99,
        ];
        let program = i32_vec_to_intcode_program(quine.clone()).run_to_termination();
        assert_eq!(program.error, None);
        assert_eq!(program.output, i32_vec_to_intcode_vec(quine));
    }
}

pub fn sort_by_key<T, K: Ord>(
    iter: impl IntoIterator<Item = T>,
    f: impl FnMut(&T) -> K,
) -> impl Iterator<Item = T> {
    let mut tmp: Vec<T> = iter.into_iter().collect();
    tmp.sort_by_key(f);
    tmp.into_iter()
}