// TODO: move AST to its own module
// TODO: remove StringMatcher cruft
// TODO: good error messages?

use crate::ast::{Ast, StringPart};
use crate::lexer::*;
use crate::spans::*;
use chumsky::{input::ValueInput, prelude::*, recursive::Recursive};
use std::fmt;

pub struct StringMatcher();

impl PartialEq for StringMatcher {
    fn eq(&self, _other: &StringMatcher) -> bool {
        true
    }
}

impl Clone for StringMatcher {
    fn clone(&self) -> StringMatcher {
        unreachable!()
    }
}

impl fmt::Display for StringMatcher {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "string matcher")
    }
}

impl fmt::Debug for StringMatcher {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "string matcher")
    }
}

fn is_word_char(c: char) -> bool {
    if c.is_ascii_alphanumeric() {
        return true;
    };
    matches!(c, '_' | '/' | '?' | '!')
}

fn parse_string(s: &'static str, span: SimpleSpan) -> Result<Vec<Spanned<StringPart>>, String> {
    // println!("parsing string pattern: {s}");
    let mut parts = vec![];
    let mut current_part = String::new();
    let mut start = span.start;
    let mut is_word = false;

    let mut chars = s.char_indices();
    while let Some((i, char)) = chars.next() {
        match char {
            '{' => {
                if is_word {
                    return Err("interpolations must only contain words".to_string());
                };
                match chars.next() {
                    None => return Err("unclosed brace".to_string()),
                    Some((_, '{')) => current_part.push('{'),
                    Some((i, c)) => {
                        if !current_part.is_empty() {
                            parts.push((
                                StringPart::Data(current_part),
                                SimpleSpan::new(span.context(), start..start + i),
                            ));
                        };
                        current_part = String::new();
                        start = i;
                        is_word = true;
                        if c.is_ascii_lowercase() {
                            current_part.push(c);
                        } else {
                            return Err("interpolations must only contain words".to_string());
                        }
                    }
                }
            }
            '}' => {
                if is_word {
                    parts.push((
                        StringPart::Word(current_part.clone()),
                        SimpleSpan::new(span.context(), start..start + i),
                    ));
                    current_part = String::new();
                    start = i;
                    is_word = false;
                } else {
                    match chars.next() {
                        None => return Err("unclosed brace".to_string()),
                        Some((_, c)) => current_part.push(c),
                    }
                }
            }
            _ => {
                if is_word {
                    if is_word_char(char) {
                        current_part.push(char)
                    } else {
                        return Err("interpolations must only contain words".to_string());
                    }
                } else {
                    current_part.push(char)
                }
            }
        }
    }

    if current_part == s {
        parts.push((
            StringPart::Inline(current_part),
            SimpleSpan::new(span.context(), start..span.end),
        ))
    } else if !current_part.is_empty() {
        let part_len = current_part.len();
        parts.push((
            StringPart::Data(current_part),
            SimpleSpan::new(span.context(), start..part_len),
        ))
    }

    Ok(parts)
}

