//! A `macro_rules!` expander.
//!
//! The alternative was translating each `macro_rules!` into a Nim `template`,
//! which is the shape-level correspondence: 36% of the 1,833 definitions in a
//! 400-crate sample are a single rule with no repetition, which is exactly a
//! template, and another 40% map to a Nim `macro` over `varargs`.
//!
//! That correspondence is real but it does not survive contact with this
//! project's approach. A Nim template body is *untyped*: it is substituted and
//! only then type-checked. Our lowering is type-directed throughout — it needs
//! a type to choose `div` over `/`, to size a `cast`, to pick an integer
//! literal's width. Translating a macro body would mean lowering Rust with no
//! type information, which is precisely the guessing the project refuses.
//!
//! Expanding at the call site does not have that problem: the expansion is
//! ordinary Rust in a context where types are known, so it lowers like any
//! other code. Same applicability, faithful output.
//!
//! Implemented here: a single rule, no repetition, no `:tt`. That is the 36%.

use proc_macro2::{Delimiter, Group, Ident, TokenStream, TokenTree};
use std::collections::HashMap;

/// One element of a matcher: a captured fragment, or a token to match exactly.
#[derive(Debug, Clone)]
enum Pat {
    Frag(String),
    Tok(String),
    /// A delimited group, matched recursively.
    Group(Delimiter, Vec<Pat>),
    /// `$( .. ) sep? op` — a repetition. `op` is `*`, `+` or `?`.
    Rep { inner: Vec<Pat>, sep: Option<String>, op: char },
}

/// What a fragment captured: once, or once per repetition.
#[derive(Debug, Clone)]
enum Cap {
    One(TokenStream),
    Seq(Vec<TokenStream>),
}

type Binds = HashMap<String, Cap>;

#[derive(Debug, Clone)]
struct Rule {
    pattern: Vec<Pat>,
    body: TokenStream,
}

#[derive(Debug, Clone)]
pub struct MacroDef {
    /// Rules in source order. Rust tries them top to bottom and the first
    /// whose matcher matches wins, so order is semantics, not style.
    rules: Vec<Rule>,
}

/// Parse `macro_rules!`'s body: a `;`-separated list of
/// `( $matcher ) => { $transcriber }`.
pub fn parse(tokens: TokenStream) -> Result<MacroDef, String> {
    let t: Vec<TokenTree> = tokens.into_iter().collect();
    let mut rules = Vec::new();
    let mut i = 0;
    while i < t.len() {
        let matcher = match &t[i] {
            TokenTree::Group(g) => g.clone(),
            other => return Err(format!("expected a matcher group, found `{other}`")),
        };
        let arrow: String = t[i + 1..].iter().take(2).map(|t| t.to_string()).collect();
        if arrow != "=>" {
            return Err("expected `=>` after a matcher".into());
        }
        let body = match t.get(i + 3) {
            Some(TokenTree::Group(g)) => g.stream(),
            _ => return Err("expected a transcriber group".into()),
        };
        rules.push(Rule { pattern: parse_pattern(matcher.stream())?, body });
        i += 4;
        // The separator is optional after the last rule.
        if i < t.len() && t[i].to_string() == ";" {
            i += 1;
        }
    }
    if rules.is_empty() {
        return Err("a macro with no rules".into());
    }
    Ok(MacroDef { rules })
}

