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

use crate::tokens::Token;

pub mod ast;
#[cfg(test)]
mod tests;
use self::ast::{
    raw_ast::{RawExpr, RawExpression},
    File,
};

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

pub fn parse(src: &str) -> ParseResult<File<'_>, Rich<'_, Token<'_>>> {
    let toks: Vec<_> = Token::lexer(src)
        .spanned()
        .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) = e => (word, e.span())};

    let expr = recursive(|expr| {
        let lit = select! {
            Token::Int(i) = e => RawExpression::new(RawExpr::Lit(ast::Lit::Int(i.parse().expect("TODO: handle better"))), e.span()),
            Token::Float(f) = e => RawExpression::new(RawExpr::Lit(ast::Lit::Float(f.parse().expect("TODO: handle better"))), e.span()),
            Token::String(s) = e => RawExpression::new(RawExpr::Lit(ast::Lit::String(s.strip_prefix('"').expect("a").strip_suffix('"').expect("b"))), e.span())
        };
        let mat = just(Token::Mat)
            .ignore_then(select! { Token::Dimensions(dimensions) = e => (dimensions, e.span())})
            .then(
                lit.separated_by(just(Token::Comma))
                    .collect::<Vec<_>>()
                    .separated_by(just(Token::Semicolon))
                    .collect::<Vec<_>>()
                    .delimited_by(just(Token::BracketOpen), just(Token::BracketClose)),
            )
            .map_with(|(dimensions, data), e| {
                // TODO: Validation and proper error handling/reporting
                // (validation = validating the matrix dimensions)
                RawExpression::new(
                    RawExpr::Matrix(dimensions, data.into_iter().flatten().collect()),
                    e.span(),
                )
            });
        let var = select! {
            Token::VarIdent(name) => (RawExpr::Var as fn(_) -> _, name),
            Token::InputIdent(name) => (RawExpr::InputVar as fn(_) -> _, name)
        }
        .map_with(|(item_type, name), extra| RawExpression::new(item_type(name), extra.span()))
        .labelled("variable");

        let attrset = word
            .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::BraceOpen), just(Token::BraceClose))
            .map_with(|v, e| (v, e.span()))
            .labelled("attrset");

        let node = word
            .repeated()
            .collect()
            .then(attrset.clone().or_not())
            .map_with(|(name, params), extra| {
                RawExpression::new(RawExpr::Node(name, params), extra.span())
            })
            // .or(var)
            // .or(attrset
            //     .map_with(|attrset, extra| Expression::new(Expr::AttrSet(attrset), extra.span())))
            // .or(lit)
            // .or(mat)
            .labelled("node");

        let atom = var
            .or(lit)
            .or(mat)
            .or(attrset.map_with(|attrset, extra| {
                RawExpression::new(RawExpr::AttrSet(attrset), extra.span())
            }))
            .or(node.clone());

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

        pipeline
    });

    let decls = just(Token::Def)
        .ignore_then(
            word.then_ignore(just(Token::Equals))
                .then(expr.clone().map(|expr| expr))
                .then_ignore(just(Token::Semicolon)),
        )
        .repeated()
        .collect::<Vec<_>>()
        .map(|decls| File {
            decls: IndexMap::from_iter(decls),
        });

    let single_expr = expr.map(|expr| File {
        decls: IndexMap::from_iter([(("main", (0..0).into()), expr)]),
    });

    just(Token::Def).rewind().ignore_then(decls).or(single_expr)
    // single_expr.or(decls)

    // 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),
    // }))
}

pub mod asg {
    use petgraph::graph::DiGraph;

    use super::Spanned;
}