use std::{collections::HashMap, iter::Peekable, rc::Rc, slice::Iter, str::Chars, vec::IntoIter};

use crate::{
    environment::Environment,
    node::{
        NodeEnum,
        add::Add,
        assign::Assign,
        call::Call,
        constant::{Constant, ConstantValue},
        divide::Divide,
        equals::Equals,
        function::{Function, FunctionType},
        multiply::Multiply,
        set::Set,
        subtract::Subtract,
        symbol::Symbol,
    },
};

#[derive(Debug)]
pub struct Token(usize, TokenType);

#[derive(Debug, PartialEq, Clone)]
pub enum TokenType {
    // Space,
    Number(ConstantValue),
    Identifier(String),

    Plus,
    Minus,
    Star,
    Slash,

    Equals,

    ColonEquals,
    LeftArrow,

    RParen,
    LParen,
    Comma,

    If,
    Then,
    Else,
    End,
}

impl TokenType {
    pub fn len(&self) -> usize {
        match self {
            TokenType::Number(n) => n.to_string().len(),
            TokenType::Identifier(s) => s.len(),

            TokenType::Plus => 1,
            TokenType::Minus => 1,
            TokenType::Star => 1,
            TokenType::Slash => 1,
            TokenType::Equals => 1,

            TokenType::ColonEquals => 2,
            TokenType::LeftArrow => 2,

            TokenType::RParen => 1,
            TokenType::LParen => 1,
            TokenType::Comma => 1,

            TokenType::If => 2,
            TokenType::Then => 4,
            TokenType::Else => 4,
            TokenType::End => 3,
        }
    }
}

pub struct Lexer<'a> {
    source: Peekable<Chars<'a>>,
}

#[derive(Debug)]
pub enum LexerError {
    UnexpectedChar(usize, String),
}

impl<'a> Lexer<'a> {
    pub fn new(source: &'a String) -> Self {
        Self {
            source: source.chars().peekable(),
        }
    }

    pub fn lex(&'a mut self) -> Result<Vec<Token>, LexerError> {
        let mut i = 0;
        let mut tokens = vec![];

        while let Some(c) = self.source.next() {
            match c {
                // Collapse spaces into a single Space token
                ' ' => {
                    while self.source.peek() == Some(&' ') {
                        self.source.next();
                        i += 1;
                    }
                    // tokens.push(Token(i, TokenType::Space));
                }
                // Comments with //
                '/' if self.source.peek() == Some(&'/') => {
                    while self.source.next() != Some('\n') {
                        i += 1;
                    }
                }
                // Numbers with decimal points
                '0'..='9' | '.' => {
                    let mut digit = String::from(c);
                    loop {
                        let d = self.source.peek();
                        let mut has_decimal = c == '.';
                        match d {
                            Some('0'..='9') => {
                                digit.push(*d.unwrap());
                                self.source.next();
                                i += 1;
                            }
                            #[allow(unused_assignments)] // For some reason it thinks has_decimal
                            // is never read
                            Some('.') => {
                                if has_decimal {
                                    return Err(LexerError::UnexpectedChar(
                                        i,
                                        "Invalid digit with multiple decimal points".into(),
                                    ));
                                }

                                digit.push(*d.unwrap());
                                has_decimal = true;
                            }
                            _ => {
                                break;
                            }
                        }
                    }
                    let number = digit.parse::<ConstantValue>().unwrap();

                    tokens.push(Token(i, TokenType::Number(number)));
                }

                // LeftArrow (->)
                '-' if self.source.peek() == Some(&'>') => {
                    self.source.next();
                    i += 1;
                    tokens.push(Token(i, TokenType::LeftArrow));
                }

                '+' => tokens.push(Token(i, TokenType::Plus)),
                '-' => tokens.push(Token(i, TokenType::Minus)),
                '*' => tokens.push(Token(i, TokenType::Star)),
                '/' => tokens.push(Token(i, TokenType::Slash)),
                '=' => tokens.push(Token(i, TokenType::Equals)),
                ',' => tokens.push(Token(i, TokenType::Comma)),

                ':' if self.source.peek() == Some(&'=') => {
                    self.source.next();
                    i += 1;
                    tokens.push(Token(i, TokenType::ColonEquals));
                }

                '(' => tokens.push(Token(i, TokenType::LParen)),
                ')' => tokens.push(Token(i, TokenType::RParen)),

                _ if c.is_alphabetic() || c == '_' => {
                    tokens.push(self.lex_identifier(&mut i, c)?);
                }

                _ => {
                    return Err(LexerError::UnexpectedChar(
                        i,
                        format!("Unexpected char {}", c),
                    ));
                }
            }
            i += 1;
        }

        Ok(tokens)
    }

