gollum-parser 0.4.0

//! Recursive-descent nom 8 parser for Gollum source text.

use chrono::{DateTime, NaiveDate, NaiveDateTime, Utc};
use nom::{
    branch::alt,
    combinator::{map, opt},
    multi::{many0, separated_list0, separated_list1},
    Parser,
};

use gollum_ast::{
    BinOpKind, BodyGoal, Directive, Expr, Fact, Interval, Item, ModalAnnotation, Objective,
    PlainClause, Predicate, Probabilistic, Query, Rule, Term,
};

use crate::error::Error;
use crate::lexer::Token;

const NANOS_PER_SECOND: i128 = 1_000_000_000;
const NANOS_PER_MINUTE: i128 = 60 * NANOS_PER_SECOND;
const NANOS_PER_HOUR: i128 = 60 * NANOS_PER_MINUTE;
const NANOS_PER_DAY: i128 = 24 * NANOS_PER_HOUR;
const NANOS_PER_WEEK: i128 = 7 * NANOS_PER_DAY;
const NANOS_PER_YEAR: i128 = 36525 * NANOS_PER_DAY / 100;

// ---------------------------------------------------------------------------
// Input wrapper
// ---------------------------------------------------------------------------

#[derive(Clone, Copy, Debug)]
struct Input<'a>(&'a [Token]);

// ---------------------------------------------------------------------------
// Iterators for nom::Input
// ---------------------------------------------------------------------------

pub struct TokenIterator<'a> {
    tokens: &'a [Token],
    position: usize,
}

impl<'a> Iterator for TokenIterator<'a> {
    type Item = Token;
    fn next(&mut self) -> Option<Token> {
        if self.position >= self.tokens.len() {
            None
        } else {
            let t = self.tokens[self.position].clone();
            self.position += 1;
            Some(t)
        }
    }
}

pub struct IndexTokenIterator<'a> {
    tokens: &'a [Token],
    position: usize,
}

impl<'a> Iterator for IndexTokenIterator<'a> {
    type Item = (usize, Token);
    fn next(&mut self) -> Option<(usize, Token)> {
        if self.position >= self.tokens.len() {
            None
        } else {
            let idx = self.position;
            let t = self.tokens[self.position].clone();
            self.position += 1;
            Some((idx, t))
        }
    }
}

impl<'a> nom::Input for Input<'a> {
    type Item = Token;
    type Iter = TokenIterator<'a>;
    type IterIndices = IndexTokenIterator<'a>;

    fn input_len(&self) -> usize {
        self.0.len()
    }

    fn iter_elements(&self) -> TokenIterator<'a> {
        TokenIterator {
            tokens: self.0,
            position: 0,
        }
    }

    fn iter_indices(&self) -> IndexTokenIterator<'a> {
        IndexTokenIterator {
            tokens: self.0,
            position: 0,
        }
    }

    fn position<P>(&self, predicate: P) -> Option<usize>
    where
        P: Fn(Self::Item) -> bool,
    {
        self.0.iter().position(|t| predicate(t.clone()))
    }

    fn slice_index(&self, count: usize) -> Result<usize, nom::Needed> {
        if count <= self.0.len() {
            Ok(count)
        } else {
            Err(nom::Needed::new(count - self.0.len()))
        }
    }

    fn take(&self, index: usize) -> Self {
        Input(&self.0[..index])
    }

    fn take_from(&self, index: usize) -> Self {
        Input(&self.0[index..])
    }

    /// Returns (suffix, prefix) = (remaining after index, first index elements)
    /// This matches nom's convention: IResult = (remaining, output)
    fn take_split(&self, index: usize) -> (Self, Self) {
        (Input(&self.0[index..]), Input(&self.0[..index]))
    }
}

// ---------------------------------------------------------------------------
// Error type for nom
// ---------------------------------------------------------------------------

#[derive(Debug)]
struct ParseErr;

impl<'a> nom::error::ParseError<Input<'a>> for ParseErr {
    fn from_error_kind(_: Input<'a>, _: nom::error::ErrorKind) -> Self {
        ParseErr
    }
    fn append(_: Input<'a>, _: nom::error::ErrorKind, other: Self) -> Self {
        other
    }
}

type IResult<'a, O> = nom::IResult<Input<'a>, O, ParseErr>;

// ---------------------------------------------------------------------------
// Basic token combinators
// ---------------------------------------------------------------------------

/// Match a specific token (by equality).
fn tok(expected: Token) -> impl Fn(Input<'_>) -> IResult<'_, Token> {
    move |input: Input<'_>| match input.0 {
        [first, rest @ ..] if *first == expected => Ok((Input(rest), first.clone())),
        _ => Err(nom::Err::Error(ParseErr)),
    }
}

