extern crate rand;

use rand::{distributions::Alphanumeric, Rng};
use std::collections::HashMap;

#[derive(Debug, Clone)]
pub struct Function {
    pub name: String,
    pub statement: Statement,
    pub args: Vec<String>,
    assembly: Option<String>,
}

#[derive(Debug, Clone)]
pub enum Statement {
    FunctionCall(String, Vec<Expression>),
    VariableDeclaration(String, Expression),
    VariableAssignment(String, Expression),
    Scope(Vec<Statement>),
}

#[derive(Debug, Clone)]
pub enum Expression {
    Literal(Literal),
    Variable(String),
    Unary(Unary, Box<Expression>),
    Binary(BinOp, Box<Expression>, Box<Expression>),
}

#[derive(Debug, Clone, PartialEq)]
pub enum Unary {
    PointerTo,
    ValueFrom,
}

#[derive(Debug, Clone, PartialEq)]
pub enum BinOp {
    Plus,
    Minus,
}

#[derive(Debug, Clone)]
pub enum Literal {
    Int(i64),
}

#[derive(Debug, Clone)]
pub struct AssemblerState {
    pub functions: HashMap<String, Function>,
    pub variables: Vec<String>,
}

pub trait Node {
    fn to_assembly(&self, state: &mut AssemblerState) -> String;
}

impl Function {
    pub fn new(name: String, args: Vec<String>, statement: Statement) -> Function {
        Function {
            name,
            statement,
            args,
            assembly: None,
        }
    }

    pub fn add(&self, state: &mut AssemblerState) -> Function {
        state.functions.insert(self.name.clone(), self.to_owned());
        self.to_owned()
    }
}

impl AssemblerState {
    pub fn new() -> AssemblerState {
        AssemblerState {
            functions: get_default_functions(),
            variables: Vec::new(),
        }
    }
}

fn get_default_functions() -> HashMap<String, Function> {
    let mut functions = HashMap::new();
    functions.insert(
        "out".to_string(),
        Function {
            name: "out".to_string(),
            statement: Statement::FunctionCall("".to_string(), vec![]),
            args: vec!["value".to_string()],
            assembly: Some(
                "
out:
;
; Pop the argument of the stack
;

SE-OUT  RAM-ADDR-U-IN 0xFF ; Set ram address to stack end
RAM-OUT OUTPUT-IN     0xFF ; Output top address onto output register

ROM-OUT FLAGS-IN      0x80 ; Subtract mode
SE-OUT  A-IN          0xFF ; Fetch SE
ROM-OUT B-IN          0x01 ; Subtract one
ALU-OUT SE-IN         0xFF ; Store in SE
ROM-OUT FLAGS-IN      0x00 ; Add mode

;
; Pop the return address of the stack
;

SE-OUT  RAM-ADDR-U-IN 0xFF ; Set ram address to stack end
RAM-OUT ROM-ADDR-U-IN 0xFF ; Store top address in ROM jump address

ROM-OUT FLAGS-IN      0x80 ; Subtract mode
SE-OUT  A-IN          0xFF ; Fetch SE
ROM-OUT B-IN          0x01 ; Subtract one
ALU-OUT SE-IN         0xFF ; Store in SE
ROM-OUT FLAGS-IN      0x00 ; Add mode

F-OUT   JMP           0xFF"
                    .to_string(),
            ),
        },
    );
    functions.insert(
        "exit".to_string(),
        Function {
            name: "exit".to_string(),
            statement: Statement::FunctionCall("".to_string(), vec![]),
            args: vec![],
            assembly: Some(
                "
exit:
F-OUT JMP 0xFF
"
                .to_string(),
            ),
        },
    );
    functions
}

impl Node for Function {
    fn to_assembly(&self, state: &mut AssemblerState) -> String {
        match &self.assembly {
            Some(v) => v.to_string(),
            None => {
                let name = &self.name;

                let mut pop_argument_assembly = "ROM-OUT FLAGS-IN 0x80".to_string();

                for argument in &self.args {
                    state.variables.push(argument.to_string());
                    pop_argument_assembly += "
SE-OUT  A-IN  0xFF
ROM-OUT B-IN  0x01
ALU-OUT SE-IN 0xFF
"
                }

                let statement_assembly = self.statement.to_assembly(state);
                format!(
                    "{name}:
{statement_assembly}
{pop_argument_assembly}
;
; Pop the return address of the stack
;

SE-OUT  RAM-ADDR-U-IN 0xFF ; Set ram address to stack end
RAM-OUT ROM-ADDR-U-IN 0xFF ; Store top address in ROM jump address

ROM-OUT FLAGS-IN      0x80 ; Subtract mode
SE-OUT  A-IN          0xFF ; Fetch SE
ROM-OUT B-IN          0x01 ; Subtract one
ALU-OUT SE-IN         0xFF ; Store in SE
ROM-OUT FLAGS-IN      0x00 ; Add mode

F-OUT   JMP           0xFF
",
                )
            }
        }
    }
}