pub fn parser<I>(
) -> impl Parser<'static, I, Spanned<Ast>, extra::Err<Rich<'static, Token, Span>>> + Clone
where
    I: ValueInput<'static, Token = Token, Span = Span>,
{
    use Ast::*;

    let mut expr = Recursive::declare();

    let mut pattern = Recursive::declare();

    let mut simple = Recursive::declare();

    let mut nonbinding = Recursive::declare();

    let separators = recursive(|separators| {
        just(Token::Punctuation(","))
            .or(just(Token::Punctuation("\n")))
            .then(separators.clone().repeated())
    });

    let terminators = recursive(|terminators| {
        just(Token::Punctuation(";"))
            .or(just(Token::Punctuation("\n")))
            .then(terminators.clone().repeated())
    });

    let placeholder_pattern =
        select! {Token::Punctuation("_") => PlaceholderPattern}.map_with(|p, e| (p, e.span()));

    let word_pattern = select! { Token::Word(w) => WordPattern(w) }.map_with(|w, e| (w, e.span()));

    let atom_pattern = select! {
        Token::Nil => NilPattern,
        Token::Boolean(b) => BooleanPattern(b),
        Token::Number(n) => NumberPattern(n),
        Token::Keyword(k) => KeywordPattern(k),
    }
    .map_with(|a, e| (a, e.span()));

    let string_pattern = select! {Token::String(s) => s}.try_map_with(|s, e| {
        let parsed = parse_string(s, e.span());
        match parsed {
            Ok(parts) => match parts[0] {
                (StringPart::Inline(_), _) => Ok((StringPattern(s), e.span())),
                _ => Ok((InterpolatedPattern(parts.clone()), e.span())),
            },
            Err(msg) => Err(Rich::custom(e.span(), msg)),
        }
    });

    let bare_splat = just(Token::Punctuation("...")).map_with(|_, e| {
        (
            Splattern(Box::new((PlaceholderPattern, e.span()))),
            e.span(),
        )
    });

    let splattable = word_pattern.or(placeholder_pattern);

    let patt_splat = just(Token::Punctuation("..."))
        .ignore_then(splattable)
        .map_with(|x, e| (Splattern(Box::new(x)), e.span()));

    let splattern = patt_splat.or(bare_splat);

    let tuple_pattern = pattern
        .clone()
        .or(splattern.clone())
        .separated_by(separators.clone())
        .allow_leading()
        .allow_trailing()
        .collect()
        .delimited_by(just(Token::Punctuation("(")), just(Token::Punctuation(")")))
        .map_with(|tuple, e| (TuplePattern(tuple), e.span()))
        .labelled("tuple pattern");

    let list_pattern = pattern
        .clone()
        .or(splattern.clone())
        .separated_by(separators.clone())
        .allow_leading()
        .allow_trailing()
        .collect()
        .delimited_by(just(Token::Punctuation("[")), just(Token::Punctuation("]")))
        .map_with(|list, e| (ListPattern(list), e.span()));

    let pair_pattern = select! {Token::Keyword(k) => k}
        .then(pattern.clone())
        .map_with(|(key, patt), e| (PairPattern(key, Box::new(patt)), e.span()));

    let shorthand_pattern = select! {Token::Word(w) => w}.map_with(|w, e| {
        (
            PairPattern(w, Box::new((WordPattern(w), e.span()))),
            e.span(),
        )
    });

    let dict_pattern = pair_pattern
        .or(shorthand_pattern)
        .or(splattern.clone())
        .separated_by(separators.clone())
        .allow_leading()
        .allow_trailing()
        .collect()
        .delimited_by(
            just(Token::Punctuation("#{")),
            just(Token::Punctuation("}")),
        )
        .map_with(|dict, e| (DictPattern(dict), e.span()));

    let keyword = select! {Token::Keyword(k) => Keyword(k),}.map_with(|k, e| (k, e.span()));

    let as_pattern = select! {Token::Word(w) => w}
        .then_ignore(just(Token::Reserved("as")))
        .then(select! {Token::Keyword(k) => k})
        .map_with(|(w, t), e| (AsPattern(w, t), e.span()));

    pattern.define(
        atom_pattern
            .or(string_pattern)
            .or(as_pattern)
            .or(word_pattern)
            .or(placeholder_pattern)
            .or(tuple_pattern.clone())
            .or(list_pattern)
            .or(dict_pattern)
            .labelled("pattern"),
    );

    let placeholder =
        select! {Token::Punctuation("_") => Placeholder}.map_with(|p, e| (p, e.span()));

    let word = select! { Token::Word(w) => Word(w) }
        .map_with(|w, e| (w, e.span()))
        .labelled("word");

    let value = select! {
        Token::Nil => Nil,
        Token::Boolean(b) => Boolean(b),
        Token::Number(n) => Number(n),
    }
    .map_with(|v, e| (v, e.span()));

    let string = select! {Token::String(s) => s}.try_map_with(|s, e| {
        let parsed = parse_string(s, e.span());
        match parsed {
            Ok(parts) => match parts[0] {
                (StringPart::Inline(_), _) => Ok((String(s), e.span())),
                _ => Ok((Interpolated(parts), e.span())),
            },
            Err(msg) => Err(Rich::custom(e.span(), msg)),
        }
    });

    let tuple = simple
        .clone()
        .separated_by(separators.clone())
        .allow_leading()
        .allow_trailing()
        .collect()
        .delimited_by(just(Token::Punctuation("(")), just(Token::Punctuation(")")))
        .map_with(|tuple, e| (Tuple(tuple), e.span()));

    let args = simple
        .clone()
        .or(placeholder)
        .separated_by(separators.clone())
        .allow_leading()
        .allow_trailing()
        .collect()
        .delimited_by(just(Token::Punctuation("(")), just(Token::Punctuation(")")))
        .map_with(|args, e| (Arguments(args), e.span()));

    let or = just(Token::Reserved("or")).map_with(|_, e| (Or, e.span()));

    let and = just(Token::Reserved("and")).map_with(|_, e| (And, e.span()));

    let synth_root = or.or(and).or(word).or(keyword);

    let synth_term = keyword.or(args);

    let synthetic = synth_root
        .then(synth_term.clone())
        .then(synth_term.clone().repeated().collect())
        .map_with(|((root, first), rest), e| {
            (Synthetic(Box::new(root), Box::new(first), rest), e.span())
        });

    let splat = just(Token::Punctuation("..."))
        .ignore_then(word)
        .map_with(|(w, _), e| {
            (
                Splat(if let Word(w) = w { w } else { unreachable!() }),
                e.span(),
            )
        });

    let list = simple
        .clone()
        .or(splat.clone())
        .separated_by(separators.clone())
        .allow_leading()
        .allow_trailing()
        .collect()
        .delimited_by(just(Token::Punctuation("[")), just(Token::Punctuation("]")))
        .map_with(|list, e| (List(list), e.span()));

    let pair = select! {Token::Keyword(k) => k}
        .then(simple.clone())
        .map_with(|(key, value), e| (Pair(key, Box::new(value)), e.span()));

    let shorthand = select! {Token::Word(w) => w}
        .map_with(|w, e| (Pair(w, Box::new((Word(w), e.span()))), e.span()));

    let dict = pair
        .or(shorthand)
        .or(splat.clone())
        .separated_by(separators.clone())
        .allow_leading()
        .allow_trailing()
        .collect()
        .delimited_by(
            just(Token::Punctuation("#{")),
            just(Token::Punctuation("}")),
        )
        .map_with(|dict, e| (Dict(dict), e.span()));

    let recur = just(Token::Reserved("recur"))
        .ignore_then(tuple.clone())
        .map_with(|args, e| {
            let (Tuple(args), _) = args else {
                unreachable!()
            };
            (Recur(args), e.span())
        });

    let block = expr
        .clone()
        .separated_by(terminators.clone())
        .allow_leading()
        .allow_trailing()
        .collect()
        .delimited_by(just(Token::Punctuation("{")), just(Token::Punctuation("}")))
        .map_with(|block, e| (Block(block), e.span()))
        .recover_with(via_parser(nested_delimiters(
            Token::Punctuation("{"),
            Token::Punctuation("}"),
            [
                (Token::Punctuation("("), Token::Punctuation(")")),
                (Token::Punctuation("["), Token::Punctuation("]")),
            ],
            |span| (Error, span),
        )));

    let r#if = just(Token::Reserved("if"))
        .ignore_then(simple.clone())
        .then_ignore(terminators.clone().or_not())
        .then_ignore(just(Token::Reserved("then")))
        .then(expr.clone())
        .then_ignore(terminators.clone().or_not())
        .then_ignore(just(Token::Reserved("else")))
        .then(expr.clone())
        .map_with(|((condition, then_branch), else_branch), e| {
            (
                If(
                    Box::new(condition),
                    Box::new(then_branch),
                    Box::new(else_branch),
                ),
                e.span(),
            )
        });

    let when_clause = simple
        .clone()
        .then_ignore(just(Token::Punctuation("->")))
        .then(expr.clone())
        .map_with(|(cond, body), e| (WhenClause(Box::new(cond), Box::new(body)), e.span()));

    let when = just(Token::Reserved("when"))
        .ignore_then(
            when_clause
                .separated_by(terminators.clone())
                .allow_trailing()
                .allow_leading()
                .collect()
                .delimited_by(just(Token::Punctuation("{")), just(Token::Punctuation("}"))),
        )
        .map_with(|clauses, e| (When(clauses), e.span()));

    let guarded_clause = pattern
        .clone()
        .then_ignore(just(Token::Reserved("if")))
        .then(simple.clone())
        .then_ignore(just(Token::Punctuation("->")))
        .then(expr.clone())
        .map_with(|((patt, guard), body), e| {
            (
                MatchClause(Box::new(patt), Box::new(Some(guard)), Box::new(body)),
                e.span(),
            )
        });

    let match_clause = pattern
        .clone()
        .then_ignore(just(Token::Punctuation("->")))
        .then(expr.clone())
        .map_with(|(patt, body), e| {
            (
                MatchClause(Box::new(patt), Box::new(None), Box::new(body)),
                e.span(),
            )
        });

    let r#match = just(Token::Reserved("match"))
        .ignore_then(simple.clone())
        .then_ignore(just(Token::Reserved("with")))
        .then(
            match_clause
                .clone()
                .or(guarded_clause.clone())
                .separated_by(terminators.clone())
                .allow_leading()
                .allow_trailing()
                .collect()
                .delimited_by(just(Token::Punctuation("{")), just(Token::Punctuation("}"))),
        )
        .map_with(|(expr, clauses), e| (Match(Box::new(expr), clauses), e.span()));

    let receive = just(Token::Reserved("receive"))
        .ignore_then(
            match_clause
                .clone()
                .or(guarded_clause)
                .separated_by(terminators.clone())
                .allow_leading()
                .allow_trailing()
                .collect()
                .delimited_by(just(Token::Punctuation("{")), just(Token::Punctuation("}"))),
        )
        .map_with(|clauses, e| (Receive(clauses), e.span()));

    let conditional = when.or(r#if).or(r#match).or(receive);

    let panic = just(Token::Reserved("panic!"))
        .ignore_then(nonbinding.clone())
        .map_with(|expr, e| (Panic(Box::new(expr)), e.span()));

    let do_ = just(Token::Reserved("do"))
        .ignore_then(
            nonbinding
                .clone()
                .separated_by(
                    just(Token::Punctuation(">")).then(just(Token::Punctuation("\n")).repeated()),
                )
                .collect(),
        )
        .map_with(|exprs, e| (Do(exprs), e.span()));

    let repeat = just(Token::Reserved("repeat"))
        .ignore_then(simple.clone())
        .then(block.clone())
        .map_with(|(count, body), e| (Repeat(Box::new(count), Box::new(body)), e.span()));

    let fn_guarded = tuple_pattern
        .clone()
        .then_ignore(just(Token::Reserved("if")))
        .then(simple.clone())
        .then_ignore(just(Token::Punctuation("->")))
        .then(nonbinding.clone())
        .map_with(|((patt, guard), body), e| {
            (
                MatchClause(Box::new(patt), Box::new(Some(guard)), Box::new(body)),
                e.span(),
            )
        })
        .labelled("function clause");

    let fn_unguarded = tuple_pattern
        .clone()
        .then_ignore(just(Token::Punctuation("->")))
        .then(nonbinding.clone())
        .map_with(|(patt, body), e| {
            (
                MatchClause(Box::new(patt), Box::new(None), Box::new(body)),
                e.span(),
            )
        })
        .labelled("function clause");

    let fn_clause = fn_guarded.clone().or(fn_unguarded.clone());

    let lambda = just(Token::Reserved("fn"))
        .ignore_then(fn_unguarded.clone())
        .map_with(|clause, e| {
            (
                Fn("", Box::new((Ast::FnBody(vec![clause]), e.span())), None),
                e.span(),
            )
        });

    let fn_clauses = fn_clause
        .clone()
        .separated_by(terminators.clone())
        .allow_leading()
        .allow_trailing()
        .collect();

    let loop_multiclause = fn_clauses
        .clone()
        .delimited_by(just(Token::Punctuation("{")), just(Token::Punctuation("}")));

    let fn_single_clause = fn_clause.clone().map_with(|c, _| vec![c]);

    let r#loop = just(Token::Reserved("loop"))
        .ignore_then(tuple.clone())
        .then_ignore(just(Token::Reserved("with")))
        .then(loop_multiclause.clone().or(fn_single_clause.clone()))
        .map_with(|(init, body), e| (Loop(Box::new(init), body), e.span()));

    simple.define(
        synthetic
            .or(recur)
            .or(word)
            .or(keyword)
            .or(value)
            .or(tuple.clone())
            .or(list)
            .or(dict)
            .or(string)
            .or(lambda.clone())
            .labelled("simple expression"),
    );

    nonbinding.define(
        simple
            .clone()
            .or(conditional)
            .or(block)
            .or(panic)
            .or(do_)
            .or(repeat)
            .or(r#loop)
            .labelled("nonbinding expression"),
    );

    let let_ = just(Token::Reserved("let"))
        .ignore_then(pattern.clone())
        .then_ignore(just(Token::Punctuation("=")))
        .then(nonbinding.clone())
        .map_with(|(pattern, expression), e| {
            (Let(Box::new(pattern), Box::new(expression)), e.span())
        });

    let box_ = just(Token::Reserved("box"))
        .ignore_then(word)
        .then_ignore(just(Token::Punctuation("=")))
        .then(nonbinding.clone())
        .map_with(|(word, expr), e| {
            let name = if let Word(w) = word.0 {
                w
            } else {
                unreachable!()
            };
            (LBox(name, Box::new(expr)), e.span())
        });

    let fn_decl = just(Token::Reserved("fn"))
        .ignore_then(word)
        .map_with(|(word, _), e| {
            let name = if let Word(w) = word {
                w
            } else {
                unreachable!()
            };
            (FnDeclaration(name), e.span())
        });

    let fn_named = just(Token::Reserved("fn"))
        .ignore_then(word)
        .then(fn_unguarded.clone())
        .map_with(|(word, clause), e| {
            let name = if let Word(word) = word.0 {
                word
            } else {
                unreachable!()
            };
            (
                Fn(name, Box::new((Ast::FnBody(vec![clause]), e.span())), None),
                e.span(),
            )
        });

    let docstr = select! {Token::String(s) => s};

    let fn_multiclause = separators
        .clone()
        .or_not()
        .ignore_then(docstr.or_not())
        .then(fn_clauses.clone())
        .delimited_by(just(Token::Punctuation("{")), just(Token::Punctuation("}")))
        .map_with(|(docstr, clauses), e| (docstr, clauses, e.span()));

    let fn_compound = just(Token::Reserved("fn"))
        .ignore_then(word)
        .then(fn_multiclause)
        .map_with(|(word, (docstr, clauses, _)), e| {
            let name = if let Word(word) = word.0 {
                word
            } else {
                unreachable!()
            };
            (
                Fn(name, Box::new((Ast::FnBody(clauses), e.span())), docstr),
                e.span(),
            )
        });

    let fn_ = fn_named.or(fn_compound).or(fn_decl);

    let binding = let_.or(box_).or(fn_);

    expr.define(binding.or(nonbinding));

    let script = expr
        .separated_by(terminators.clone())
        .allow_trailing()
        .allow_leading()
        .collect()
        .map_with(|exprs, e| (Block(exprs), e.span()));

    script
}