/*

mondo.mjs - <short description TODO>

Copyright (C) 2022 Strudel contributors - see <https://github.com/tidalcycles/strudel/blob/main/packages/mini/test/mini.test.mjs>

This program is free software: you can redistribute it and/or modify it under the terms of the GNU Affero General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Affero General Public License for more details. You should have received a copy of the GNU Affero General Public License along with this program.  If not, see <https://www.gnu.org/licenses/>.

*/

// evolved from https://garten.salat.dev/lisp/parser.html

export class MondoParser {

  // these are the tokens we expect

  token_types = {

    comment: /^\/\/(.*?)(?=\n|$)/,

    quotes_double: /^"(.*?)"/,

    quotes_single: /^'(.*?)'/,

    open_list: /^\(/,

    close_list: /^\)/,

    open_angle: /^</,

    close_angle: /^>/,

    open_square: /^\[/,

    close_square: /^\]/,

    open_curly: /^\{/,

    close_curly: /^\}/,

    number: /^-?[0-9]*\.?[0-9]+/, // before pipe!

    // TODO: better error handling when "-" is used as rest, e.g "s [- bd]"

    op: /^[*/:!@%?+\-&]|^\.{2}/, // * / : ! @ % ? ..

    // dollar: /^\$/,

    pipe: /^#/,

    stack: /^[,$]/,

    or: /^[|]/,

    plain: /^[a-zA-Z0-9-~_^#]+/,

  };

  op_precedence = [['*', '/', ':', '!', '@', '%', '?', '+', '-', '..'], ['&']];

  // matches next token

  next_token(code, offset = 0) {

    for (let type in this.token_types) {

      const match = code.match(this.token_types[type]);

      if (match) {

        let token = { type, value: match[0] };

        if (offset !== -1) {

          // add location

          token.loc = [offset, offset + match[0].length];

        }

        return token;

      }

    }

    throw new Error(`mondo: could not match '${code}'`);

  }

  // takes code string, returns list of matched tokens (if valid)

  tokenize(code, offset = 0) {

    let tokens = [];

    let locEnabled = offset !== -1;

    let trim = () => {

      // trim whitespace at start, update offset

      offset += code.length - code.trimStart().length;

      // trim start and end to not confuse parser

      return code.trim();

    };

    code = trim();

    while (code.length > 0) {

      code = trim();

      const token = this.next_token(code, locEnabled ? offset : -1);

      code = code.slice(token.value.length);

      offset += token.value.length;

      tokens.push(token);

    }

    return tokens;

  }

  // take code, return abstract syntax tree

  parse(code, offset) {

    this.code = code;

    this.offset = offset;

    this.tokens = this.tokenize(code, offset);

    const expressions = [];

    while (this.tokens.length) {

      expressions.push(this.parse_expr());

    }

    if (expressions.length === 0) {

      // empty case

      return { type: 'list', children: [] };

    }

    // do we have multiple top level expressions or a single non list?

    if (expressions.length > 1 || expressions[0].type !== 'list') {

      return {

        type: 'list',

        children: this.desugar(expressions),

      };

    }

    // we have a single list

    return expressions[0];

  }

  // parses any valid expression

  parse_expr() {

    if (!this.tokens[0]) {

      throw new Error(`unexpected end of file`);

      // TODO: could we allow that? like (((((((( s bd

      // return { type: 'list', children: [] };

    }

    let next = this.tokens[0]?.type;

    if (next === 'open_list') {

      return this.parse_list();

    }

    if (next === 'open_angle') {

      return this.parse_angle();

    }

    if (next === 'open_square') {

      return this.parse_square();

    }

    if (next === 'open_curly') {

      return this.parse_curly();

    }

    return this.consume(next);

  }

  // Token[] => Token[][], e.g. (x , y z) => [['x'],['y','z']]

  split_children(children, split_type) {

    const chunks = [];

    while (true) {

      let splitIndex = children.findIndex((child) => child.type === split_type);

      if (splitIndex === -1) break;

      const chunk = children.slice(0, splitIndex);

      chunks.push(chunk);

      children = children.slice(splitIndex + 1);

    }

    chunks.push(children);

    return chunks;

  }

  desugar_split(children, split_type, next) {

    const chunks = this.split_children(children, split_type);

    if (chunks.length === 1) {

      return next(children);

    }

    // collect args of stack function

    const args = chunks

      .map((chunk) => {

        if (!chunk.length) {

          return; // useful for things like "$ s bd $ s hh*8" (first chunk is empty)

        }

        if (chunk.length === 1) {

          // chunks of one element can be added to the stack as is

          return chunk[0];

        }

        // chunks of multiple args

        chunk = next(chunk);

        return { type: 'list', children: chunk };

      })

      .filter(Boolean); // ignore empty chunks

    return [{ type: 'plain', value: split_type }, ...args];

  }

  // prevents to get a list, e.g. ((x y)) => (x y)

  unwrap_children(children) {

    if (children.length === 1) {

      return children[0].children;

    }

    return children;

  }

  desugar_ops(children, types) {

    while (true) {

      let opIndex = children.findIndex((child) => child.type === 'op' && types.includes(child.value));

      if (opIndex === -1) break;

      const op = { type: 'plain', value: children[opIndex].value };

      if (opIndex === children.length - 1) {

        //throw new Error(`cannot use operator as last child.`);

        children[opIndex] = op; // ignore operator if last child.. e.g. "note [c -]"

        continue;

      }

      if (opIndex === 0) {

        // regular function call (assuming each operator exists as function)

        children[opIndex] = op;

        continue;

      }

      // convert infix to prefix notation

      const left = children[opIndex - 1];

      const right = children[opIndex + 1];

      if (left.type === 'pipe') {

        // "x !* 2" => (* 2 x)

        children[opIndex] = op;

        continue;

      }

      // some careful error handling

      if (left.type === 'op') {

        throw new Error(`got 2 ops in a row: "${left.value}${op.value}"`);

      }

      if (right.type === 'op') {

        let err = `got 2 ops in a row: "${op.value}${right.value}"`;

        if (op.value === '-') {

          // yes i know this file is not supposed to know about rests x.X

          err += '. you probably want a rest, which is "_" in mondo!';

        }

        throw new Error(err);

      }

      const call = { type: 'list', children: [op, right, left] };

      // insert call while keeping other siblings

      children = [...children.slice(0, opIndex - 1), call, ...children.slice(opIndex + 2)];

      children = this.unwrap_children(children);

    }

    return children;

  }

  get_lambda(args, children) {

    // (.fast 2) = (fn (_) (fast _ 2))

    children = this.desugar(children);

    const body = children.length === 1 ? children[0] : { type: 'list', children };

    return [{ type: 'plain', value: 'fn' }, { type: 'list', children: args }, body];

  }

  // returns location range of given ast (even if desugared)

  get_range(ast, range = [Infinity, 0]) {

    let union = (a, b) => [Math.min(a[0], b[0]), Math.max(a[1], b[1])];

    if (ast.loc) {

      return union(range, ast.loc);

    }

    if (ast.type !== 'list') {

      return range;

    }

    return ast.children.reduce((range, child) => {

      const childrange = this.get_range(child, range);

      return union(range, childrange);

    }, range);

  }

  errorhead(ast) {

    return `[mondo ${this.get_range(ast)?.join(':') || '?'}]`;

  }

  // returns original user code where the given ast originates (even if desugared)

  get_code_snippet(ast) {

    const [min, max] = this.get_range(ast);

    return this.code.slice(min - this.offset, max - this.offset);

  }

  desugar_pipes(children) {

    let chunks = this.split_children(children, 'pipe');

    while (chunks.length > 1) {

      let [left, right, ...rest] = chunks;

      if (!left.length) {

        const arg = { type: 'plain', value: '_' };

        return this.get_lambda([arg], [arg, ...children]);

      }

      // s jazz hh.fast 2 => (fast 2 (s jazz hh))

      const call = left.length > 1 ? { type: 'list', children: left } : left[0];

      chunks = [[...right, call], ...rest];

    }

    // return next(chunks[0]);

    return chunks[0];

  }

  parse_pair(open_type, close_type) {

    const begin = this.tokens[0].loc?.[0];

    this.consume(open_type);

    const children = [];

    while (this.tokens[0]?.type !== close_type) {

      children.push(this.parse_expr());

    }

    const end = this.tokens[0].loc?.[1];

    this.consume(close_type);

    const node = { type: 'list', children };

    if (begin !== undefined) {

      node.loc = [begin, end];

      node.raw = this.code.slice(begin, end);

    }

    return node;

  }

  desugar(children, type) {

    // if type is given, the first element is expected to contain it as plain value

    // e.g. with (square a b, c), we want to split (a b, c) and ignore "square"

    children = type ? children.slice(1) : children;

    children = this.desugar_split(children, 'stack', (children) =>

      this.desugar_split(children, 'or', (children) => {

        // chunks of multiple args

        if (type) {

          // the type we've removed before splitting needs to be added back

          children = [{ type: 'plain', value: type }, ...children];

        }

        // for each precendence group, call desugar_ops once

        this.op_precedence.forEach((ops) => {

          children = this.desugar_ops(children, ops);

        });

        children = this.desugar_pipes(children);

        return children;

      }),

    );

    return children;

  }

  parse_list() {

    let node = this.parse_pair('open_list', 'close_list');

    node.children = this.desugar(node.children);

    return node;

  }

  parse_angle() {

    let node = this.parse_pair('open_angle', 'close_angle');

    node.children.unshift({ type: 'plain', value: 'angle' });

    node.children = this.desugar(node.children, 'angle');

    return node;

  }

  parse_square() {

    let node = this.parse_pair('open_square', 'close_square');

    node.children.unshift({ type: 'plain', value: 'square' });

    node.children = this.desugar(node.children, 'square');

    return node;

  }

  parse_curly() {

    let node = this.parse_pair('open_curly', 'close_curly');

    node.children.unshift({ type: 'plain', value: 'curly' });

    node.children = this.desugar(node.children, 'curly');

    return node;

  }

  consume(type) {

    // shift removes first element and returns it

    const token = this.tokens.shift();

    if (token.type !== type) {

      throw new Error(`expected token type ${type}, got ${token.type}`);

    }

    return token;

  }

  get_locations(code, offset = 0) {

    let walk = (ast, locations = []) => {

      if (ast.type === 'list') {

        return ast.children.forEach((child) => walk(child, locations));

      }

      if (ast.loc) {

        locations.push(ast.loc);

      }

    };

    const ast = this.parse(code, offset);

    let locations = [];

    walk(ast, locations);

    return locations;

  }

}

export function printAst(ast, compact = false, lvl = 0) {

  const br = compact ? '' : '\n';

  const spaces = compact ? '' : Array(lvl).fill(' ').join('');

  if (ast.type === 'list') {

    return `${lvl ? br : ''}${spaces}(${ast.children.map((child) => printAst(child, compact, lvl + 1)).join(' ')}${

      ast.children.find((child) => child.type === 'list') ? `${br}${spaces})` : ')'

    }`;

  }