    fn lex_identifier<'b>(&'b mut self, i: &mut usize, c: char) -> Result<Token, LexerError> {
        let mut identifier = c.to_string();

        while let Some(c) = self.source.peek() {
            if c.is_alphanumeric() || c == &'_' || c == &'\'' {
                identifier.push(*c);
                self.source.next();
                *i += 1;
            } else {
                break;
            }
        }

        Ok(Token(
            *i,
            match identifier.to_lowercase().as_str() {
                "if" => TokenType::If,
                "then" => TokenType::Then,
                "else" => TokenType::Else,
                "end" => TokenType::End,
                _ => TokenType::Identifier(identifier),
            },
        ))
    }
}

pub enum ParserError {
    UnexpectedEndOfTokens(String),
    UnexpectedToken(usize, usize, String),
    Unimplemented(usize, usize, String),
    UnexpectedNode(usize, String),
}

/// Recursive descent parser
pub struct Parser<'a> {
    tokens: Peekable<IntoIter<Token>>,
    environment: &'a mut Environment,
    previous: Option<&'a Token>,
}

type Tokens<'a> = Peekable<Iter<'a, Token>>;

impl<'a> Parser<'a> {
    pub fn new(tokens: Vec<Token>, env: &'a mut Environment) -> Self {
        // #[cfg(debug_assertions)]
        // println!("\r{tokens:?}");
        Self {
            tokens: tokens.into_iter().peekable(),
            environment: env,
            previous: None,
        }
    }

    // Parse tokens recursively and descendentantly
    pub fn parse(&mut self) -> Result<Vec<Rc<NodeEnum>>, ParserError> {
        let mut expressions = vec![];

        while self.tokens.peek().is_some() {
            expressions.push(self.expression()?);
        }

        Ok(expressions)
    }

