#[derive(Debug, Clone, PartialEq, Eq)]
enum TokenType {
    Lparen,
    Rparen,
    Lbrace,
    Rbrace,
    Lbracket,
    Rbracket,
    Arrow,
    Comma,
    Equals,
    Semicolon,
    Colon,
    Identifier(String),
    Eof,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct Blame {
    line: usize,
    char: usize,
}

#[derive(Debug, PartialEq, Eq)]
struct Token {
    token: TokenType,
    blame: Blame,
}

struct Lexer;

impl Lexer {
    pub fn analyze(program: &str) -> Vec<Token> {
        let mut tokens: Vec<Token> = program
            .lines()
            .enumerate()
            .flat_map(|(line, text)| Self::lex_line(line, text))
            .collect();

        tokens.push(Token {
            token: TokenType::Eof,
            blame: Blame {
                line: program.lines().count(),
                char: 0,
            },
        });
        tokens
    }

    fn lex_line(line: usize, text: &str) -> Vec<Token> {
        let mut tokens = Vec::new();
        let mut ident_start: Option<usize> = None;

        let flush = |tokens: &mut Vec<Token>, ident_start: &mut Option<usize>, end: usize| {
            if let Some(start) = ident_start.take() {
                if end > start {
                    tokens.push(Token {
                        token: TokenType::Identifier(text[start..end].to_string()),
                        blame: Blame { line, char: start },
                    });
                }
            }
        };

        for (i, c) in text.char_indices() {
            let lookahead = text.chars().nth(i.saturating_add(1));
            let singlechar = match c {
                '(' => Some(TokenType::Lparen),
                ')' => Some(TokenType::Rparen),
                '[' => Some(TokenType::Lbracket),
                ']' => Some(TokenType::Rbracket),
                '{' => Some(TokenType::Lbrace),
                '}' => Some(TokenType::Rbrace),
                ':' => Some(TokenType::Colon),
                ';' => Some(TokenType::Semicolon),
                '-' => {
                    Some(TokenType::Arrow).take_if(|_| lookahead.filter(|it| *it == '>').is_some())
                }
                ',' => Some(TokenType::Comma),
                '=' => Some(TokenType::Equals),
                _ => None,
            };

            if c == ' ' || c == '\t' {
                flush(&mut tokens, &mut ident_start, i);
                continue;
            }

            if let Some(tt) = singlechar {
                flush(&mut tokens, &mut ident_start, i);
                tokens.push(Token {
                    token: tt,
                    blame: Blame { line, char: i },
                });
                continue;
            }

            if ident_start.is_none() {
                ident_start = Some(i);
            }
        }

        flush(&mut tokens, &mut ident_start, text.len());
        tokens
    }
}

enum Atomic {
    Unit,
    Bool,
    Int,
}

enum SimpleType {
    Arrow((Box<SimpleType>, Box<SimpleType>)),
    Atom(Atomic),
}

type Identifier = String;

struct Binding {
    name: Identifier,
    arg_type: Option<Box<TypeDeclaration>>, // TODO: Gradual typing!
}

enum TypeDeclaration {
    BaseType(Identifier),
    Arrow(Box<(TypeDeclaration, TypeDeclaration)>),
}

struct Assignment {
    binding: Binding,
    computation: Computation,
}

struct Abstraction {
    arguments: Vec<Binding>,
    body: Box<Computation>,
}

struct Application {
    apply: Box<(Computation, Computation)>,
}

enum Computation {
    Application(Application),
    Abstraction(Abstraction),
    Identifier(Identifier),
    LetStatement(LetStatement), // TODO: I think we can flatten these to just a long Abstraction
}

enum ASTNode {
    Binding(Binding),
    TypeDeclaration(TypeDeclaration),
    Assignment(Assignment),
    Abstraction(Abstraction),
    Application(Application),
    Computation(Computation),
    LetStatement(LetStatement),
}

struct Parser {
    position: usize,
    tokens: Vec<Token>,
}

#[derive(Debug)]
enum ParsingFailure {
    Blame(Blame, String),
    PastTokens,
}

impl Parser {
    fn peek(&mut self) -> Result<&Token, ParsingFailure> {
        return self
            .tokens
            .get(self.position + 1)
            .ok_or(ParsingFailure::PastTokens);
    }