  return `${ast.value}`;

}

// lisp runner

export class MondoRunner {

  constructor({ evaluator } = {}) {

    this.parser = new MondoParser();

    this.evaluator = evaluator;

    this.assert(typeof evaluator === 'function', `expected an evaluator function to be passed to new MondoRunner`);

  }

  // a helper to check conditions and throw if they are not met

  assert(condition, error) {

    if (!condition) {

      throw new Error(error);

    }

  }

  run(code, scope, offset = 0) {

    const ast = this.parser.parse(code, offset);

    //console.log(printAst(ast));

    return this.evaluate(ast, scope);

  }

  evaluate_let(ast, scope) {

    // (let ((x 3) (y 4)) ...body)

    // = ((fn (x y) ...body) 3 4)

    const defs = ast.children[1].children;

    const args = defs.map((pair) => pair.children[0]);

    const vals = defs.map((pair) => pair.children[1]);

    const body = ast.children.slice(2);

    const lambda = {

      type: 'list',

      children: [{ type: 'plain', value: 'fn' }, { type: 'list', children: args }, ...body],

    };

    return this.evaluate({ type: 'list', children: [lambda, ...vals] }, scope);

  }

  evaluate_def(ast, scope) {

    // function definition special form?

    if (ast.children[1].type === 'list') {

      // (def (add a b) (+ a b))

      // => (def add (fn (a b) (+ a b)) )

      const args = ast.children[1].children.slice(1);

      const lambda = {

        // lambda

        type: 'list',

        children: [

          { type: 'plain', value: 'fn' },

          { type: 'list', children: args },

          ...ast.children.slice(2), // body

        ],

      };

      // we mutate to make sure the old ast wont make a mess later

      ast.children[1] = ast.children[1].children[0];

      ast.children[2] = lambda;

      ast.children = ast.children.slice(0, 3); // throw away rest

    }

    // (def name body)

    if (ast.children.length !== 3) {

      throw new Error(`expected "def" to have 3 children, but got ${ast.children.length}`);

    }

    const name = ast.children[1].value;

    const body = this.evaluate(ast.children[2], scope);

    scope[name] = body;

    // def with fall through

  }