    #[inline]
    fn matchType<'b>(&'b mut self, t: TokenType) -> bool {
        if let Some(Token(_, token_type)) = self.tokens.peek() {
            if *token_type == t {
                self.tokens.next();
                return true;
            }
        }

        false
    }

    #[inline]
    fn matchOrErr(&mut self, t: TokenType) -> Result<bool, ParserError> {
        let (i, tt) = if let Some(Token(i, tt)) = self.tokens.peek() {
            (*i, tt.clone())
        } else {
            return Err(ParserError::UnexpectedEndOfTokens(format!(
                "Expected {t:?} but found nothing instead"
            )));
        };

        if self.matchType(t.clone()) {
            Ok(true)
        } else {
            Err(ParserError::UnexpectedToken(
                i,
                tt.len(),
                format!("Expected {t:?} but found {tt:?} instead"),
            ))
        }
    }

    fn expression(&mut self) -> Result<Rc<NodeEnum>, ParserError> {
        self.assignment()
    }

    fn assignment(&mut self) -> Result<Rc<NodeEnum>, ParserError> {
        let expr = self.equality();

        if let Some(Token(_, TokenType::ColonEquals)) = self.tokens.peek() {
            self.tokens.next();
            return Ok(Rc::new(Assign::new(expr?, self.equality()?).into()));
        }

        expr
    }

    fn equality(&mut self) -> Result<Rc<NodeEnum>, ParserError> {
        // TODO: Implement equality
        let expr = self.comparison();
        if self.matchType(TokenType::Equals) {
            return Ok(Rc::new(Equals::new(expr?, self.equality()?).into()));
        }
        expr
    }

    fn comparison(&mut self) -> Result<Rc<NodeEnum>, ParserError> {
        // TODO: Implement comparison
        self.term()
    }

    fn term(&mut self) -> Result<Rc<NodeEnum>, ParserError> {
        let expr = self.factor()?;
        if self.matchType(TokenType::Plus) {
            Ok(Rc::new(Add::new(expr, self.comparison()?).into()))
        } else if let Some(Token(_, TokenType::Minus)) = self.tokens.peek() {
            self.tokens.next();
            Ok(Rc::new(Subtract::new(expr, self.comparison()?).into()))
        } else {
            Ok(expr)
        }
    }

    fn factor(&mut self) -> Result<Rc<NodeEnum>, ParserError> {
        let expr = self.unary()?;
        if let Some(Token(_, TokenType::Star)) = self.tokens.peek() {
            self.tokens.next();
            Ok(Rc::new(Multiply::new(expr, self.comparison()?).into()))
        } else if let Some(Token(_, TokenType::Slash)) = self.tokens.peek() {
            self.tokens.next();
            Ok(Rc::new(Divide::new(expr, self.comparison()?).into()))
        } else {
            Ok(expr)
        }
    }

    fn unary(&mut self) -> Result<Rc<NodeEnum>, ParserError> {
        self.exponent()
    }

    fn exponent(&mut self) -> Result<Rc<NodeEnum>, ParserError> {
        self.call()
    }

    fn call(&mut self) -> Result<Rc<NodeEnum>, ParserError> {
        let mut expr = self.function();

        loop {
            let (i, t) = if let Some(Token(i, x)) = self.tokens.peek() {
                (*i, x.clone())
            } else {
                return expr;
            };

            if t == TokenType::LParen {
                let potential_parameters = self.primary()?;
                let parameters = if let NodeEnum::Set(set) = potential_parameters.as_ref() {
                    potential_parameters
                } else {
                    // return Err(ParserError::UnexpectedNode(
                    //     i,
                    //     format!("Expected a Set here, but got a {potential_parameters:?}"),
                    // ));
                    Set::new(vec![potential_parameters])
                };

             expr = Ok(Call::new(expr?, parameters));
            } else {
                break;
            }
        }

        expr
    }

    fn function(&mut self) -> Result<Rc<NodeEnum>, ParserError> {
        let error_loc = if let Some(Token(i, _)) = self.tokens.peek() {
            *i
        } else {
            0
        };

        let expr = self.primary()?;

        if self.matchType(TokenType::LeftArrow) {
            let right = self.equality()?;
            match expr.clone().as_ref() {
                NodeEnum::Symbol(symbol) => {
                    return Ok(Function::new(
                        FunctionType::UserFunction(right),
                        vec![symbol.clone()],
                    ));
                }
                NodeEnum::Set(set) => {
                    let mut symbols = vec![];
                    for (i, value) in set.get_values().into_iter().enumerate() {
                        match value.as_ref() {
                            _ => {
                                return Err(ParserError::UnexpectedNode(
                                    error_loc,
                                    format!(
                                        "Expected set of Identifiers, but argument #{i} is a {value:?}"
                                    ),
                                ));
                            }

                            NodeEnum::Symbol(symbol) => symbols.push(symbol.clone()),
                        }
                    }

                    return Ok(Function::new(FunctionType::UserFunction(right), symbols));
                }
                _ => {
                    return Err(ParserError::UnexpectedNode(
                        error_loc,
                        format!("Expected Set, got {:?}", expr),
                    ));
                }
            }
        }

        Ok(expr)
    }

    fn primary(&mut self) -> Result<Rc<NodeEnum>, ParserError> {
        let (i, token) = if let Some(Token(i, token)) = self.tokens.next() {
            (i, token)
        } else {
            return Err(ParserError::UnexpectedEndOfTokens(
                "Expected a Primary here".into(),
            ));
        };

        match token {
            TokenType::Number(value) => Ok(Rc::new(Constant::new(value).into())),
            TokenType::Identifier(string) => Ok(Rc::new(
                Symbol::new_from_str(string, self.environment).into(),
            )),
            TokenType::LParen => {
                let expr = self.expression()?;
                let (i, t) = if let Some(Token(i, x)) = self.tokens.peek() {
                    (i, x)
                } else {
                    return Err(ParserError::UnexpectedEndOfTokens(
                        "Unclosed right parenthesis".into(),
                    ));
                };

                match t {
                    TokenType::RParen => {
                        self.tokens.next();
                        Ok(expr)
                    }
                    TokenType::Comma => {
                        let mut values = vec![expr];
                        while {
                            if let Some(Token(_, TokenType::RParen)) = self.tokens.peek() {
                                self.tokens.next();
                                false
                            } else {
                                true
                            }
                        } {
                            let (i, token) = if let Some(Token(i, x)) = self.tokens.peek() {
                                (i, x)
                            } else {
                                return Err(ParserError::UnexpectedEndOfTokens(
                                    "Expected comma here".into(),
                                ));
                            };

                            if *token == TokenType::Comma {
                                self.tokens.next();
                            } else {
                                return Err(ParserError::UnexpectedToken(
                                    *i,
                                    token.len(),
                                    format!("Expected comma here, but got {token:?}"),
                                ));
                            }

                            values.push(self.equality()?);
                        }

                        Ok(Set::new(values))
                    }
                    _ => Err(ParserError::Unimplemented(
                        *i,
                        t.len(),
                        format!("Expected either a comma or a right parenthesis here. Got {t:?}"),
                    )),
                }

                // if t != TokenType::RParen {
                //     Err(ParserError::UnexpectedToken(i, t.len(), format!("")))
                // } else {
                //     Ok(expr)
                // }
            }
            _ => Err(ParserError::UnexpectedToken(
                i,
                token.len(),
                format!("Unexpected token {token:?}"),
            )),
        }
    }
}
