use crate::lexer::*;
use crate::spans::*;
use chumsky::{input::ValueInput, prelude::*, recursive::Recursive};
use std::fmt;
use struct_scalpel::Dissectible;

#[derive(Clone, Debug, PartialEq)]
pub struct WhenClause<'src> {
    pub cond: Spanned<Ast<'src>>,
    pub body: Spanned<Ast<'src>>,
}

impl<'src> fmt::Display for WhenClause<'src> {
    fn fmt(self: &WhenClause<'src>, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "cond: {}, body: {}", self.cond.0, self.body.0)
    }
}

#[derive(Clone, Debug, PartialEq)]
pub struct MatchClause<'src> {
    pub patt: Spanned<Pattern<'src>>,
    pub guard: Option<Spanned<Ast<'src>>>,
    pub body: Spanned<Ast<'src>>,
}

impl<'src> fmt::Display for MatchClause<'src> {
    fn fmt(self: &MatchClause<'src>, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "pattern: {}, guard: {:?} body: {}",
            self.patt.0, self.guard, self.body.0
        )
    }
}

#[derive(Clone, Debug, PartialEq, Dissectible)]
pub enum Ast<'src> {
    Error,
    Placeholder,
    Nil,
    Boolean(bool),
    Number(f64),
    Keyword(&'src str),
    Word(&'src str),
    String(&'src str),
    Block(Vec<Spanned<Self>>),
    If(Box<Spanned<Self>>, Box<Spanned<Self>>, Box<Spanned<Self>>),
    Tuple(Vec<Spanned<Self>>),
    Arguments(Vec<Spanned<Self>>),
    List(Vec<Spanned<Self>>),
    Dict(Vec<Spanned<Self>>),
    Let(Box<Spanned<Pattern<'src>>>, Box<Spanned<Self>>),
    Box(&'src str, Box<Spanned<Self>>),
    Synthetic(Box<Spanned<Self>>, Box<Spanned<Self>>, Vec<Spanned<Self>>),
    When(Vec<Spanned<WhenClause<'src>>>),
    Match(Box<Spanned<Self>>, Vec<MatchClause<'src>>),
    Fn(&'src str, Vec<MatchClause<'src>>),
    FnDeclaration(&'src str),
    Panic(Box<Spanned<Self>>),
    Do(Vec<Spanned<Self>>),
    Repeat(Box<Spanned<Self>>, Box<Spanned<Self>>),
    Splat(&'src str),
    Pair(&'src str, Box<Spanned<Self>>),
    Loop(Box<Spanned<Self>>, Vec<MatchClause<'src>>),
    Recur(Vec<Spanned<Self>>),
}

impl<'src> fmt::Display for Ast<'src> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Ast::Error => write!(f, "Error"),
            Ast::Nil => write!(f, "nil"),
            Ast::String(s) => write!(f, "String: \"{}\"", s),
            Ast::Boolean(b) => write!(f, "Boolean: {}", b),
            Ast::Number(n) => write!(f, "Number: {}", n),
            Ast::Keyword(k) => write!(f, "Keyword: :{}", k),
            Ast::Word(w) => write!(f, "Word: {}", w),
            Ast::Block(b) => write!(
                f,
                "Block: <{}>",
                b.iter()
                    .map(|(line, _)| line.to_string())
                    .collect::<Vec<_>>()
                    .join("\n")
            ),
            Ast::If(cond, then_branch, else_branch) => write!(
                f,
                "If: {} Then: {} Else: {}",
                cond.0, then_branch.0, else_branch.0
            ),
            Ast::Let(pattern, expression) => {
                write!(f, "Let: {} = {}", pattern.0, expression.0)
            }
            Ast::Dict(entries) => write!(
                f,
                "#{{{}}}",
                entries
                    .iter()
                    .map(|pair| (*pair).0.to_string())
                    .collect::<Vec<_>>()
                    .join(", ")
            ),
            Ast::List(l) => write!(
                f,
                "List: [{}]",
                l.iter()
                    .map(|(line, _)| line.to_string())
                    .collect::<Vec<_>>()
                    .join("\n")
            ),
            Ast::Tuple(t) | Ast::Arguments(t) => write!(
                f,
                "Tuple: ({})",
                t.iter()
                    .map(|(line, _)| line.to_string())
                    .collect::<Vec<_>>()
                    .join("\n")
            ),
            Ast::Synthetic(root, first, rest) => write!(
                f,
                "Synth: [{}, {}, {}]",
                root.0,
                first.0,
                rest.iter()
                    .map(|(term, _)| term.to_string())
                    .collect::<Vec<_>>()
                    .join("\n")
            ),
            Ast::When(clauses) => write!(
                f,
                "When: [{}]",
                clauses
                    .iter()
                    .map(|clause| clause.0.to_string())
                    .collect::<Vec<_>>()
                    .join("\n")
            ),
            Ast::Placeholder => todo!(),
            Ast::Box(_name, _rhs) => todo!(),
            Ast::Match(value, clauses) => {
                write!(
                    f,
                    "match: {} with {}",
                    &value.0.to_string(),
                    clauses
                        .iter()
                        .map(|clause| clause.to_string())
                        .collect::<Vec<_>>()
                        .join("\n")
                )
            }
            Ast::Fn(name, clauses) => {
                write!(
                    f,
                    "fn: {}\n{}",
                    name,
                    clauses
                        .iter()
                        .map(|clause| clause.to_string())
                        .collect::<Vec<_>>()
                        .join("\n")
                )
            }
            Ast::FnDeclaration(_name) => todo!(),
            Ast::Panic(_expr) => todo!(),
            Ast::Do(terms) => {
                write!(
                    f,
                    "do: {}",
                    terms
                        .iter()
                        .map(|(term, _)| term.to_string())
                        .collect::<Vec<_>>()
                        .join(" > ")
                )
            }
            Ast::Repeat(_times, _body) => todo!(),
            Ast::Splat(word) => {
                write!(f, "splat: {}", word)
            }
            Ast::Pair(k, v) => {
                write!(f, "pair: {} {}", k, (*v).0.to_string())
            }
            Ast::Loop(init, body) => {
                write!(
                    f,
                    "loop: {} with {}",
                    (*init).0.to_string(),
                    body.iter()
                        .map(|clause| clause.to_string())
                        .collect::<Vec<_>>()
                        .join("\n")
                )
            }
            Ast::Recur(args) => {
                write!(
                    f,
                    "recur: {}",
                    args.iter()
                        .map(|(arg, _)| arg.to_string())
                        .collect::<Vec<_>>()
                        .join(", ")
                )
            }
        }
    }
}

#[derive(Clone, Debug, PartialEq)]
pub struct PairPattern<'src> {
    pub key: &'src str,
    pub patt: Spanned<Pattern<'src>>,
}

impl<'src> fmt::Display for PairPattern<'src> {
    fn fmt(self: &PairPattern<'src>, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "pair pattern: {}: {}", self.key, self.patt.0)
    }
}

#[derive(Clone, Debug, PartialEq)]
pub enum Pattern<'src> {
    Nil,
    Boolean(bool),
    Number(f64),
    String(&'src str),
    Keyword(&'src str),
    Word(&'src str),
    As(&'src str, &'src str),
    Splattern(Box<Spanned<Self>>),
    Placeholder,
    Tuple(Vec<Spanned<Self>>),
    List(Vec<Spanned<Self>>),
    Pair(&'src str, Box<Spanned<Self>>),
    Dict(Vec<Spanned<Self>>),
}

impl<'src> fmt::Display for Pattern<'src> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Pattern::Nil => write!(f, "nil"),
            Pattern::Boolean(b) => write!(f, "{}", b),
            Pattern::Number(n) => write!(f, "{}", n),
            Pattern::String(s) => write!(f, "{}", s),
            Pattern::Keyword(k) => write!(f, ":{}", k),
            Pattern::Word(w) => write!(f, "{}", w),
            Pattern::As(w, t) => write!(f, "{} as {}", w, t),
            Pattern::Splattern(p) => write!(f, "...{}", (*p).0.to_string()),
            Pattern::Placeholder => write!(f, "_"),
            Pattern::Tuple(t) => write!(
                f,
                "({})",
                t.iter()
                    .map(|x| x.0.to_string())
                    .collect::<Vec<_>>()
                    .join(", ")
            ),
            Pattern::List(l) => write!(
                f,
                "({})",
                l.iter()
                    .map(|x| x.0.to_string())
                    .collect::<Vec<_>>()
                    .join(", ")
            ),
            Pattern::Dict(entries) => write!(
                f,
                "#{{{}}}",
                entries
                    .iter()
                    .map(|(pair, _)| pair.to_string())
                    .collect::<Vec<_>>()
                    .join(", ")
            ),
            Pattern::Pair(key, value) => write!(f, ":{} {}", key, value.0),
        }
    }
}

pub fn parser<'src, I>(
) -> impl Parser<'src, I, Spanned<Ast<'src>>, extra::Err<Rich<'src, Token<'src>, Span>>> + Clone
where
    I: ValueInput<'src, Token = Token<'src>, Span = Span>,
{
    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("_") => Pattern::Placeholder}.map_with(|p, e| (p, e.span()));

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

    let atom_pattern = select! {
        Token::Nil => Pattern::Nil,
        Token::Boolean(b) => Pattern::Boolean(b),
        Token::Number(n) => Pattern::Number(n),
        Token::Keyword(k) => Pattern::Keyword(k),
        // todo: actual string patterns
        Token::String(s) => Pattern::String(s)
    }
    .map_with(|a, e| (a, e.span()));

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

    let splattable = word_pattern.clone().or(placeholder_pattern.clone());

    let patt_splat = just(Token::Punctuation("..."))
        .ignore_then(splattable)
        .map_with(|x, e| (Pattern::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| (Pattern::Tuple(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| (Pattern::List(list), e.span()));

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

    let shorthand_pattern = select! {Token::Word(w) => w}.map_with(|w, e| {
        (
            Pattern::Pair(w, Box::new((Pattern::Word(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| (Pattern::Dict(dict), e.span()));

    let keyword = select! {Token::Keyword(k) => Ast::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| (Pattern::As(w, t), e.span()));

    pattern.define(
        atom_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("_") => Ast::Placeholder}.map_with(|p, e| (p, e.span()));

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

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

    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| (Ast::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| (Ast::Arguments(args), e.span()));

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

    let synth_term = keyword.clone().or(args);

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

    let splat = just(Token::Punctuation("..."))
        .ignore_then(word.clone())
        .map_with(|(w, _), e| {
            (
                Ast::Splat(if let Ast::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| (Ast::List(list), e.span()));

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

    let shorthand = select! {Token::Word(w) => w}
        .map_with(|w, e| (Ast::Pair(w, Box::new((Ast::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| (Ast::Dict(dict), e.span()));

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

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

    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| (Ast::Block(block), e.span()))
        .recover_with(via_parser(nested_delimiters(
            Token::Punctuation("{"),
            Token::Punctuation("}"),
            [
                (Token::Punctuation("("), Token::Punctuation(")")),
                (Token::Punctuation("["), Token::Punctuation("]")),
            ],
            |span| (Ast::Error, span),
        )));

    let if_ = just(Token::Reserved("if"))
        .ignore_then(simple.clone())
        .then_ignore(just(Token::Reserved("then")))
        .then(expr.clone())
        .then_ignore(just(Token::Reserved("else")))
        .then(expr.clone())
        .map_with(|((condition, then_branch), else_branch), e| {
            (
                Ast::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 { cond, 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| (Ast::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), _| MatchClause {
            patt,
            guard: Some(guard),
            body,
        });

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

    let match_ = just(Token::Reserved("match"))
        .ignore_then(simple.clone())
        .then_ignore(just(Token::Reserved("with")))
        .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(|(expr, clauses), e| (Ast::Match(Box::new(expr), clauses), e.span()));

    let conditional = when.or(if_).or(match_);

    //todo:
    // * [x] do
    // * [ ] loop
    // * [x] repeat
    // * [x] panic!

    let panic = just(Token::Reserved("panic!"))
        .ignore_then(nonbinding.clone())
        .map_with(|expr, e| (Ast::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| (Ast::Do(exprs), e.span()));

    let repeat = just(Token::Reserved("repeat"))
        .ignore_then(simple.clone())
        .then(block.clone())
        .map_with(|(count, body), e| (Ast::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), _| MatchClause {
            patt,
            body,
            guard: Some(guard),
        })
        .labelled("function clause");

    let fn_unguarded = tuple_pattern
        .clone()
        .then_ignore(just(Token::Punctuation("->")))
        .then(nonbinding.clone())
        .map_with(|(patt, body), _| MatchClause {
            patt,
            body,
            guard: None,
        })
        .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| (Ast::Fn("anonymous", vec![clause]), e.span()));

    let fn_multiclause = fn_clause
        .clone()
        .separated_by(terminators.clone())
        .allow_leading()
        .allow_trailing()
        .collect()
        .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(fn_multiclause.clone().or(fn_single_clause.clone()))
        .map_with(|(init, body), e| (Ast::Loop(Box::new(init), body), e.span()));

    nonbinding.define(
        simple
            .clone()
            .or(conditional)
            .or(block)
            .or(lambda)
            .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| {
            (Ast::Let(Box::new(pattern), Box::new(expression)), e.span())
        });

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

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

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

    let fn_compound = just(Token::Reserved("fn"))
        .ignore_then(word.clone())
        .then(fn_multiclause.clone())
        .map_with(|(word, clauses), e| {
            let name = if let Ast::Word(word) = word.0 {
                word
            } else {
                unreachable!()
            };
            (Ast::Fn(name, clauses), 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| (Ast::Block(exprs), e.span()));

    script
}