  evaluate_match(ast, scope) {

    // (match (p1 e1) (p2 e2) ... (pn en))

    // = cond in lisp

    if (ast.children.length < 2) {

      return;

    }

    const [_, ...body] = ast.children;

    for (let i = 0; i < body.length; ++i) {

      const [predicate, exp] = body[i].children;

      if (predicate.value === 'else') {

        return this.evaluate(exp, scope);

      }

      const outcome = this.evaluate(predicate, scope);

      if (outcome) {

        return this.evaluate(exp, scope);

      }

    }

    return undefined; // nothing was matched

  }

  evaluate_if(ast, scope) {

    // if is a special case of match

    if (ast.children.length !== 4) {

      return;

    }

    // (if predicate consequent alternative)

    const [_, predicate, consequent, alternative] = ast.children;

    // (match (predicate consequent) (else alternative))

    const matcher = {

      type: 'list',

      children: [

        { type: 'plain', value: 'match' },

        { type: 'list', children: [predicate, consequent] },

        { type: 'list', children: [{ type: 'plain', value: 'else' }, alternative] },

      ],

    };

    return this.evaluate_match(matcher, scope);

  }

  evaluate_lambda(ast, scope) {

    // (fn (_)   (ply 2 _)

    //     ^args ^ body

    const [_, formalArgs, ...body] = ast.children;

    return (...args) => {

      const params = Object.fromEntries(formalArgs.children.map((arg, i) => [arg.value, args[i]]));

      const closure = {

        ...scope,

        ...params,

      };

      // body can have multiple expressions

      const res = body.map((exp) => this.evaluate(exp, closure));

      // last expression is the return value

      return res[res.length - 1];

    };

  }