fn parse_pattern(ts: TokenStream) -> Result<Vec<Pat>, String> {
    let mut out = Vec::new();
    let mut it = ts.into_iter().peekable();
    while let Some(t) = it.next() {
        match t {
            TokenTree::Punct(p) if p.as_char() == '$' => {
                match it.next() {
                    Some(TokenTree::Ident(name)) => {
                        // `$name:fragment`
                        match it.next() {
                            Some(TokenTree::Punct(c)) if c.as_char() == ':' => {}
                            _ => return Err("expected `:` after a fragment name".into()),
                        }
                        let kind = match it.next() {
                            Some(TokenTree::Ident(k)) => k.to_string(),
                            _ => return Err("expected a fragment specifier".into()),
                        };
                        if kind == "tt" {
                            return Err(
                                "`:tt` makes a macro a token-tree interpreter, which \
                                 has no mechanical translation"
                                    .into(),
                            );
                        }
                        out.push(Pat::Frag(name.to_string()));
                    }
                    Some(TokenTree::Group(g)) => {
                        let inner = parse_pattern(g.stream())?;
                        if inner.iter().any(|p| matches!(p, Pat::Rep { .. })) {
                            return Err("nested `$(..)` repetition is not implemented yet".into());
                        }
                        // `$( .. ) sep? op`: an optional separator token, then
                        // the operator.
                        let mut sep = None;
                        let op = loop {
                            match it.next() {
                                Some(TokenTree::Punct(p))
                                    if matches!(p.as_char(), '*' | '+' | '?') =>
                                {
                                    break p.as_char()
                                }
                                Some(t) => {
                                    if sep.is_some() {
                                        return Err(
                                            "a repetition separator must be one token".into()
                                        );
                                    }
                                    sep = Some(t.to_string());
                                }
                                None => return Err("a repetition needs `*`, `+` or `?`".into()),
                            }
                        };
                        out.push(Pat::Rep { inner, sep, op });
                    }
                    _ => return Err("unexpected token after `$`".into()),
                }
            }
            TokenTree::Group(g) => {
                out.push(Pat::Group(g.delimiter(), parse_pattern(g.stream())?))
            }
            other => out.push(Pat::Tok(other.to_string())),
        }
    }
    Ok(out)
}

impl MacroDef {
    /// Try each rule in order and substitute using the first that matches,
    /// which is what Rust does.
    pub fn expand(&self, input: TokenStream) -> Result<TokenStream, String> {
        let toks: Vec<TokenTree> = input.into_iter().collect();
        let mut why = Vec::new();
        for (n, rule) in self.rules.iter().enumerate() {
            let mut binds = Binds::new();
            match match_seq(&rule.pattern, &toks, &mut binds, None) {
                Ok(used) if used == toks.len() => {
                    return Ok(substitute(rule.body.clone(), &binds))
                }
                Ok(_) => why.push(format!(
                    "rule {}: matched, but left tokens over",
                    n + 1
                )),
                Err(e) => why.push(format!("rule {}: {e}", n + 1)),
            }
        }
        Err(format!("no rule matched ({})", why.join("; ")))
    }
}