/// Match Token::Atom and return the string.
fn atom_tok(input: Input<'_>) -> IResult<'_, String> {
    match input.0 {
        [Token::Atom(s), rest @ ..] => Ok((Input(rest), s.clone())),
        [Token::QuotedAtom(s), rest @ ..] => Ok((Input(rest), s.clone())),
        _ => Err(nom::Err::Error(ParseErr)),
    }
}

/// Match any name: Atom, QuotedAtom, or keyword tokens used as atoms.
fn name_tok(input: Input<'_>) -> IResult<'_, String> {
    match input.0 {
        [Token::Atom(s), rest @ ..] => Ok((Input(rest), s.clone())),
        [Token::QuotedAtom(s), rest @ ..] => Ok((Input(rest), s.clone())),
        [Token::Before, rest @ ..] => Ok((Input(rest), "before".to_string())),
        [Token::After, rest @ ..] => Ok((Input(rest), "after".to_string())),
        [Token::During, rest @ ..] => Ok((Input(rest), "during".to_string())),
        [Token::Until, rest @ ..] => Ok((Input(rest), "until".to_string())),
        [Token::Using, rest @ ..] => Ok((Input(rest), "using".to_string())),
        [Token::Is, rest @ ..] => Ok((Input(rest), "is".to_string())),
        [Token::In, rest @ ..] => Ok((Input(rest), "in".to_string())),
        _ => Err(nom::Err::Error(ParseErr)),
    }
}

/// Match Token::Var and return the string.
fn var_tok(input: Input<'_>) -> IResult<'_, String> {
    match input.0 {
        [Token::Var(s), rest @ ..] => Ok((Input(rest), s.clone())),
        _ => Err(nom::Err::Error(ParseErr)),
    }
}

/// Match Token::Integer and return the value.
fn int_tok(input: Input<'_>) -> IResult<'_, i64> {
    match input.0 {
        [Token::Integer(n), rest @ ..] => Ok((Input(rest), *n)),
        _ => Err(nom::Err::Error(ParseErr)),
    }
}

/// Match Token::Float and return the value.
fn float_tok(input: Input<'_>) -> IResult<'_, f64> {
    match input.0 {
        [Token::Float(n), rest @ ..] => Ok((Input(rest), *n)),
        _ => Err(nom::Err::Error(ParseErr)),
    }
}

/// Match Token::Str and return the string.
fn str_tok(input: Input<'_>) -> IResult<'_, String> {
    match input.0 {
        [Token::Str(s), rest @ ..] => Ok((Input(rest), s.clone())),
        _ => Err(nom::Err::Error(ParseErr)),
    }
}

/// Parse a tensor literal: `#[f1, f2, ...]` → `Term::Tensor(Vec<f32>)`.
///
/// After `TensorOpen` (`#[`), reads zero or more float/integer literals separated
/// by commas, then a closing `]`.  Integers are losslessly cast to `f32`.
fn parse_tensor_literal(input: Input<'_>) -> IResult<'_, Term> {
    let (mut input, _) = tok(Token::TensorOpen)(input)?;
    let mut elems: Vec<f32> = Vec::new();
    // Parse optional leading minus for the first element
    loop {
        // Try closing bracket first (handles empty tensor)
        if let Ok((rest, _)) = tok(Token::RBracket)(input) {
            return Ok((rest, Term::Tensor(elems)));
        }
        // Parse optional unary minus
        let (inp2, neg) = if let Ok((r, _)) = tok(Token::Minus)(input) {
            (r, true)
        } else {
            (input, false)
        };
        // Accept float or integer
        let (inp3, val) = if let Ok((r, f)) = float_tok(inp2) {
            (r, f as f32)
        } else if let Ok((r, n)) = int_tok(inp2) {
            (r, n as f32)
        } else {
            return Err(nom::Err::Error(ParseErr));
        };
        elems.push(if neg { -val } else { val });
        input = inp3;
        // Try comma (continue) or close bracket (done)
        if let Ok((rest, _)) = tok(Token::Comma)(input) {
            input = rest;
        } else {
            let (rest, _) = tok(Token::RBracket)(input)?;
            return Ok((rest, Term::Tensor(elems)));
        }
    }
}


fn type_name_tok(input: Input<'_>) -> IResult<'_, String> {
    match input.0 {
        [Token::Var(s), rest @ ..] => Ok((Input(rest), s.clone())),
        [Token::Atom(s), rest @ ..] => Ok((Input(rest), s.clone())),
        [Token::QuotedAtom(s), rest @ ..] => Ok((Input(rest), s.clone())),
        _ => Err(nom::Err::Error(ParseErr)),
    }
}

// ---------------------------------------------------------------------------
// Term parser
// ---------------------------------------------------------------------------

fn parse_term(input: Input<'_>) -> IResult<'_, Term> {
    parse_term_with_type(input)
}

/// Parse a single term from a token slice. Returns an error if parsing fails or leftover tokens remain.
pub fn parse_single_term(tokens: &[Token]) -> crate::Result<Term> {
    let input = Input(tokens);
    match parse_term(input) {
        Ok((remaining, term)) => {
            if remaining.0.is_empty() {
                Ok(term)
            } else {
                Err(Error::UnexpectedToken)
            }
        }
        Err(_) => Err(Error::UnexpectedToken),
    }
}

/// Parse a term, optionally followed by a type annotation or neural gradient annotation.
/// Handles `X:Type` and `X :: "model" :: "grad"` syntax.
fn parse_term_with_type(input: Input<'_>) -> IResult<'_, Term> {
    let (input, term) = parse_base_term(input)?;

    // Try type annotation (':')
    let (input, type_opt) = opt(parse_type_annotation).parse(input)?;

    // Try neural gradient annotation ('::')
    let (input, grad_opt) = opt(parse_neural_gradient_annotation).parse(input)?;

    match (type_opt, grad_opt) {
        (Some(type_name), Some((model_id, grad_id))) => {
            // Both type and neural gradient: X:Type :: "model" :: "grad"
            Ok((
                input,
                Term::NeuralGradient {
                    term: Box::new(Term::TypeAnnotated {
                        term: Box::new(term),
                        type_name,
                    }),
                    model_id,
                    grad_id,
                },
            ))
        }
        (Some(type_name), None) => {
            // Just type annotation: X:Type
            Ok((
                input,
                Term::TypeAnnotated {
                    term: Box::new(term),
                    type_name,
                },
            ))
        }
        (None, Some((model_id, grad_id))) => {
            // Just neural gradient: X :: "model" :: "grad"
            Ok((
                input,
                Term::NeuralGradient {
                    term: Box::new(term),
                    model_id,
                    grad_id,
                },
            ))
        }
        (None, None) => Ok((input, term)),
    }
}

/// Parse a range term: `Low..High` → `Compound("..", [Low, High])`.
fn parse_range_term(input: Input<'_>) -> IResult<'_, Term> {
    let (input, lo) = int_tok(input)?;
    let (input, _) = tok(Token::DotDot)(input)?;
    let (input, hi) = int_tok(input)?;
    Ok((
        input,
        Term::Compound("..".into(), vec![Term::Integer(lo), Term::Integer(hi)]),
    ))
}

/// Parse a base term without type annotation.
fn parse_base_term(input: Input<'_>) -> IResult<'_, Term> {
    alt((
        parse_range_term,
        parse_paren_term,
        parse_tensor_literal,
        parse_compound,
        parse_list,
        map(var_tok, Term::Variable),
        map(int_tok, Term::Integer),
        map(float_tok, Term::Float),
        map(str_tok, Term::Str),
        map(tok(Token::Anon), |_| Term::Anon),
        map(name_tok, Term::Atom),
    ))
    .parse(input)
}

/// Parse a parenthesised term or conjunction: `(T)` or `(A,B,C,...)`.
///
/// In Prolog, `(A,B)` used as a *term* (e.g. as an argument to a functor) is
/// syntactic sugar for the compound `','(A,B)`.  Multiple conjuncts are
/// right-associated: `(A,B,C)` → `','(A, ','(B,C))`.
///
/// A single-element group `(T)` is transparent — it just returns `T`.
fn parse_paren_term(input: Input<'_>) -> IResult<'_, Term> {
    let (input, _) = tok(Token::LParen)(input)?;
    let (input, first) = parse_term(input)?;

    let mut terms = vec![first];
    let mut rest = input;
    while let Ok((inp, _)) = tok(Token::Comma)(rest) {
        match parse_term(inp) {
            Ok((inp2, t)) => {
                terms.push(t);
                rest = inp2;
            }
            Err(_) => break,
        }
    }
    let (rest, _) = tok(Token::RParen)(rest)?;

    let result = if terms.len() == 1 {
        terms.into_iter().next().unwrap()
    } else {
        // Build right-associative conjunction tree: (A,B,C) → ','(A, ','(B,C))
        let mut iter = terms.into_iter().rev();
        let last = iter.next().unwrap();
        iter.fold(last, |acc, t| Term::Compound(",".into(), vec![t, acc]))
    };

    Ok((rest, result))
}

/// Parse a type annotation: `:` followed by a type name.
fn parse_type_annotation(input: Input<'_>) -> IResult<'_, String> {
    let (input, _) = tok(Token::Colon)(input)?;
    type_name_tok(input)
}

/// Parse a neural gradient annotation: `:: "model_id" :: "grad_id"`
fn parse_neural_gradient_annotation(input: Input<'_>) -> IResult<'_, (String, String)> {
    let (input, _) = tok(Token::ColonColon)(input)?;
    let (input, model_id) = str_tok(input)?;
    let (input, _) = tok(Token::ColonColon)(input)?;
    let (input, grad_id) = str_tok(input)?;
    Ok((input, (model_id, grad_id)))
}

/// Parse compound term: `name(args...)`
fn parse_compound(input: Input<'_>) -> IResult<'_, Term> {
    let (input, name) = name_tok(input)?;
    let (input, _) = tok(Token::LParen)(input)?;
    let (input, args) = separated_list0(tok(Token::Comma), parse_term).parse(input)?;
    let (input, _) = tok(Token::RParen)(input)?;
    Ok((input, Term::Compound(name, args)))
}

/// Parse list: `[...]`
fn parse_list(input: Input<'_>) -> IResult<'_, Term> {
    let (input, _) = tok(Token::LBracket)(input)?;
    // empty list
    if let Ok((input, _)) = tok(Token::RBracket)(input) {
        return Ok((input, Term::List(vec![])));
    }
    // non-empty: head elements
    let (input, first) = parse_term(input)?;
    let mut heads = vec![first];
    let mut rest_input = input;
    while let Ok((inp, _)) = tok(Token::Comma)(rest_input) {
        let (inp, t) = parse_term(inp)?;
        heads.push(t);
        rest_input = inp;
    }
    // check for |tail
    if let Ok((inp, _)) = tok(Token::Pipe)(rest_input) {
        let (inp, tail) = parse_term(inp)?;
        let (inp, _) = tok(Token::RBracket)(inp)?;
        Ok((inp, Term::ListCons(heads, Box::new(tail))))
    } else {
        let (inp, _) = tok(Token::RBracket)(rest_input)?;
        Ok((inp, Term::List(heads)))
    }
}

// ---------------------------------------------------------------------------
// Expression parser (precedence climbing)
// ---------------------------------------------------------------------------

fn parse_expr(input: Input<'_>) -> IResult<'_, Expr> {
    parse_comparison(input)
}

fn parse_comparison(input: Input<'_>) -> IResult<'_, Expr> {
    let (input, lhs) = parse_additive(input)?;
    let op_result: IResult<'_, BinOpKind> = parse_cmp_op(input);
    if let Ok((input, op)) = op_result {
        let (input, rhs) = parse_additive(input)?;
        Ok((input, Expr::BinOp(op, Box::new(lhs), Box::new(rhs))))
    } else {
        Ok((input, lhs))
    }
}

fn parse_cmp_op(input: Input<'_>) -> IResult<'_, BinOpKind> {
    alt((
        map(tok(Token::ClpGte), |_| BinOpKind::ClpGte),
        map(tok(Token::ClpLte), |_| BinOpKind::ClpLte),
        map(tok(Token::ClpNeq), |_| BinOpKind::ClpNeq),
        map(tok(Token::ClpGt), |_| BinOpKind::ClpGt),
        map(tok(Token::ClpLt), |_| BinOpKind::ClpLt),
        map(tok(Token::ClpEq), |_| BinOpKind::ClpEq),
        map(tok(Token::In), |_| BinOpKind::ClpIn),
        map(tok(Token::Is), |_| BinOpKind::Is),
        map(tok(Token::ArithEq), |_| BinOpKind::ArithEq),
        map(tok(Token::ArithNeq), |_| BinOpKind::ArithNeq),
        map(tok(Token::NotEq), |_| BinOpKind::NotUnify),
        map(tok(Token::Lte), |_| BinOpKind::Lte),
        map(tok(Token::Gte), |_| BinOpKind::Gte),
        map(tok(Token::Lt), |_| BinOpKind::Lt),
        map(tok(Token::Gt), |_| BinOpKind::Gt),
        map(tok(Token::NeuralUnify), |_| BinOpKind::NeuralUnify),
        map(tok(Token::Eq), |_| BinOpKind::Unify),
    ))
    .parse(input)
}

fn parse_additive(input: Input<'_>) -> IResult<'_, Expr> {
    let (mut input, mut lhs) = parse_multiplicative(input)?;
    loop {
        let op_result: IResult<'_, BinOpKind> = alt((
            map(tok(Token::Plus), |_| BinOpKind::Add),
            map(tok(Token::Minus), |_| BinOpKind::Sub),
        ))
        .parse(input);
        if let Ok((inp, op)) = op_result {
            let (inp, rhs) = parse_multiplicative(inp)?;
            lhs = Expr::BinOp(op, Box::new(lhs), Box::new(rhs));
            input = inp;
        } else {
            break;
        }
    }
    Ok((input, lhs))
}

fn parse_multiplicative(input: Input<'_>) -> IResult<'_, Expr> {
    let (mut input, mut lhs) = parse_primary(input)?;
    loop {
        let op_result: IResult<'_, BinOpKind> = alt((
            map(tok(Token::Star), |_| BinOpKind::Mul),
            map(tok(Token::Slash), |_| BinOpKind::Div),
            map(tok(Token::Mod), |_| BinOpKind::Mod),
        ))
        .parse(input);
        if let Ok((inp, op)) = op_result {
            let (inp, rhs) = parse_primary(inp)?;
            lhs = Expr::BinOp(op, Box::new(lhs), Box::new(rhs));
            input = inp;
        } else {
            break;
        }
    }
    Ok((input, lhs))
}

fn parse_primary(input: Input<'_>) -> IResult<'_, Expr> {
    // Parenthesized expression
    if let Ok((inp, _)) = tok(Token::LParen)(input) {
        let (inp, e) = parse_expr(inp)?;
        let (inp, _) = tok(Token::RParen)(inp)?;
        return Ok((inp, e));
    }
    // Term
    map(parse_term, Expr::Term).parse(input)
}

// ---------------------------------------------------------------------------
// Body goal parser
// ---------------------------------------------------------------------------

fn parse_body_goal(input: Input<'_>) -> IResult<'_, BodyGoal> {
    // cut
    if let Ok((inp, _)) = tok(Token::Cut)(input) {
        return Ok((inp, BodyGoal::Cut));
    }
    // negation: \+ <goal> or not <goal>
    if let Ok((inp, _)) = tok(Token::NotPlus)(input) {
        let (inp, inner) = parse_body_goal(inp)?;
        return Ok((inp, BodyGoal::Not(Box::new(inner))));
    }
    if let Ok((inp, _)) = tok(Token::Not)(input) {
        let (inp, inner) = parse_body_goal(inp)?;
        return Ok((inp, BodyGoal::Not(Box::new(inner))));
    }
    // Parse as expression, then promote
    let (input, expr) = parse_expr(input)?;
    let goal = expr_to_goal(expr);
    Ok((input, goal))
}

/// Promote an expression to a body goal.
fn expr_to_goal(expr: Expr) -> BodyGoal {
    match expr {
        Expr::Term(Term::Compound(name, args)) => BodyGoal::Call(name, args),
        Expr::Term(Term::Atom(name)) => BodyGoal::Call(name, vec![]),
        other => BodyGoal::Expr(Box::new(other)),
    }
}

fn parse_body(input: Input<'_>) -> IResult<'_, Vec<BodyGoal>> {
    separated_list1(tok(Token::Comma), parse_body_goal).parse(input)
}

// ---------------------------------------------------------------------------
// Head predicate (for fact/rule head)
// ---------------------------------------------------------------------------

/// Parse predicate head: `name(args...)` or just `name`
fn parse_predicate(input: Input<'_>) -> IResult<'_, Predicate> {
    let (input, name) = name_tok(input)?;
    // optional args
    if let Ok((inp, _)) = tok(Token::LParen)(input) {
        let (inp, args) = separated_list0(tok(Token::Comma), parse_term).parse(inp)?;
        let (inp, _) = tok(Token::RParen)(inp)?;
        Ok((inp, Predicate { name, args }))
    } else {
        Ok((input, Predicate { name, args: vec![] }))
    }
}

// ---------------------------------------------------------------------------
// Directive body parsers
// ---------------------------------------------------------------------------

fn parse_directive_body(input: Input<'_>) -> IResult<'_, Directive> {
    alt((
        parse_dir_diff_neural,
        parse_dir_diff_neural_model,
        parse_dir_neural_model,
        parse_dir_neural_gen,
        parse_dir_neural_unify,
        parse_dir_neural,
        parse_dir_differentiable,
        parse_dir_type,
        parse_dir_pred,
        parse_dir_table,
        parse_dir_optimize,
    ))
    .parse(input)
}

fn parse_dir_type(input: Input<'_>) -> IResult<'_, Directive> {
    let (input, kw) = atom_tok(input)?;
    if kw != "type" {
        return Err(nom::Err::Error(ParseErr));
    }
    let (input, name) = type_name_tok(input)?;
    Ok((input, Directive::TypeDecl { name }))
}

fn parse_dir_pred(input: Input<'_>) -> IResult<'_, Directive> {
    let (input, kw) = atom_tok(input)?;
    if kw != "pred" {
        return Err(nom::Err::Error(ParseErr));
    }
    let (input, functor) = atom_tok(input)?;
    let (input, _) = tok(Token::LParen)(input)?;
    let (input, arg_types) = separated_list0(tok(Token::Comma), type_name_tok).parse(input)?;
    let (input, _) = tok(Token::RParen)(input)?;
    Ok((input, Directive::PredDecl { functor, arg_types }))
}

fn parse_dir_table(input: Input<'_>) -> IResult<'_, Directive> {
    let (input, kw) = atom_tok(input)?;
    if kw != "table" {
        return Err(nom::Err::Error(ParseErr));
    }
    let (input, functor) = atom_tok(input)?;
    let (input, _) = tok(Token::Slash)(input)?;
    let (input, arity) = int_tok(input)?;
    Ok((
        input,
        Directive::Table {
            functor,
            arity: arity as u32,
        },
    ))
}

fn parse_dir_neural_unify(input: Input<'_>) -> IResult<'_, Directive> {
    let (input, kw) = atom_tok(input)?;
    if kw != "neural_unify" {
        return Err(nom::Err::Error(ParseErr));
    }
    // Expect: threshold(<f64-or-int>)
    let (input, kw2) = atom_tok(input)?;
    if kw2 != "threshold" {
        return Err(nom::Err::Error(ParseErr));
    }
    let (input, _) = tok(Token::LParen)(input)?;
    // Accept either a float or an integer literal.
    let (input, threshold) = alt((
        map(float_tok, |v| v),
        map(int_tok, |v| v as f64),
    ))
    .parse(input)?;
    let (input, _) = tok(Token::RParen)(input)?;
    Ok((input, Directive::NeuralUnify { threshold }))
}

fn parse_dir_neural(input: Input<'_>) -> IResult<'_, Directive> {
    let (input, kw) = atom_tok(input)?;
    if kw != "neural" {
        return Err(nom::Err::Error(ParseErr));
    }
    let (input, functor) = atom_tok(input)?;

    // Accept either slash-form `name/arity` or typed `name(type, ...)`.
    if let Ok((input2, _)) = tok(Token::Slash)(input) {
        let (input2, arity) = int_tok(input2)?;
        // Parse optional trailing options: model(...), input_template(...), etc.
        let (input3, options) = parse_neural_options(input2)?;
        return Ok((
            input3,
            Directive::Neural {
                functor,
                arg_types: Vec::new(),
                arity: Some(arity as u32),
                options,
            },
        ));
    }

    let (input, _) = tok(Token::LParen)(input)?;
    let (input, arg_types) = separated_list0(tok(Token::Comma), type_name_tok).parse(input)?;
    let (input, _) = tok(Token::RParen)(input)?;
    // Parse optional trailing options.
    let (input, options) = parse_neural_options(input)?;
    Ok((
        input,
        Directive::Neural {
            functor,
            arg_types,
            arity: None,
            options,
        },
    ))
}

// Parse zero or more options following a neural directive.  Options are
// `key(value)` and may be separated by commas or whitespace.  The parser stops
// when it cannot parse a `key(` pattern.
fn parse_neural_options(input: Input<'_>) -> IResult<'_, Vec<(String, Term)>> {
    let mut opts: Vec<(String, Term)> = Vec::new();
    let mut cur = input;
    loop {
        // Consume any leading commas between options.
        if let Ok((after_comma, _)) = tok(Token::Comma)(cur) {
            cur = after_comma;
        }
        // Try to parse an option name followed by `(`.  If this fails, stop.
        let res = atom_tok(cur);
        if res.is_err() {
            break;
        }
        let (after_name, name) = res?;
        // Require a `(` after the name to be considered an option; otherwise stop.
        if let Ok((after_lparen, _)) = tok(Token::LParen)(after_name) {
            let (after_val, val) = parse_term(after_lparen)?;
            let (after_rparen, _) = tok(Token::RParen)(after_val)?;
            opts.push((name, val));
            cur = after_rparen;
            continue;
        } else {
            break;
        }
    }
    Ok((cur, opts))
}

fn parse_dir_neural_model(input: Input<'_>) -> IResult<'_, Directive> {
    let (input, kw) = atom_tok(input)?;
    if kw != "neural_model" {
        return Err(nom::Err::Error(ParseErr));
    }
    let (input, predicate) = atom_tok(input)?;
    let (input, _) = tok(Token::Using)(input)?;
    let (input, model_name) = str_tok(input)?;
    Ok((
        input,
        Directive::NeuralModel {
            predicate,
            model_name,
        },
    ))
}

fn parse_dir_neural_gen(input: Input<'_>) -> IResult<'_, Directive> {
    let (input, kw) = atom_tok(input)?;
    if kw != "neural_gen" {
        return Err(nom::Err::Error(ParseErr));
    }
    let (input, functor) = atom_tok(input)?;
    let (input, _) = tok(Token::Slash)(input)?;
    let (input, arity) = int_tok(input)?;
    Ok((
        input,
        Directive::NeuralGen {
            functor,
            arity: arity as u32,
        },
    ))
}

fn parse_dir_diff_neural_model(input: Input<'_>) -> IResult<'_, Directive> {
    let (input, kw) = atom_tok(input)?;
    if kw != "diff_neural_model" {
        return Err(nom::Err::Error(ParseErr));
    }
    let (input, predicate) = atom_tok(input)?;
    let (input, _) = tok(Token::Using)(input)?;
    let (input, model_name) = str_tok(input)?;
    Ok((
        input,
        Directive::DiffNeuralModel {
            predicate,
            model_name,
        },
    ))
}

fn parse_dir_differentiable(input: Input<'_>) -> IResult<'_, Directive> {
    let (input, kw) = atom_tok(input)?;
    if kw != "differentiable" {
        return Err(nom::Err::Error(ParseErr));
    }
    let (input, kw2) = atom_tok(input)?;
    if kw2 != "predicate" {
        return Err(nom::Err::Error(ParseErr));
    }
    let (input, functor) = atom_tok(input)?;
    let (input, _) = tok(Token::Slash)(input)?;
    let (input, arity) = int_tok(input)?;
    Ok((
        input,
        Directive::Differentiable {
            functor,
            arity: arity as u32,
        },
    ))
}

/// Parse `:- differentiable neural functor/arity using "model_name".`
fn parse_dir_diff_neural(input: Input<'_>) -> IResult<'_, Directive> {
    let (input, kw1) = atom_tok(input)?;
    if kw1 != "differentiable" {
        return Err(nom::Err::Error(ParseErr));
    }
    let (input, kw2) = atom_tok(input)?;
    if kw2 != "neural" {
        return Err(nom::Err::Error(ParseErr));
    }
    let (input, functor) = atom_tok(input)?;
    let (input, _) = tok(Token::Slash)(input)?;
    let (input, arity) = int_tok(input)?;
    let (input, _) = tok(Token::Using)(input)?;
    let (input, model_name) = str_tok(input)?;
    Ok((
        input,
        Directive::DiffNeural {
            functor,
            arity: arity as u32,
            model_name,
        },
    ))
}

fn parse_dir_optimize(input: Input<'_>) -> IResult<'_, Directive> {
    let (input, kw) = atom_tok(input)?;
    if kw != "optimize" {
        return Err(nom::Err::Error(ParseErr));
    }
    let (input, objective) = alt((
        map(tok(Token::Minimize), |_| Objective::Minimize),
        map(tok(Token::Maximize), |_| Objective::Maximize),
    ))
    .parse(input)?;
    let (input, target) = parse_term(input)?;
    Ok((input, Directive::Optimize { objective, target }))
}

// ---------------------------------------------------------------------------
// Top-level item parsers
// ---------------------------------------------------------------------------

fn parse_fact(input: Input<'_>) -> IResult<'_, Item> {
    // allow optional prefix modality (e.g. `□ pred(...)`) or postfix modality after the clause
    let (input, prefix_modality) = opt(parse_modal_annotation).parse(input)?;
    let (input, predicate) = parse_predicate(input)?;
    let (input, temporal_interval) = opt(parse_temporal_annotation).parse(input)?;
    let (input, trailing_modality) = opt(parse_modal_annotation).parse(input)?;
    let (input, diff_neural_ref) = opt(parse_neural_gradient_annotation).parse(input)?;
    let (input, _) = tok(Token::Dot)(input)?;
    let modality = prefix_modality.or(trailing_modality);
    Ok((
        input,
        Item::Fact(Fact {
            name: predicate.name,
            args: predicate.args,
            temporal_interval,
            modality,
            diff_neural_ref,
        }),
    ))
}

fn parse_temporal_annotation(input: Input<'_>) -> IResult<'_, Interval> {
    let (input, _) = tok(Token::At)(input)?;
    let (input, _) = tok(Token::LBracket)(input)?;
    let (input, start_ns) = parse_timestamp(input)?;
    let (input, _) = tok(Token::Comma)(input)?;
    let (input, end_ns) = parse_timestamp(input)?;
    let (input, _) = tok(Token::RBracket)(input)?;
    match Interval::new(start_ns, end_ns) {
        Some(interval) => Ok((input, interval)),
        None => Err(nom::Err::Error(ParseErr)),
    }
}

/// Parse a modal annotation following a clause (e.g. `□`, `◇`, optionally with an agent name).
fn parse_modal_annotation(input: Input<'_>) -> IResult<'_, ModalAnnotation> {
    // Try necessity
    if let Ok((input, _)) = tok(Token::Box)(input) {
        // optional agent: accept atom or quoted/str (not variables - agents must be ground)
        let (input, agent_opt) =
            opt(alt((map(atom_tok, |s| s), map(str_tok, |s| s)))).parse(input)?;
        let modality = match agent_opt {
            Some(a) => ModalAnnotation::necessity_for(a),
            None => ModalAnnotation::necessity(),
        };
        return Ok((input, modality));
    }
    // Try possibility
    if let Ok((input, _)) = tok(Token::Diamond)(input) {
        let (input, agent_opt) =
            opt(alt((map(atom_tok, |s| s), map(str_tok, |s| s)))).parse(input)?;
        let modality = match agent_opt {
            Some(a) => ModalAnnotation::possibility_for(a),
            None => ModalAnnotation::possibility(),
        };
        return Ok((input, modality));
    }
    Err(nom::Err::Error(ParseErr))
}

fn parse_timestamp(input: Input<'_>) -> IResult<'_, i128> {
    match input.0 {
        [Token::UnitLiteral((num, unit)), rest @ ..] => {
            let ns = unit_literal_to_ns(*num, unit).ok_or(nom::Err::Error(ParseErr))?;
            Ok((Input(rest), ns))
        }
        [Token::Atom(s), rest @ ..]
        | [Token::QuotedAtom(s), rest @ ..]
        | [Token::Str(s), rest @ ..] => {
            let ns = parse_iso_date(s).ok_or(nom::Err::Error(ParseErr))?;
            Ok((Input(rest), ns))
        }
        _ => Err(nom::Err::Error(ParseErr)),
    }
}

fn parse_iso_date(input: &str) -> Option<i128> {
    let input = input.trim();

    if let Ok(dt) = DateTime::parse_from_rfc3339(input) {
        return Some(ns_from_datetime(dt.with_timezone(&Utc)));
    }
    if let Ok(dt) = NaiveDateTime::parse_from_str(input, "%Y-%m-%dT%H:%M:%S%.f") {
        return Some(ns_from_naivedt(dt));
    }
    if let Ok(dt) = NaiveDateTime::parse_from_str(input, "%Y-%m-%dT%H:%M:%S") {
        return Some(ns_from_naivedt(dt));
    }
    if let Ok(dt) = NaiveDateTime::parse_from_str(input, "%Y-%m-%d %H:%M:%S") {
        return Some(ns_from_naivedt(dt));
    }
    if let Ok(d) = NaiveDate::parse_from_str(input, "%Y-%m-%d") {
        return Some(ns_from_naivedate(d));
    }
    None
}

fn ns_from_datetime(dt: DateTime<Utc>) -> i128 {
    dt.timestamp() as i128 * NANOS_PER_SECOND + dt.timestamp_subsec_nanos() as i128
}

fn ns_from_naivedt(dt: NaiveDateTime) -> i128 {
    ns_from_datetime(dt.and_utc())
}

fn ns_from_naivedate(d: NaiveDate) -> i128 {
    let dt = d.and_hms_opt(0, 0, 0).unwrap();
    ns_from_naivedt(dt)
}

fn unit_literal_to_ns(num: i64, unit: &str) -> Option<i128> {
    let nanos_per_unit: i128 = match unit {
        "ns" => 1,
        "us" => 1_000,
        "ms" => 1_000_000,
        "s" => NANOS_PER_SECOND,
        "min" => NANOS_PER_MINUTE,
        "h" => NANOS_PER_HOUR,
        "d" => NANOS_PER_DAY,
        "w" => NANOS_PER_WEEK,
        "y" => NANOS_PER_YEAR,
        _ => return None,
    };
    (num as i128).checked_mul(nanos_per_unit)
}

fn parse_rule(input: Input<'_>) -> IResult<'_, Item> {
    // optional prefix modality
    let (input, prefix_modality) = opt(parse_modal_annotation).parse(input)?;
    let (input, head) = parse_predicate(input)?;
    let (input, temporal_interval) = opt(parse_temporal_annotation).parse(input)?;
    let (input, _) = tok(Token::Neck)(input)?;
    let (input, body) = parse_body(input)?;
    let (input, trailing_modality) = opt(parse_modal_annotation).parse(input)?;
    let (input, _) = tok(Token::Dot)(input)?;
    let modality = prefix_modality.or(trailing_modality);
    Ok((
        input,
        Item::Rule(Rule {
            head,
            body,
            temporal_interval,
            modality,
        }),
    ))
}

fn parse_query(input: Input<'_>) -> IResult<'_, Item> {
    let (input, _) = tok(Token::QueryNeck)(input)?;
    let (input, goals) = parse_body(input)?;
    let (input, temporal_interval) = opt(parse_temporal_annotation).parse(input)?;
    let (input, _) = tok(Token::Dot)(input)?;
    Ok((
        input,
        Item::Query(Query {
            goals,
            temporal_interval,
        }),
    ))
}

fn parse_directive(input: Input<'_>) -> IResult<'_, Item> {
    let (input, _) = tok(Token::Neck)(input)?;
    let (input, dir) = parse_directive_body(input)?;
    let (input, _) = tok(Token::Dot)(input)?;
    Ok((input, Item::Directive(dir)))
}

fn parse_probabilistic(input: Input<'_>) -> IResult<'_, Item> {
    // probability :: clause
    let (input, prob): (Input<'_>, f64) =
        alt((map(float_tok, |f| f), map(int_tok, |n| n as f64))).parse(input)?;
    let (input, _) = tok(Token::ColonColon)(input)?;
    // head predicate
    // optional prefix modality (after the probability prefix)
    let (input, prefix_modality) = opt(parse_modal_annotation).parse(input)?;
    let (input, head) = parse_predicate(input)?;
    // check for rule (:-) or fact (.)
    if let Ok((inp, _)) = tok(Token::Neck)(input) {
        let (inp, body) = parse_body(inp)?;
        let (inp, temporal_interval) = opt(parse_temporal_annotation).parse(inp)?;
        let (inp, trailing_modality) = opt(parse_modal_annotation).parse(inp)?;
        let (inp, _) = tok(Token::Dot)(inp)?;
        let modality = prefix_modality.or(trailing_modality);
        let clause = PlainClause::Rule(Rule {
            head,
            body,
            temporal_interval,
            modality,
        });
        Ok((
            inp,
            Item::Probabilistic(Probabilistic {
                probability: prob,
                clause,
            }),
        ))
    } else {
        let (input2, temporal_interval) = opt(parse_temporal_annotation).parse(input)?;
        let (inp, trailing_modality) = opt(parse_modal_annotation).parse(input2)?;
        let (inp, diff_neural_ref) = opt(parse_neural_gradient_annotation).parse(inp)?;
        let (inp, _) = tok(Token::Dot)(inp)?;
        let modality = prefix_modality.or(trailing_modality);
        let clause = PlainClause::Fact(Fact {
            name: head.name,
            args: head.args,
            temporal_interval,
            modality,
            diff_neural_ref,
        });
        Ok((
            inp,
            Item::Probabilistic(Probabilistic {
                probability: prob,
                clause,
            }),
        ))
    }
}

fn parse_item(input: Input<'_>) -> IResult<'_, Item> {
    alt((
        parse_probabilistic,
        parse_query,
        parse_directive,
        parse_rule,
        parse_fact,
    ))
    .parse(input)
}

// ---------------------------------------------------------------------------
// Public entry point
// ---------------------------------------------------------------------------

/// Parse a token slice into a vector of Items.
pub fn parse_program(tokens: &[Token]) -> crate::Result<Vec<Item>> {
    let input = Input(tokens);
    let (remaining, items) = many0(parse_item)
        .parse(input)
        .map_err(|_| Error::UnexpectedToken)?;
    if !remaining.0.is_empty() {
        return Err(Error::UnexpectedToken);
    }
    Ok(items)
}

// ---------------------------------------------------------------------------
// Unit tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;
    use logos::Logos;

    fn lex(src: &str) -> Vec<Token> {
        Token::lexer(src)
            .collect::<Result<Vec<_>, _>>()
            .expect("lex failed")
    }

    fn parse(src: &str) -> Vec<Item> {
        let tokens = lex(src);
        parse_program(&tokens).expect("parse failed")
    }

    fn parse_err(src: &str) -> Error {
        let tokens = Token::lexer(src).collect::<Vec<_>>();
        // If lex produces errors, return LexError
        if tokens.iter().any(|r| r.is_err()) {
            return Error::LexError;
        }
        let tokens: Vec<Token> = tokens.into_iter().map(|r| r.unwrap()).collect();
        parse_program(&tokens).expect_err("expected parse failure")
    }

    #[test]
    fn test_parse_fact() {
        let items = parse("parent(alice, bob).");
        assert_eq!(items.len(), 1);
        assert!(matches!(&items[0], Item::Fact(f)
            if f.name == "parent"
            && f.args == vec![Term::Atom("alice".into()), Term::Atom("bob".into())]
        ));
    }

    #[test]
    fn test_parse_rule() {
        let items = parse("grandparent(X, Y) :- parent(X, Z), parent(Z, Y).");
        assert_eq!(items.len(), 1);
        assert!(matches!(&items[0], Item::Rule(_)));
    }

    #[test]
    fn test_parse_query_single_goal() {
        let items = parse("?- grandparent(alice, Y).");
        assert_eq!(items.len(), 1);
        assert!(matches!(&items[0], Item::Query(_)));
    }

    #[test]
    fn test_parse_query_multi_goal() {
        let items = parse("?- parent(alice, X), parent(X, Y).");
        assert_eq!(items.len(), 1);
        if let Item::Query(q) = &items[0] {
            assert_eq!(q.goals.len(), 2);
        } else {
            panic!("expected Query");
        }
    }

    #[test]
    fn test_parse_term_integer() {
        let items = parse("age(alice, 42).");
        assert!(matches!(&items[0], Item::Fact(f)
            if f.args.contains(&Term::Integer(42))
        ));
    }

    #[test]
    fn test_parse_term_float() {
        let items = parse("score(alice, 3.14).");
        assert!(matches!(&items[0], Item::Fact(f)
            if f.args.iter().any(|t| matches!(t, Term::Float(_)))
        ));
    }

    #[test]
    fn test_parse_term_list() {
        let items = parse("items(alice, [a, b, c]).");
        assert!(matches!(&items[0], Item::Fact(f)
            if f.args.iter().any(|t| matches!(t, Term::List(_)))
        ));
    }

    #[test]
    fn test_parse_term_list_cons() {
        let items = parse("first(H, [H|T]).");
        assert!(matches!(&items[0], Item::Fact(f)
            if f.args.iter().any(|t| matches!(t, Term::ListCons(_, _)))
        ));
    }

    #[test]
    fn test_parse_compound() {
        let items = parse("event(login, alice, t(2023, 3, 14)).");
        assert!(matches!(&items[0], Item::Fact(f)
            if f.args.iter().any(|t| matches!(t, Term::Compound(n, _) if n == "t"))
        ));
    }

    #[test]
    fn test_parse_expr_arithmetic() {
        let items = parse("double(X, Y) :- Y is X * 2.");
        assert!(matches!(&items[0], Item::Rule(r)
            if r.body.iter().any(|g| matches!(g,
                BodyGoal::Expr(e) if matches!(e.as_ref(), Expr::BinOp(BinOpKind::Is, _, _))
            ))
        ));
    }

    #[test]
    fn test_parse_expr_comparison() {
        let items = parse("positive(X) :- X > 0.");
        assert!(matches!(&items[0], Item::Rule(r)
            if r.body.iter().any(|g| matches!(g,
                BodyGoal::Expr(e) if matches!(e.as_ref(), Expr::BinOp(BinOpKind::Gt, _, _))
            ))
        ));
    }

    #[test]
    fn test_parse_body_not() {
        let items = parse("not_parent(X, Y) :- not parent(X, Y).");
        assert!(matches!(&items[0], Item::Rule(r)
            if r.body.iter().any(|g| matches!(g, BodyGoal::Not(_)))
        ));
    }

    #[test]
    fn test_parse_probabilistic_fact() {
        let items = parse("0.8 :: rain.");
        assert!(matches!(&items[0], Item::Probabilistic(p)
            if (p.probability - 0.8).abs() < f64::EPSILON
        ));
    }

    #[test]
    fn test_parse_directive_table() {
        let items = parse(":- table ancestor/2.");
        assert!(matches!(
            &items[0],
            Item::Directive(Directive::Table { .. })
        ));
    }

    #[test]
    fn test_parse_directive_optimize() {
        let items = parse(":- optimize minimize loss(classification).");
        assert!(matches!(
            &items[0],
            Item::Directive(Directive::Optimize {
                objective: Objective::Minimize,
                ..
            })
        ));
    }

    #[test]
    fn test_error_missing_paren() {
        parse_err("parent(alice, bob");
        // just checks it errors
    }

    #[test]
    fn test_error_missing_dot() {
        parse_err("parent(alice, bob)");
    }

    #[test]
    fn test_parse_diff_neural_fact() {
        let items = parse(r#"predict(X) :: "mlp" :: "loss"."#);
        match &items[0] {
            Item::Fact(f) => {
                assert_eq!(f.name, "predict");
                assert_eq!(f.args.len(), 1);
                assert_eq!(
                    f.diff_neural_ref,
                    Some(("mlp".to_string(), "loss".to_string()))
                );
            }
            _ => panic!("Expected Fact, got: {:?}", items),
        }
    }

    #[test]
    fn test_parse_diff_neural_directive() {
        let items = parse(r#":- differentiable neural classify/0 using "cnn"."#);
        match &items[0] {
            Item::Directive(Directive::DiffNeural {
                functor,
                arity,
                model_name,
            }) => {
                assert_eq!(functor, "classify");
                assert_eq!(*arity, 0);
                assert_eq!(model_name, "cnn");
            }
            _ => panic!("Expected Directive::DiffNeural, got: {:?}", items),
        }
    }

    #[test]
    fn test_parse_type_and_neural_gradient() {
        let items = parse(r#"predict(X:foo :: "mlp" :: "loss")."#);
        match &items[0] {
            Item::Fact(f) => {
                assert_eq!(f.name, "predict");
                assert_eq!(f.args.len(), 1);
                match &f.args[0] {
                    Term::NeuralGradient {
                        term,
                        model_id,
                        grad_id,
                    } => {
                        assert_eq!(model_id, "mlp");
                        assert_eq!(grad_id, "loss");
                        match term.as_ref() {
                            Term::TypeAnnotated {
                                term: inner,
                                type_name,
                            } => {
                                assert_eq!(type_name, "foo");
                                assert!(matches!(inner.as_ref(), Term::Variable(_)));
                            }
                            _ => panic!("Expected TypeAnnotated, got: {:?}", term),
                        }
                    }
                    _ => panic!("Expected NeuralGradient, got: {:?}", f.args[0]),
                }
            }
            _ => panic!("Expected Fact, got: {:?}", items),
        }
    }
}