  evaluate_list(ast, scope) {

    // evaluate all children before evaluating list (dont mutate!!!)

    const args = ast.children

      .filter((child) => child.type !== 'comment') // ignore comments

      .map((arg) => this.evaluate(arg, scope));

    const node = { type: 'list', children: args };

    return this.evaluator(node, scope);

  }

  evaluate_leaf(ast, scope) {

    if (ast.type === 'number') {

      ast.value = Number(ast.value);

    } else if (['quotes_double', 'quotes_single'].includes(ast.type)) {

      ast.value = ast.value.slice(1, -1);

      ast.type = 'string';

    }

    return this.evaluator(ast, scope);

  }

  evaluate(ast, scope = {}) {

    if (ast.type !== 'list') {

      return this.evaluate_leaf(ast, scope);

    }

    const name = ast.children[0]?.value;

    if (name === 'fn') {

      return this.evaluate_lambda(ast, scope);

    }

    if (name === 'match') {

      return this.evaluate_match(ast, scope);

    }

    if (name === 'if') {

      return this.evaluate_if(ast, scope);

    }

    if (name === 'let') {

      return this.evaluate_let(ast, scope);

    }

    if (name === 'def') {

      this.evaluate_def(ast, scope);

    }

    return this.evaluate_list(ast, scope);

  }

}