use std::collections::{HashMap, HashSet};

use crate::ast::*;
use crate::tokeniser::Keyword;
use crate::tokeniser::Token;

#[derive(Clone, Debug)]
enum ParseState {
    Function(Function),
    Statement(Function),
    None(Vec<Function>),
}

impl Function {
    fn push_statement(&mut self, statement: Statement) -> Function {
        self.statements.push(statement);
        self.to_owned()
    }
}

fn parse_math(tokens: &Vec<Token>) -> ValueNode {
    let mut prioritizing_order: HashMap<i64, HashSet<Token>> = HashMap::new();
    prioritizing_order.insert(0, HashSet::from([Token::Plus, Token::Minus]));

    let mut ordered_binops: Vec<Token> = Vec::new();
    let mut ordered_binops_indexes: Vec<usize> = Vec::new();

    if tokens.len() == 1 {
        return match &tokens[0] {
            Token::Number(n) => ValueNode::Literal(Literal::Int(n.parse::<i64>().unwrap())),
            Token::Identifier(identifier) => ValueNode::Variable(identifier.to_owned()),
            _ => todo!(),
        };
    }

    for j in prioritizing_order.keys() {
        let mut i = 0;
        let current_order_binops = prioritizing_order.get(j).unwrap();

        while i < tokens.len() {
            for binop in current_order_binops {
                if &tokens[i] == binop {
                    ordered_binops.push(binop.to_owned());
                    ordered_binops_indexes.push(i);
                }
            }

            i += 1;
        }
    }

    let binop = ordered_binops.pop().unwrap();
    let binop_index = ordered_binops_indexes.pop().unwrap();

    let mut a: Vec<Token> = Vec::new();
    let mut b: Vec<Token> = Vec::new();

    let mut i = 0;
    while i < tokens.len() {
        if i < binop_index {
            a.push(tokens[i].clone());
        }
        if i > binop_index {
            b.push(tokens[i].clone());
        }

        i += 1;
    }

    return match binop {
        Token::Plus => ValueNode::Operator(Operator::Add(
            Box::new(parse_math(&a)),
            Box::new(parse_math(&b)),
        )),
        Token::Minus => ValueNode::Operator(Operator::Sub(
            Box::new(parse_math(&a)),
            Box::new(parse_math(&b)),
        )),
        _ => todo!(),
    };
}

fn parse_statement(
    parse_state: &mut ParseState,
    i: &mut usize,
    tokens: &Vec<Token>,
    functions: &Vec<Function>,
) -> ParseState {
    let mut current_function = match parse_state {
        ParseState::Statement(f) => f.to_owned(),
        _ => todo!(),
    };

    if tokens[i.to_owned()] == Token::CloseSquiglyBracket {
        let mut functions = functions.to_owned();
        functions.push(current_function);

        return ParseState::None(functions);
    }

    let statement = match &tokens[i.to_owned()] {
        Token::Keyword(Keyword::Let) => {
            *i += 1;
            let name = match tokens[*i].clone() {
                Token::Identifier(identifier) => identifier,
                _ => todo!(),
            };

            *i += 1;
            match tokens[*i].clone() {
                Token::Equals => {}
                _ => todo!(),
            };

            *i += 1;

            let mut declaration_tokens: Vec<Token> = Vec::new();
            while tokens[*i] != Token::Semi {
                declaration_tokens.push(tokens[*i].clone());
                *i += 1;
            }
            Statement::VariableDeclaration(name, parse_math(&declaration_tokens))
        }
        Token::Keyword(Keyword::Assign) => {
            *i += 1;
            let name = match tokens[*i].clone() {
                Token::Identifier(identifier) => identifier,
                _ => todo!(),
            };

            *i += 1;
            match tokens[*i].clone() {
                Token::Equals => {}
                _ => todo!(),
            };

            *i += 1;

            let mut declaration_tokens: Vec<Token> = Vec::new();
            while tokens[*i] != Token::Semi {
                declaration_tokens.push(tokens[*i].clone());
                *i += 1;
            }
            Statement::VariableAssignment(name, parse_math(&declaration_tokens))
        }
        Token::Identifier(identifier) => {
            *i += 1;

            let mut declaration_tokens: Vec<Vec<Token>> = Vec::new();

            *i += 1;
            while tokens[*i] != Token::CloseParanthesis {
                let mut argument_tokens: Vec<Token> = Vec::new();
                while tokens[*i] != Token::Comma {
                    argument_tokens.push(tokens[*i].clone());
                    *i += 1;
                }
                declaration_tokens.push(argument_tokens);
                *i += 1;
            }

            *i += 1;

            let mut arguments: Vec<ValueNode> = Vec::new();

            for token_vec in &declaration_tokens {
                if token_vec == &Vec::new() {
                    continue;
                }
                arguments.push(parse_math(token_vec));
            }

            Statement::FunctionCall(identifier.to_owned(), arguments)
        }
        _other => {
            todo!()
        }
    };

    current_function.push_statement(statement);

    return ParseState::Statement(current_function);
}

pub fn ast(tokens: Vec<Token>) -> Vec<Function> {
    let mut parse_state = ParseState::None(Vec::new());
    let mut functions: Vec<Function> = Vec::new();

    let mut i = 0;
    while i < tokens.len() {
        parse_state = match &parse_state {
            ParseState::None(funcs) => {
                functions = funcs.to_owned();
                // Expect function and shit
                match &tokens[i] {
                    Token::Keyword(keyword) => match keyword {
                        Keyword::Function => {
                            ParseState::Function(Function::new("".to_string(), Vec::new()))
                        }
                        _ => todo!(),
                    },
                    _ => todo!(),
                }
            }
            ParseState::Function(func) => match &tokens[i] {
                Token::Identifier(identifier) => {
                    ParseState::Function(Function::new(identifier.to_owned(), Vec::new()))
                }
                Token::OpenSquiglyBracket => ParseState::Statement(func.to_owned()),
                _ => todo!(),
            },
            ParseState::Statement(_function) => {
                parse_statement(&mut parse_state, &mut i, &tokens, &functions)
            }
        };

        i += 1;
    }

    match parse_state {
        ParseState::None(v) => v,
        _ => todo!(),
    }
}