use std::{iter::Peekable, rc::Rc};

use crate::{
    ast::{Ast, Atom, Error, Expr, models::LetVar},
    lex::Token,
    span::{Pos, Span, Spanned},
};

pub(super) const MAX_DEPTH: usize = 256; // TODO: make it a compile flag

fn parse_number(number: &str) -> Result<Atom, Error> {
    match number.parse() {
        Ok(ok) => Ok(Atom::Int(ok)),
        Err(err) => Err(Error::InvalidIntegerLiteral(number.into(), err)),
    }
}

fn parse_string(string: &str) -> Result<Atom, Error> {
    let mut result = String::new();
    let mut is_escape = false;

    for ch in string.chars() {
        if !is_escape {
            match ch {
                '\\' => is_escape = true,
                _ => result.push(ch),
            }
        } else {
            match ch {
                '"' => result.push('"'),
                'n' => result.push('\n'),
                '\\' => result.push('\\'),
                '\n' => {}
                _ => return Err(Error::UnexpectedEscapeChar(ch)),
            }
            is_escape = false;
        }
    }

    if is_escape {
        return Err(Error::UnclosedString(string.into()));
    }

    Ok(Atom::Str(result.into()))
}

fn parse_symbol(symbol: &str) -> Atom {
    match symbol {
        "true" => Atom::Bool(true),
        "false" => Atom::Bool(false),
        "nil" => Atom::Nil,
        _ => Atom::Sym(symbol.into()),
    }
}

pub struct Parser<'a, I>
where
    I: Iterator<Item = Spanned<Token<'a>>>,
{
    tokens: Peekable<I>,
    last_token_span: Span,
    depth: usize,
}