    fn consume_if_and_map<T>(
        &mut self,
        f: fn(&Token) -> Option<T>,
        msg: String,
    ) -> Result<T, ParsingFailure> {
        if self.position == self.tokens.len() {
            return Result::Err(ParsingFailure::PastTokens);
        }
        let val = f(&self.tokens[self.position]);
        if val.is_some() {
            self.position += 1;
            return Result::Ok(val.unwrap());
        }
        return Result::Err(ParsingFailure::Blame(self.tokens[self.position].blame, msg));
    }

    fn consume_if_match(
        &mut self,
        f: fn(&Token) -> bool,
        msg: String,
    ) -> Result<&Token, ParsingFailure> {
        if self.position == self.tokens.len() {
            return Result::Err(ParsingFailure::PastTokens);
        }
        if f(&self.tokens[self.position]) {
            self.position += 1;
            return Result::Ok(&self.tokens[self.position - 1]);
        }
        return Result::Err(ParsingFailure::Blame(self.tokens[self.position].blame, msg));
    }

    fn consume_if_match_nomsg(&mut self, f: fn(&Token) -> bool) -> Result<&Token, ParsingFailure> {
        return self.consume_if_match(f, String::from("Unexpected error"));
    }

    fn parse_type_declaration(&mut self) -> Result<TypeDeclaration, ParsingFailure> {}

    fn parse_binding(&mut self) -> Result<Binding, ParsingFailure> {
        let name = self.consume_if_and_map(
            |it| match ((*it).token).clone() {
                // unfortunately, gotta clone to make the borrow checker happy. maybe i'll come back to this :3
                TokenType::Identifier(name) => Some(name),
                _ => None,
            },
            String::from("Expected identifier"),
        )?;

        let val = self.consume_if_match_nomsg(|it| (*it).token == TokenType::Colon);
        if val.is_ok() {
            let blame = val.unwrap().blame;
            if let Ok(type_decl) = self.parse_type_declaration() {
                return Ok(Binding {
                    name,
                    arg_type: Some(Box::new(type_decl)),
                });
            }
            return Err(ParsingFailure::Blame(
                blame,
                String::from("Expected type declaration"),
            ));
        }
        return Ok(Binding {
            name,
            arg_type: None,
        });
    }

    fn parse_assignment(&mut self) -> Result<Assignment, ParsingFailure> {
        let binding = self.parse_binding()?;
        let _ = self.consume_if_match(
            |it| (*it).token == TokenType::Equals,
            String::from("Expected equals after binding"),
        )?;

        return Ok(Assignment {
            binding,
            computation: self.parse_computation()?,
        });
    }

    fn parse_let_statement(&mut self) -> Result<LetStatement, ParsingFailure> {
        self.consume_if_match_nomsg(|it| (*it).token == TokenType::Identifier(String::from("let")));
        let mut assignments: Vec<Assignment> = Vec::new();
        loop {
            if let Ok(_done) = self.consume_if_match_nomsg(|it| {
                (*it).token == TokenType::Identifier(String::from("in"))
            }) {
                let body = self.parse_computation()?;

                return Ok(LetStatement { assignments });
            }
            let assignment = self.parse_assignment()?;
            self.consume_if_match_nomsg(|it| (*it).token == TokenType::Semicolon)?;
            assignments.push(assignment);
        }
    }

    fn parse_abstraction(&mut self) -> Result<Abstraction, ParsingFailure> {
        self.consume_if_match_nomsg(|it| (*it).token == TokenType::Lbracket)?;
        let mut arguments: Vec<Binding> = Vec::new();
        loop {
            if let Ok(_done) = self.consume_if_match_nomsg(|it| (*it).token == TokenType::Rbracket)
            {
                return Ok(Abstraction {
                    arguments,
                    body: Box::new(self.parse_computation()?),
                });
            }
            arguments.push(self.parse_binding()?);
            self.consume_if_match_nomsg(|it| (*it).token == TokenType::Comma);
        }
    }