/// Match `pats` against the front of `toks`, returning how many were consumed.
/// `outer_stop` is where a trailing fragment must stop when the pattern
/// itself does not say: inside `$(, $rest:expr)*` that is the repetition's own
/// leading comma, which is where the next iteration begins.
fn match_seq(
    pats: &[Pat],
    toks: &[TokenTree],
    binds: &mut Binds,
    outer_stop: Option<&str>,
) -> Result<usize, String> {
    let mut i = 0;
    for (k, p) in pats.iter().enumerate() {
        match p {
            Pat::Tok(s) => {
                let t = toks.get(i).ok_or("the invocation ends before the matcher does")?;
                if t.to_string() != *s {
                    return Err(format!("expected `{s}`, found `{t}`"));
                }
                i += 1;
            }
            Pat::Group(d, inner) => {
                let Some(TokenTree::Group(g)) = toks.get(i) else {
                    return Err("expected a delimited group".into());
                };
                if g.delimiter() != *d {
                    return Err("mismatched delimiter".into());
                }
                let sub: Vec<TokenTree> = g.stream().into_iter().collect();
                let n = match_seq(inner, &sub, binds, None)?;
                if n != sub.len() {
                    return Err("group has tokens the matcher does not consume".into());
                }
                i += 1;
            }
            Pat::Rep { inner, sep, op } => {
                // Everything after the repetition that is a literal token
                // marks where it has to stop.
                let stop = next_literal(&pats[k + 1..]).or_else(|| outer_stop.map(str::to_string));
                // A trailing fragment inside the repetition stops at the
                // separator if there is one, or otherwise at whatever starts
                // the next iteration.
                let inner_stop = sep.clone().or_else(|| next_literal(inner));
                let names = frag_names(inner);
                let mut count = 0usize;
                for n in &names {
                    binds.insert(n.clone(), Cap::Seq(Vec::new()));
                }
                while i < toks.len() {
                    if let Some(stop) = &stop {
                        if toks[i].to_string() == *stop {
                            break;
                        }
                    }
                    let mut one = Binds::new();
                    let used = match_seq(inner, &toks[i..], &mut one, inner_stop.as_deref())?;
                    if used == 0 {
                        break;
                    }
                    for n in &names {
                        let v = match one.remove(n) {
                            Some(Cap::One(ts)) => ts,
                            _ => return Err(format!("`${n}` did not capture in a repetition")),
                        };
                        match binds.get_mut(n) {
                            Some(Cap::Seq(v0)) => v0.push(v),
                            _ => unreachable!(),
                        }
                    }
                    count += 1;
                    i += used;
                    match sep {
                        Some(sp) if i < toks.len() && toks[i].to_string() == *sp => i += 1,
                        Some(_) => break,
                        None => {}
                    }
                    if *op == '?' {
                        break;
                    }
                }
                if *op == '+' && count == 0 {
                    return Err("`$(..)+` needs at least one repetition".into());
                }
            }
            Pat::Frag(name) => {
                // A fragment runs to the next literal token in the matcher, or
                // to the end. That is what makes `$a:expr, $b:expr` split on
                // the comma rather than swallowing it.
                let stop = next_literal(&pats[k + 1..])
                    .or_else(|| outer_stop.map(str::to_string));
                let start = i;
                let mut depth = 0i32;
                while i < toks.len() {
                    let s = toks[i].to_string();
                    if depth == 0 {
                        if let Some(stop) = &stop {
                            if s == *stop {
                                break;
                            }
                        }
                        // A fragment never spans a top-level `,` or `;`: an
                        // `expr` is one expression, and a comma at this level
                        // separates arguments rather than belonging to one.
                        // A comma inside brackets is within a `Group` token,
                        // so it is not at this level at all.
                        if s == "," || s == ";" {
                            break;
                        }
                    }
                    // `<`/`>` are counted so a generic argument list is not
                    // split, but `>` also closes `=>` and `->`, so the depth
                    // never goes negative.
                    match &toks[i] {
                        TokenTree::Punct(p) if p.as_char() == '<' => depth += 1,
                        TokenTree::Punct(p) if p.as_char() == '>' => depth = (depth - 1).max(0),
                        _ => {}
                    }
                    i += 1;
                }
                if i == start {
                    return Err(format!("nothing matched fragment `${name}`"));
                }
                binds.insert(
                    name.clone(),
                    Cap::One(toks[start..i].iter().cloned().collect::<TokenStream>()),
                );
            }
        }
    }
    Ok(i)
}

/// The first literal token that can follow, looking through a repetition:
/// after `$first:expr` in `($first:expr $(, $rest:expr)*)` the next literal is
/// the repetition's leading comma, which is where `$first` has to stop.
fn next_literal(pats: &[Pat]) -> Option<String> {
    for p in pats {
        match p {
            Pat::Tok(s) => return Some(s.clone()),
            Pat::Rep { inner, sep, .. } => {
                if let Some(s) = next_literal(inner) {
                    return Some(s);
                }
                if let Some(s) = sep {
                    return Some(s.clone());
                }
            }
            Pat::Group(..) | Pat::Frag(_) => return None,
        }
    }
    None
}