impl Node for Statement {
    fn to_assembly(&self, state: &mut AssemblerState) -> String {
        match self {
            Statement::Scope(statements) => {
                let dropped_length = state.variables.len();

                let statement_assembly = statements
                    .iter()
                    .map(|x| x.to_assembly(state))
                    .collect::<Vec<String>>()
                    .join("\n");

                while state.variables.len() > dropped_length {
                    state.variables.pop();
                }
                dbg!(&state.variables);

                statement_assembly
            }
            Statement::FunctionCall(name, arguments) => {
                if !state.functions.contains_key(name) {
                    todo!()
                }

                state.variables.push(name.to_owned());
                dbg!(&state.variables);
                let argument_assembly: String = arguments
                    .into_iter()
                    .map(|x| x.to_assembly(state))
                    .collect::<Vec<String>>()
                    .join("\n");

                let call_id: String = rand::thread_rng()
                    .sample_iter(&Alphanumeric)
                    .take(30)
                    .map(char::from)
                    .collect();

                for _ in &state.functions.get(name).unwrap().args {
                    state.variables.pop();
                }

                state.variables.pop();

                format!(
                    ";
; Stack
;

; Increment Stack Pointer

SE-OUT  A-IN  0xFF ; Fetch SE
ROM-OUT B-IN  0x01 ; Add 1
ALU-OUT SE-IN 0xFF ; Store

; Store in top stack pointer

ALU-OUT RAM-ADDR-U-IN 0xFF ; Store in RAM address
ROM-OUT RAM-IN        ${name}-{call_id}  ; Store address after jump in ram

;
; Arguments
; 

{argument_assembly}

;
; Jump
;
ROM-OUT ROM-ADDR-U-IN ${name}  ; Jump to function label
F-OUT   JMP           0xFF ; Jump to function

{name}-{call_id}:
"
                )
            }
            Statement::VariableDeclaration(name, value) => {
                let assembly = value.to_assembly(state);
                state.variables.pop();
                state.variables.push(name.to_owned());
                assembly
            }
            Statement::VariableAssignment(name, value) => {
                let value_assembly = value.to_assembly(state);
                state.variables.pop();

                let stack_pos = state.variables.len()
                    - state.variables.iter().position(|x| x == name).unwrap()
                    - 1;

                format!(
                    "{value_assembly}
SE-OUT  RAM-ADDR-U-IN 0xFF ; Move RAM Pointer To Correct Position
RAM-OUT C-IN          0xFF ; Move The Contents Of The RAM To The C Register

ROM-OUT FLAGS-IN 0x80 ; Subtract Mode
SE-OUT  A-IN     0xFF ; Fetch SE
ROM-OUT B-IN     0x01 ; Subtract 1
ALU-OUT SE-IN    0xFF ; Output SE-1 into SE

SE-OUT  A-IN          0xFF ; Fetch SE
ROM-OUT B-IN          {stack_pos} ; Subtract Stack Position From SE
ALU-OUT RAM-ADDR-U-IN 0xFF ; Move RAM Pointer To Correct Position
C-OUT   RAM-IN        0xFF
ROM-OUT FLAGS-IN      0x00 ; Add Mode
"
                )
            }
        }
    }
}

impl Node for Expression {
    fn to_assembly(&self, state: &mut AssemblerState) -> String {
        match self {
            Expression::Literal(v) => v.to_assembly(state),
            Expression::Binary(o, expr_1, expr_2) => {
                expr_1.to_assembly(state) + &expr_2.to_assembly(state) + &o.to_assembly(state)
            }
            Expression::Variable(name) => {
                if !state.variables.contains(name) {
                    dbg!(&name);
                    dbg!(&state.variables);
                    todo!();
                }

                let stack_position = state.variables.len()
                    - state.variables.iter().rposition(|x| x == name).unwrap()
                    - 1;

                dbg!(&state.variables);
                state.variables.push(name.to_owned());

                format!(
                    ";
; Get From Stack
;

; Get value of variable {name}

ROM-OUT FLAGS-IN      0x80 ; Subtract Mode
SE-OUT  A-IN          0xFF ; Fetch SE
ROM-OUT B-IN          {stack_position} ; Subtract Position Of Variable In Stack
ALU-OUT RAM-ADDR-U-IN 0xFF ; Set RAM Address To Position Of Variable In RAM
ROM-OUT FLAGS-IN      0x00 ; Add Mode
RAM-OUT C-IN          0xFF ; Put Variable Value In C Register

;
; Push To Top Of Stack
;

; Add One To SE

ROM-OUT B-IN          0x01
ALU-OUT SE-IN         0xFF
SE-OUT  RAM-ADDR-U-IN 0xFF
C-OUT   RAM-IN        0xFF
"
                )
            }
            Expression::Unary(operator, expr) => {
                expr.to_assembly(state) + &operator.to_assembly(state)
            }
        }
    }
}

