use std::ops::Range;

use chumsky::{
    error::Rich,
    input::{Stream, ValueInput},
    prelude::*,
    primitive::just,
    recursive::recursive,
    span::SimpleSpan,
    IterParser,
};
use indexmap::IndexMap;
use logos::{Logos, Source};

use crate::tokens::Token;

pub mod ast;
#[cfg(test)]
mod tests;
use self::ast::{Expr, Expression, File};

pub type Span = SimpleSpan;
pub type Spanned<T> = (T, Span);

pub fn parse<'src>(src: &'src str) -> ParseResult<File<'_>, Rich<'_, Token<'_>>> {
    let toks: Vec<_> = Token::lexer(src)
        .spanned()
        .into_iter()
        .map(|(t, s)| (t.expect("TODO: add lexer error(s)"), Span::from(s)))
        .collect();
    let tok_stream = Stream::from_iter(toks).spanned((src.len()..src.len()).into());
    parser().parse(tok_stream)
}
pub(crate) fn parser<
    'tokens,
    'src: 'tokens,
    I: ValueInput<'tokens, Token = Token<'src>, Span = Span>,
>() -> impl Parser<'tokens, I, File<'src>, extra::Err<Rich<'tokens, Token<'src>, Span>>> {
    let word = select! { Token::Word(word) => word };

    let expr = recursive(|expr| {
        let var = select! {
            Token::VarIdent(name) => (Expr::Var as fn(_) -> _, name),
            Token::InputIdent(name) => (Expr::InputVar as fn(_) -> _, name)
        }
        .map_with(|(item_type, name), extra| Expression::new(item_type(name), extra.span()))
        .labelled("variable");

        let attrset = word
            .map_with(|n, e| (n, e.span()))
            .labelled("attr name")
            .then_ignore(just(Token::Colon))
            .then(expr)
            .labelled("attr body")
            .separated_by(just(Token::Comma))
            .collect::<Vec<_>>()
            .map(IndexMap::from_iter)
            .delimited_by(just(Token::BracketOpen), just(Token::BracketClose))
            .map_with(|v, e| (v, e.span()))
            .labelled("attrset");

        let node = word
            .map_with(|v, e| (v, e.span()))
            .then(attrset.clone().or_not())
            .map_with(|(name, params), extra| {
                Expression::new(Expr::Node(name, params), extra.span())
            })
            .or(var)
            .or(attrset
                .map_with(|attrset, extra| Expression::new(Expr::AttrSet(attrset), extra.span())))
            .labelled("node");

        #[allow(clippy::let_and_return)]
        let pipeline = node
            .clone()
            .then(choice((
                just(Token::Pipe).to(Expr::SimplePipe as fn(_, _) -> _),
                just(Token::MappingPipe).to(Expr::MappingPipe as fn(_, _) -> _),
                just(Token::NullPipe).to(Expr::NullPipe as fn(_, _) -> _),
            )))
            .repeated()
            .foldr_with(node, |(curr, pipe), next, extra| {
                Expression::new(pipe(Box::new(curr), Box::new(next)), extra.span())
            });

        pipeline
    });

    let decl = just(Token::Def).ignore_then(
        word.map_with(|n, e| (n, e.span()))
            .then_ignore(just(Token::Equals))
            .then(expr.clone().map(|expr| expr))
            .then_ignore(just(Token::SemiColon)),
    );

    expr.map(|expr| File {
        decls: IndexMap::from_iter([(("main", (0..0).into()), expr)]),
    })
    .or(decl.repeated().collect::<Vec<_>>().map(|decls| File {
        decls: IndexMap::from_iter(decls),
    }))
}