/// The fragment names a pattern captures.
fn frag_names(pats: &[Pat]) -> Vec<String> {
    let mut out = Vec::new();
    for p in pats {
        match p {
            Pat::Frag(n) => out.push(n.clone()),
            Pat::Group(_, inner) => out.extend(frag_names(inner)),
            Pat::Rep { inner, .. } => out.extend(frag_names(inner)),
            Pat::Tok(_) => {}
        }
    }
    out
}

/// Replace every `$name` in the body with what it captured, and expand every
/// `$( .. ) sep? op` once per repetition.
fn substitute(ts: TokenStream, binds: &Binds) -> TokenStream {
    let mut out = Vec::new();
    let mut it = ts.into_iter().peekable();
    while let Some(t) = it.next() {
        match t {
            TokenTree::Punct(p) if p.as_char() == '$' => {
                // `$( .. ) sep? op` in the transcriber.
                if let Some(TokenTree::Group(g)) = it.peek() {
                    let g = g.clone();
                    it.next();
                    let mut sep = None;
                    loop {
                        match it.peek() {
                            Some(TokenTree::Punct(p))
                                if matches!(p.as_char(), '*' | '+' | '?') =>
                            {
                                it.next();
                                break;
                            }
                            Some(t) => {
                                sep = Some(t.clone());
                                it.next();
                            }
                            None => break,
                        }
                    }
                    let names: Vec<String> = fragments_in(g.stream())
                        .into_iter()
                        .filter(|n| matches!(binds.get(n), Some(Cap::Seq(_))))
                        .collect();
                    let n = names
                        .iter()
                        .filter_map(|n| match binds.get(n) {
                            Some(Cap::Seq(v)) => Some(v.len()),
                            _ => None,
                        })
                        .max()
                        .unwrap_or(0);
                    for idx in 0..n {
                        let mut one: Binds = binds.clone();
                        for nm in &names {
                            if let Some(Cap::Seq(v)) = binds.get(nm) {
                                if let Some(x) = v.get(idx) {
                                    one.insert(nm.clone(), Cap::One(x.clone()));
                                }
                            }
                        }
                        if idx > 0 {
                            if let Some(s) = &sep {
                                out.push(s.clone());
                            }
                        }
                        out.extend(substitute(g.stream(), &one));
                    }
                    continue;
                }
                match it.peek() {
                    Some(TokenTree::Ident(name)) => {
                        let n = name.to_string();
                        it.next();
                        match binds.get(&n) {
                            // Parenthesised so that a captured expression keeps
                            // its own precedence, as Rust's `expr` fragments do.
                            Some(Cap::One(v)) => out.push(TokenTree::Group(Group::new(
                                Delimiter::Parenthesis,
                                v.clone(),
                            ))),
                            Some(Cap::Seq(_)) => {
                                // Used outside a `$(..)`; Rust rejects this too.
                                out.push(TokenTree::Punct(p));
                                out.push(TokenTree::Ident(Ident::new(
                                    &n,
                                    proc_macro2::Span::call_site(),
                                )));
                            }
                            None => {
                                out.push(TokenTree::Punct(p));
                                out.push(TokenTree::Ident(Ident::new(
                                    &n,
                                    proc_macro2::Span::call_site(),
                                )));
                            }
                        }
                    }
                    _ => out.push(TokenTree::Punct(p)),
                }
            }
            TokenTree::Group(g) => {
                let inner = substitute(g.stream(), binds);
                out.push(TokenTree::Group(Group::new(g.delimiter(), inner)));
            }
            other => out.push(other),
        }
    }
    out.into_iter().collect()
}

/// The `$name` references appearing in a token stream.
fn fragments_in(ts: TokenStream) -> Vec<String> {
    let mut out = Vec::new();
    let mut it = ts.into_iter().peekable();
    while let Some(t) = it.next() {
        match t {
            TokenTree::Punct(p) if p.as_char() == '$' => {
                if let Some(TokenTree::Ident(n)) = it.peek() {
                    out.push(n.to_string());
                    it.next();
                }
            }
            TokenTree::Group(g) => out.extend(fragments_in(g.stream())),
            _ => {}
        }
    }
    out
}