impl<'a, I> Parser<'a, I>
where
    I: Iterator<Item = Spanned<Token<'a>>>,
{
    pub fn new(tokens: I) -> Self {
        Self {
            tokens: tokens.peekable(),
            last_token_span: Span::new(Pos::new(1, 0, 0), Pos::new(1, 0, 0)),
            depth: 0,
        }
    }

    fn peek(&mut self) -> Option<Spanned<Token<'a>>> {
        self.tokens.peek().copied()
    }

    fn peek_token(&mut self) -> Result<Spanned<Token<'a>>, Spanned<Error>> {
        match self.peek() {
            Some(token) => Ok(token),
            None => Err(Spanned::new(Error::UnexpectedEof, self.last_token_span)),
        }
    }

    fn consume(&mut self) -> Result<(), Spanned<Error>> {
        match self.tokens.next() {
            Some(token) => {
                let (token, span) = token.into_parts();

                if let Token::LeftPar = token {
                    self.depth = self.depth.saturating_add(1);
                    if self.depth > MAX_DEPTH {
                        return Err(Spanned::new(Error::RecursionLimit, span));
                    }
                } else if let Token::RightPar = token {
                    self.depth = self.depth.saturating_sub(1);
                }

                self.last_token_span = span;
                Ok(())
            }
            None => panic!("no tokens to consume"),
        }
    }

    fn require_sym(&mut self, symbol: &str) -> Result<Span, Spanned<Error>> {
        let (token, span) = self.peek_token()?.into_parts();
        self.consume()?;
        match token {
            Token::Symbol(s) if s == symbol => Ok(span),
            _ => Err(Spanned::new(Error::UnexpectedToken, span)),
        }
    }

    fn require_left_par(&mut self) -> Result<Span, Spanned<Error>> {
        let (token, span) = self.peek_token()?.into_parts();
        self.consume()?;
        match token {
            Token::LeftPar => Ok(span),
            _ => Err(Spanned::new(Error::UnexpectedToken, span)),
        }
    }

    fn require_right_par(&mut self) -> Result<Span, Spanned<Error>> {
        let (token, span) = self.peek_token()?.into_parts();
        self.consume()?;
        match token {
            Token::RightPar => Ok(span),
            _ => Err(Spanned::new(Error::UnexpectedToken, span)),
        }
    }

    fn parse_sym(&mut self) -> Result<Spanned<&'a str>, Spanned<Error>> {
        let (token, span) = self.peek_token()?.into_parts();
        self.consume()?;
        match token {
            Token::Symbol(symbol) => Ok(Spanned::new(symbol, span)),
            _ => Err(Spanned::new(Error::UnexpectedToken, span)),
        }
    }

    fn parse_args(&mut self, open_span: Span) -> Result<Vec<Spanned<Rc<str>>>, Spanned<Error>> {
        let mut args = Vec::new();
        while let Some(token) = self.peek() {
            let (token, span) = token.into_parts();

            self.consume()?;
            match token {
                Token::Symbol(symbol) => args.push(Spanned::new(symbol.into(), span)),
                Token::RightPar => return Ok(args),
                _ => return Err(Spanned::new(Error::UnexpectedToken, span)),
            }
        }

        Err(Spanned::new(Error::UnclosedPar, open_span))
    }

    fn parse_body(
        &mut self,
        open_span: Span,
        non_empty: bool,
    ) -> Result<Vec<Spanned<Expr>>, Spanned<Error>> {
        let mut body = Vec::new();
        if non_empty {
            body.push(self.parse_expr()?);
        }
        while let Some(token) = self.peek() {
            match token.inner {
                Token::RightPar => {
                    self.consume()?;
                    return Ok(body);
                }
                _ => body.push(self.parse_expr()?),
            }
        }

        Err(Spanned::new(Error::UnclosedPar, open_span))
    }

    fn parse_var(&mut self) -> Result<Spanned<LetVar>, Spanned<Error>> {
        let open_span = self.require_left_par()?;
        let name = self.parse_sym()?.map(Into::into);
        let expr = self.parse_expr()?;
        let close_span = self.require_right_par()?;

        let let_var = LetVar { name, expr };
        let span = Span::new(open_span.start, close_span.end);
        Ok(Spanned::new(let_var, span))
    }

    fn parse_vars(&mut self, open_span: Span) -> Result<Vec<Spanned<LetVar>>, Spanned<Error>> {
        let mut vars = Vec::new();
        vars.push(self.parse_var()?);

        while let Some(token) = self.peek() {
            let (token, span) = token.into_parts();
            match token {
                Token::LeftPar => vars.push(self.parse_var()?),
                Token::RightPar => {
                    self.consume()?;
                    return Ok(vars);
                }
                _ => return Err(Spanned::new(Error::UnexpectedToken, span)),
            }
        }

        Err(Spanned::new(Error::UnclosedPar, open_span))
    }

    fn parse_fn(&mut self, open_span: Span) -> Result<Spanned<Expr>, Spanned<Error>> {
        self.consume()?;

        let name = self.parse_sym()?.map(Into::into);
        let args_open_span = self.require_left_par()?;
        let args = self.parse_args(args_open_span)?;

        let body = self.parse_body(open_span, true)?;
        let function = Expr::Fn { name, args, body };
        let span = Span::new(open_span.start, self.last_token_span.end);

        Ok(Spanned::new(function, span))
    }

    fn parse_const(&mut self, open_span: Span) -> Result<Spanned<Expr>, Spanned<Error>> {
        self.consume()?;

        let vars = self.parse_vars(open_span)?;
        let constant = Expr::Const { vars };
        let span = Span::new(open_span.start, self.last_token_span.end);
        Ok(Spanned::new(constant, span))
    }

    fn parse_let(&mut self, open_span: Span) -> Result<Spanned<Expr>, Spanned<Error>> {
        self.consume()?;

        let vars_span = self.require_left_par()?;
        let vars = self.parse_vars(vars_span)?;
        let body = self.parse_body(open_span, true)?;
        let let_vars = Expr::Let { vars, body };
        let span = Span::new(open_span.start, self.last_token_span.end);
        Ok(Spanned::new(let_vars, span))
    }

    fn parse_if(&mut self, open_span: Span) -> Result<Spanned<Expr>, Spanned<Error>> {
        self.consume()?;

        let condition = self.parse_expr()?.map(Into::into);
        let then_expr = self.parse_expr()?.map(Into::into);
        let else_expr = match self.peek_token()?.inner {
            Token::RightPar => {
                self.consume()?;
                None
            }
            _ => {
                let expr = self.parse_expr()?.map(Into::into);
                self.require_right_par()?;
                Some(expr)
            }
        };

        let if_stmt = Expr::If {
            condition,
            then_expr,
            else_expr,
        };
        let span = Span::new(open_span.start, self.last_token_span.end);
        Ok(Spanned::new(if_stmt, span))
    }

    fn parse_for(&mut self, open_span: Span) -> Result<Spanned<Expr>, Spanned<Error>> {
        self.consume()?;

        let loop_var = self.parse_sym()?.map(Into::into);
        self.require_sym("from")?;
        let from = self.parse_expr()?.map(Into::into);
        self.require_sym("to")?;
        let to = self.parse_expr()?.map(Into::into);
        let body = self.parse_body(open_span, true)?;

        let for_loop = Expr::For {
            loop_var,
            from,
            to,
            body,
        };
        let span = Span::new(open_span.start, self.last_token_span.end);
        Ok(Spanned::new(for_loop, span))
    }

    fn parse_set(&mut self, open_span: Span) -> Result<Spanned<Expr>, Spanned<Error>> {
        self.consume()?;

        let target_var = self.parse_sym()?.map(Into::into);
        let expr = self.parse_expr()?.map(Into::into);
        self.require_right_par()?;

        let set = Expr::Set { target_var, expr };
        let span = Span::new(open_span.start, self.last_token_span.end);
        Ok(Spanned::new(set, span))
    }

    fn parse_do(&mut self, open_span: Span) -> Result<Spanned<Expr>, Spanned<Error>> {
        self.consume()?;

        let body = self.parse_body(open_span, true)?;

        let do_body = Expr::Do { body };
        let span = Span::new(open_span.start, self.last_token_span.end);
        Ok(Spanned::new(do_body, span))
    }

    fn parse_call(&mut self, open_span: Span) -> Result<Spanned<Expr>, Spanned<Error>> {
        let fn_name = self.parse_sym()?.map(Into::into);
        let args = self.parse_body(open_span, false)?;

        let call = Expr::Call { fn_name, args };
        let span = Span::new(open_span.start, self.last_token_span.end);
        Ok(Spanned::new(call, span))
    }

    fn parse_special_form(&mut self, open_span: Span) -> Result<Spanned<Expr>, Spanned<Error>> {
        let (token, span) = self.peek_token()?.into_parts();

        let symbol = match token {
            Token::Symbol(symbol) => symbol,
            Token::RightPar => {
                self.consume()?;
                let span = Span::new(open_span.start, span.end);
                return Ok(Spanned::new(Expr::Atom(Atom::Nil), span));
            }
            _ => return Err(Spanned::new(Error::UnexpectedToken, span)),
        };

        match symbol {
            "fn" => self.parse_fn(open_span),
            "const" => self.parse_const(open_span),
            "let" => self.parse_let(open_span),
            "if" => self.parse_if(open_span),
            "for" => self.parse_for(open_span),
            "set" => self.parse_set(open_span),
            "do" => self.parse_do(open_span),
            _ => self.parse_call(open_span),
        }
    }

    fn parse_expr(&mut self) -> Result<Spanned<Expr>, Spanned<Error>> {
        let (token, span) = self.peek_token()?.into_parts();

        let expr = match token {
            Token::LeftPar => {
                self.consume()?;
                self.parse_special_form(span)?
            }
            Token::RightPar => {
                self.consume()?;
                return Err(Spanned::new(Error::UnexpectedClosePar, span));
            }
            Token::Number(number) => {
                self.consume()?;
                let atom = parse_number(number).map_err(|e| Spanned::new(e, span))?;
                Spanned::new(Expr::Atom(atom), span)
            }
            Token::String(string) => {
                self.consume()?;
                let atom = parse_string(string).map_err(|e| Spanned::new(e, span))?;
                Spanned::new(Expr::Atom(atom), span)
            }
            Token::UnclosedString(string) => {
                self.consume()?;
                return Err(Spanned::new(Error::UnclosedString(string.into()), span));
            }
            Token::Symbol(symbol) => {
                self.consume()?;
                let atom = parse_symbol(symbol);
                Spanned::new(Expr::Atom(atom), span)
            }
        };

        Ok(expr)
    }

    pub fn parse(mut self) -> Result<Ast, Spanned<Error>> {
        let mut ast = Vec::new();

        while self.peek().is_some() {
            ast.push(self.parse_expr()?)
        }

        Ok(Ast::new(ast))
    }
}