use std::ops::Range;

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

use crate::tokens::Token;

pub mod ast;
#[cfg(test)]
mod tests;
use self::ast::{Expr, 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());
    expr_parser().parse(tok_stream)
}

fn expr_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| 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(|(name, params)| Expr::Node(name, params))
            .or(var)
            .or(attrset.map(Expr::AttrSet))
            .labelled("node");

        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(node, |(curr, pipe), next| {
                pipe(Box::new(curr), Box::new(next))
            });

        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_with(|expr, extra| (expr, extra.span())))
            .then_ignore(just(Token::SemiColon)),
    );

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