impl Node for Literal {
    fn to_assembly(&self, state: &mut AssemblerState) -> String {
        match self {
            Literal::Int(v) => {
                state.variables.push(v.to_string());
                format!(
                    ";
; Stack
;

; Increment Stack Pointer

SE-OUT  A-IN  0xFF ; Fetch SE
ROM-OUT B-IN  0x01 ; Add 1
ALU-OUT SE-IN 0xFF ; Store

; Store in top stack pointer

ALU-OUT RAM-ADDR-U-IN 0xFF ; Store in RAM address
ROM-OUT RAM-IN        {v}; Store literal value in RAM"
                )
            }
        }
    }
}

impl Node for BinOp {
    fn to_assembly(&self, state: &mut AssemblerState) -> String {
        match self {
            BinOp::Plus => {
                let assembly = format!(
                    "
SE-OUT  RAM-ADDR-U-IN 0xFF ; Set ram address to stack end
RAM-OUT C-IN          0xFF ; Output top address onto output register

ROM-OUT FLAGS-IN      0x80 ; Subtract mode
SE-OUT  A-IN          0xFF ; Fetch SE
ROM-OUT B-IN          0x01 ; Subtract one
ALU-OUT SE-IN         0xFF ; Store in SE

SE-OUT  RAM-ADDR-U-IN 0xFF ; Set ram address to stack end
RAM-OUT D-IN          0xFF ; Output top address onto output register

ROM-OUT FLAGS-IN      0x80 ; Subtract mode
SE-OUT  A-IN          0xFF ; Fetch SE
ROM-OUT B-IN          0x01 ; Subtract one
ALU-OUT SE-IN         0xFF ; Store in SE
ROM-OUT FLAGS-IN      0x00 ; Add mode

C-OUT   A-IN          0xFF ; Move temporary values into a register for alu computation 
D-OUT   B-IN          0xFF ; Move temporary values into a register for alu computation 
ALU-OUT C-IN          0xFF ; Move ALU output into temporary register

SE-OUT  A-IN  0xFF ; Fetch SE
ROM-OUT B-IN  0x01 ; Add 1
ALU-OUT SE-IN 0xFF ; Store

ALU-OUT RAM-ADDR-U-IN 0xFF ; Store in RAM address
C-OUT   RAM-IN        0xFF ; Store address after jump in ram
",
                );
                state.variables.pop();
                assembly
            }
            BinOp::Minus => {
                format!(
                    "
SE-OUT  RAM-ADDR-U-IN 0xFF ; Set ram address to stack end
RAM-OUT C-IN          0xFF ; Output top address onto output register

ROM-OUT FLAGS-IN      0x80 ; Subtract mode
SE-OUT  A-IN          0xFF ; Fetch SE
ROM-OUT B-IN          0x01 ; Subtract one
ALU-OUT SE-IN         0xFF ; Store in SE

SE-OUT  RAM-ADDR-U-IN 0xFF ; Set ram address to stack end
RAM-OUT D-IN          0xFF ; Output top address onto output register

ROM-OUT FLAGS-IN      0x80 ; Subtract mode
SE-OUT  A-IN          0xFF ; Fetch SE
ROM-OUT B-IN          0x01 ; Subtract one
ALU-OUT SE-IN         0xFF ; Store in SE
ROM-OUT FLAGS-IN      0x00 ; Add mode

ROM-OUT FLAGS-IN      0x80 ; Subtract mode
C-OUT   B-IN          0xFF ; Move temporary values into a register for alu computation 
D-OUT   A-IN          0xFF ; Move temporary values into a register for alu computation 
ALU-OUT C-IN          0xFF ; Move ALU output into temporary register
ROM-OUT FLAGS-IN      0x00 ; Add mode

SE-OUT  A-IN  0xFF ; Fetch SE
ROM-OUT B-IN  0x01 ; Add 1
ALU-OUT SE-IN 0xFF ; Store

ALU-OUT RAM-ADDR-U-IN 0xFF ; Store in RAM address
C-OUT   RAM-IN        0xFF ; Store address after jump in ram
",
                )
            }
        }
    }
}

impl Node for Unary {
    fn to_assembly(&self, _state: &mut AssemblerState) -> String {
        todo!()
    }
}

pub fn to_assembly(functions: Vec<Function>) -> String {
    let mut final_product: String = "".to_owned();
    let mut std_function_assembly = "".to_owned();
    let mut state = AssemblerState::new();

    for std_function in state.clone().functions {
        std_function_assembly += &std_function.1.to_assembly(&mut state);
    }

    for standard_function in &functions {
        standard_function.add(&mut state);
    }

    for standard_function in &functions {
        final_product += &standard_function.to_assembly(&mut state);
    }

    final_product + &std_function_assembly
}