    fn parse_application(&mut self) -> Result<Application, ParsingFailure> {
        self.consume_if_match_nomsg(|it| (*it).token == TokenType::Lparen)?;
        let left = self.parse_computation()?;
        let right = self.parse_computation()?;
        self.consume_if_match_nomsg(|it| (*it).token == TokenType::Rparen)?;
        return Ok(Application {
            apply: Box::new((left, right)),
        });
    }

    fn parse_computation(&mut self) -> Result<Computation, ParsingFailure> {
        let next = *(self.peek()?);
        match next.token {
            TokenType::Identifier(identifier) => {
                Computation::LetStatement(self.parse_let_statement()?)
            }
            TokenType::Lbrace => Computation::Abstraction(self.parse_abstraction()?),
            TokenType::Lparen => Computation::Application(self.parse_application()?),
        };
    }

    fn new(tokens: Vec<Token>) -> Parser {
        return Parser {
            position: 0,
            tokens,
        };
    }
    //    fn new() -> Parser {}
}

pub enum StepResult<T> {
    Terminal(T),
    Continue,
}

pub trait Steppable<T> {
    fn small_step(&mut self) -> StepResult<T>;
}

pub struct Environment<'a, T> {
    parent_scope: Option<&'a Environment<'a, T>>,
    capture: Identifier,
    substitution: &'a T,
}

impl<T> Environment<'_, T> {
    fn add<'a>(&'a self, identifier: Identifier, substitution: &'a T) -> Environment<'a, T> {
        return Environment {
            parent_scope: Some(&self),
            capture: identifier,
            substitution: substitution,
        };
    }

    fn get<'a>(&'a self, identifier: Identifier) -> Option<&'a T> {
        if *identifier == *(self.capture) {
            return Some(self.substitution);
        }
        return match self.parent_scope {
            Some(scope) => (*scope).get(identifier),
            None => None,
        };
    }

    fn root<'a>(identifier: Identifier, substitution: &'a T) -> Environment<'a, T> {
        return Environment {
            parent_scope: None,
            capture: identifier,
            substitution: substitution,
        };
    }
}

fn main() {}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn test_env() {
        let ident = Identifier::from("ident");
        let ident2 = Identifier::from("ident2");
        let val: i32 = 1;
        let val2: i32 = 2;
        let val3: i32 = 3;
        let env = Environment::root(ident.clone(), &val);
        assert_eq!(env.get(ident.clone()), Some(&val));
        let env2 = env.add(ident2.clone(), &val2);
        assert_eq!(env2.get(ident2.clone()), Some(&val2));
        let env3 = env2.add(ident.clone(), &val3);
        assert_eq!(env3.get(ident2.clone()), Some(&val2));
        assert_eq!(env3.get(ident.clone()), Some(&val3));
        assert_eq!(env.get(ident.clone()), Some(&val));
    }

    #[test]
    fn test_lex() {
        let prog = "let identity: ((Int) \n-> (Int)) = [a] { a };";
        let analysis = Lexer::analyze(prog);
        assert_eq!(
            analysis[0],
            Token {
                token: TokenType::Identifier(String::from("let")),
                blame: Blame { line: 0, char: 0 }
            }
        );
        assert_eq!(
            analysis[1],
            Token {
                token: TokenType::Identifier(String::from("identity")),
                blame: Blame { line: 0, char: 4 }
            }
        );
        assert_eq!(
            analysis[2],
            Token {
                token: TokenType::Colon,
                blame: Blame { line: 0, char: 12 }
            }
        );
        assert_eq!(
            analysis[7],
            Token {
                token: TokenType::Arrow,
                blame: Blame { line: 1, char: 0 }
            }
        );
        assert_eq!(
            analysis[21],
            Token {
                token: TokenType::Eof,
                blame: Blame { line: 2, char: 0 }
            }
        );
    }

    #[test]
    fn test_parse() {
        let prog = "let Identity: ((Int) -> (Int)) = [a] { a };";
        let tokens = Lexer::analyze(prog);
        let mut parser = Parser::new(tokens);
        let ast = parser.parse().ok();
        ast.unwrap();
    }
}