#!/usr/bin/env python3
# The CAVE Machine

import re
import sys

Symbol = str
sym = sys.intern

class Pair(object):
   def __init__(self, head, tail):
      self.head = head
      self.tail = tail

   def __repr__(self):
      return "Pair(%r, %r)" % (self.head, self.tail)

   def __str__(self):
      return "(%s . %s)" % (self.head, self.tail)

pair = Pair

TRUE = sym("true")
FALSE = sym("false")
EMPTY = sym("()")
NULL = sym("null")
NULLARY = sym("(())")

def auto(sy):
   return sy.startswith(":")

auto_extra = {TRUE, FALSE, EMPTY, NULL}

def stail(obj):
   return isinstance(obj, Pair) or (obj is EMPTY)

# Array to SLL
def a2s(ar):
   si = EMPTY
   pos = len(ar) - 1
   while pos >= 0:
      si = pair(ar[pos], si)
      pos = pos - 1
   return si

# SLL to Array
def s2a(si):
   ar = []
   while si is not EMPTY:
      ar.append(si.head)
      if not stail(si.tail):
         raise ValueError("Input is not a SLL")
      si = si.tail
   return ar

def last(si):
   while si is not EMPTY:
      if not stail(si.tail):
         raise ValueError("Input is not a SLL")
      if si.tail is EMPTY:
         break
      si = si.tail
   return si.head

# Pop an object from a SLL
# obj, si = pop(si)
def pop(si):
   return si.head, si.tail

def nth(si, n):
   while si is not EMPTY:
      if n == 0:
         return si.head
      if not stail(si.tail):
         raise ValueError("Input is not a SLL")
      si = si.tail
      n = n - 1
   raise ValueError("Out of bounds access of a SLL")

def env_mappings(env):
   return s2a(env)

def env_get(env, sy):
   for bindings in s2a(env):
      if sy in bindings:
         return bindings[sy]
   raise KeyError("Symbol %r not found in environment" % sy)

def env_set_local(env, sy, obj):
   env.head[sy] = obj
   return env

def env_set_global(env, sy, obj):
   top = env
   while env.tail is not EMPTY:
      env = env.tail
   env.head[sy] = obj
   return top

r_token = re.compile(r'[()]|"[^"]*"|[0-9]+|[^ "()]+|[(][)]')
paren_token = lambda token: token in "()"
string_token = lambda token: token.startswith('"')
number_token = lambda token: token[0] in "0123456789"

def parse(text):
   stack = [[]]
   for token in r_token.findall(text):
      if paren_token(token):
         if token == "(":
            stack.append([])
         elif token == ")":
            value = a2s(stack.pop())
            stack[-1].append(value)

      elif string_token(token):
         value = token[1:-1].encode("utf-8")
         stack[-1].append(value)

      elif number_token(token):
         value = int(token)
         stack[-1].append(value)

      else: # symbol
         value = sym(token)
         stack[-1].append(value)

   return stack[0][0]

def serialise(expr):
   if isinstance(expr, int):
      return str(expr)

   elif isinstance(expr, str):
      return expr

   elif isinstance(expr, bytes):
      return '"' + expr.decode("utf-8") + '"'

   elif isinstance(expr, Pair):
      children = []
      try: c = s2a(expr)
      except ValueError:
         h = serialise(expr.head)
         t = serialise(expr.tail)
         return "(: " + h + " " + t + ")"

      for child in c:
         children.append(serialise(child))
      return "(" + " ".join(children) + ")"

   elif isinstance(expr, dict):
      kvs = []
      for k, v in sorted(expr.items()):
         kvs.append(serialise(k))
         kvs.append(serialise(v))
      return "(:: " + " ".join(kvs) + ")"

   # expr could be a tuple, so don't use %s
   raise ValueError("Cannot be serialised: " + str(expr))

VALUE = sym("VALUE")
EVAL = sym("EVAL")

basic = {}
primitives = {}
closures = {}

def argspec_instructions(argspec):
   # print("ARGSPEC:", argspec)
   aas = s2a(argspec)

   variadic = FALSE

   if aas:
      end = aas[-1]
      if end.startswith("*"):
         end = sym(end[1:])
         if end.startswith("'"):
            ins = VALUE
            end = sym(end[1:])
         else:
            ins = EVAL
         variadic = pair(ins, end)
         aas = aas[:-1]

   strategy = []
   symbols = []

   for aa in aas:
      if aa.startswith("'"):
         ins = VALUE
         aa = sym(aa[1:])
      else:
         ins = EVAL
      strategy.append(ins)
      symbols.append(aa)

   return a2s(strategy), a2s(symbols), variadic

def primitive(name, argspec):
   sy = sym(name)
   argspec = parse(argspec)

   if name.startswith(":"):
      asy = sym(name)
   else:
      asy = sym(":" + name)

   def decorate(function):
      basic[sy] = asy
      primitives[asy] = function
      strategy, named, variadic = argspec_instructions(argspec)
      function.strategy = strategy
      function.variadic = variadic
   return decorate

# python3 -c 'import cave; print(cave.tests.examples)'

def tests(text):
   tests.examples += text
   lines = text.strip("\n").split("\n")
   while lines:
      test, expected = lines[:2]
      lines = lines[2:]
      test = parse(test[2:])
      tests.tests.append((test, expected))
tests.tests = []
tests.examples = ""

# PURE, CAVE, STATE, I/O

### Primitives

# *

@primitive("*", "(a b)")
def multiply(CAVE, a, b):
   assert isinstance(a, int)
   assert isinstance(b, int)
   return a * b

tests("""
> (* 1 2)
2
> (* 1 2 3)
ERROR
> (type (* 1 2))
Number
> ((* 1) 5)
5
> (* (id 1) (id 2))
2
""")

# +

@primitive("+", "(a b)")
def add(CAVE, a, b):
   assert isinstance(a, int)
   assert isinstance(b, int)
   return a + b

tests("""
> (+ 1 2)
3
> (+ 1 2 3)
ERROR
> (type (+ 1 2))
Number
> ((+ 1) 5)
6
> (+ (id 1) (id 2))
3
""")

# -
@primitive("-", "(a b)")
def subtract(CAVE, a, b):
   assert isinstance(a, int)
   assert isinstance(b, int)
   return a - b

# TODO: (- 1 2)
tests("""
> (- 2 1)
1
> (- 3 2 1)
ERROR
> (type (- 2 1))
Number
> ((- 2) 1)
1
> (- (id 2) (id 1))
1
""")

# <

@primitive("<", "(a b)")
def lt(CAVE, a, b):
   assert isinstance(a, int)
   assert isinstance(b, int)
   if a < b:
      return TRUE
   return FALSE

tests("""
> (< 1 2)
true
> (< 2 1)
false
> (< (id 1) (id 2))
true
> (< (id 2) (id 1))
false
> (< "a" "b")
ERROR
> (type (< 2 1))
Symbol
> (type (< 2))
Pair
""")

# :

@primitive(":", "('a 'b)")
def create_pair(CAVE, a, b):
   return pair(a, b)

# TODO: reverse curry, swap, with (: ())
tests("""
> (: a b)
(: a b)
> (type (: a b))
Pair
> (: (id 1) (id 2))
((id 1) id 2)
> (curry? (head (: (id) (id 2))))
false
> (closure? (head (: (id 1) (id 2))))
false
""")

# (= swap (-> (a 'b 'c) (a c b)))
# (= swap (-> (a 'b 'c) (a (eval c) (eval b))))
# (swap : x y)

# (= swap (-> (a 'b 'c) (a c b)))
# (= .x (swap : x))
# (.x (' y))

# =

@primitive("=", "('a b)")
def assign(CAVE, a, b):
   if auto(a) or (a in auto_extra):
      raise Exception("Cannot assign to primitive self-evaluating symbol")
   env_set_global(CAVE["E"], a, b)
   return b

tests("""
> (= ten 10)
10
> ten
10
> (del ten)
ten
> (def? ten)
false
> (= a b)
ERROR
> (= x 3)
3
> (= y 5)
5
> (+ x y)
8
> (del x)
x
> (del y)
y
""")

# =v

@primitive("=v", "(a b)")
def assignv(CAVE, a, b):
   if auto(a) or (a in auto_extra):
      raise Exception("Cannot assign to primitive self-evaluating symbol")
   env_set_global(CAVE["E"], a, b)
   return b

tests("""
> (=v (' five) 5)
5
> five
5
> (del five)
five
""")

# closure?

@primitive("closure?", "(a)")
def closure_p(CAVE, a):
   if a in closures:
      return TRUE
   return FALSE

tests("""
> (closure? (-> (a) a))
true
> (closure? (' (-> (a) a)))
false
""")

# curry?

@primitive("curry?", "(a)")
def curry_p(CAVE, a):
   if a in curries:
      return TRUE
   return FALSE

tests("""
> (curry? (+))
true
> (curry? (+ 1))
true
> (curry? ((+) 1))
true
> (curry? (((+) 1) 2))
false
> (curry? ((+ 1) 2))
false
> (curry? (+ 1 2))
false
""")

# continuation

@primitive("continuation", "()")
def continuation(CAVE):
   return CAVE["C"]

tests("""
> (continuation)
(ARGS () STOP)
> (id (continuation))
(ARGS () ARG APPLY CLOSURE (-> (obj) obj) ARGS () STOP)
""")

# =>

@primitive("=>", "('sy 'argspec *'exprs)")
def _fatarrow(CAVE, sy, argspec, *exprs):
   exprs = exprs[0]
   body = a2s([sym(":do")] + s2a(exprs))
   print("ARGSPEC-INNER:", argspec)
   clo = a2s([sym(":->"), argspec, body])
   e = a2s([sym(":eval"), clo])
   eq = a2s([sym(":="), sy, e])
   CAVE["C"] = push([EVAL, eq], CAVE["C"])
   return NULL # discarded

tests_ = """
> (=> . (a) (+ 1 a))
(-> (a) (:do (+ 1 a)))
> (. 5)
6
"""

# def

@primitive("def", "('sy 'argspec 'exprs)")
def _def(CAVE, sy, argspec, exprs):
   body = a2s([sym(":do")] + s2a(exprs))
   clo = a2s([sym(":->"), argspec, body])
   e = a2s([sym(":eval"), clo])
   eq = a2s([sym(":="), sy, e])
   CAVE["C"] = push([EVAL, eq], CAVE["C"])
   return NULL # discarded

tests("""
> (= . (def . (a)))
(:def . (a))
> (. ((+ 1 a))) 
(-> (a) (:do (+ 1 a)))
> (. 5)
6
""")

# def?

@primitive("def?", "('a)")
def defined(CAVE, a):
   e = CAVE["E"]
   while e is not EMPTY:
      if a in e.head:
         return TRUE
      e = e.tail
   return FALSE

tests("""
> (def? ->)
true
> (def? def?)
true
> (def? def!)
false
""")

# del

@primitive("del", "('a)")
def delete(CAVE, a):
   del CAVE["E"].head[a]
   return a

tests("""
> (= a "a")
"a"
> (def? a)
true
> (del a)
a
> (def? a)
false
""")

# ::

@primitive("::", "(*'kvs)")
def _vdict(CAVE, *kvs):
   if not kvs:
      return {}
   kvs = kvs[0]

   result = {}
   while kvs is not EMPTY:
      key = kvs.head
      assert isinstance(key, Symbol)
      value = kvs.tail.head
      kvs = kvs.tail.tail
      result[key] = value

   return result

# dict

@primitive("dict", "(*kvs)")
def _dict(CAVE, *kvs):
   if not kvs:
      return {}
   kvs = kvs[0]

   result = {}
   while kvs is not EMPTY:
      key = kvs.head
      assert isinstance(key, Symbol)
      value = kvs.tail.head
      kvs = kvs.tail.tail
      result[key] = value

   return result

example = """
names = (dict "a" 1 "b" 2)

names =
  dict "a" 1 "b" 2

names =
  dict
    "a"
    1
    "b"
    2

names =
  dict \
    "a" 1 \
    "b" 2

names =
  dict "a" 1 \
       "b" 2

names = dict \
  "a" 1 \
  "b" 2

names =
  dict-pairs
    "a" : 1
    "b" : 2

names ::=
  "a" : 1
  "b" : 2

# @@ evaluated pairs

names ::=
  pair "a" 1
  pair "b" 2

# ::= could call eval...
"""

# do

# (do ((pro 1) (continuation)))
# (do ((pro (continuation)) (pro 1)))
# (' (APPLY PRIMITIVE :pro RETURN ((pro 1)) APPLY PRIMITIVE :do STOP))
# (' (APPLY PRIMITIVE :pro RETURN ((pro 1)) APPLY PRIMITIVE :do STOP))
@primitive("do", "(*args)")
def do(CAVE, args):
   return last(args)

# TODO: ltr evaluation order
tests("""
> (do (id 1) (id 2))
2
> (do (= one 1) (+ one 2))
3
> (del one)
one
""")

# eval

@primitive("eval", "(a)")
def eval(CAVE, a):
   CAVE["C"] = push([EVAL, a], CAVE["C"])
   return NULL # discarded

tests("""
> (eval (' (id "a")))
"a"
""")

# head

@primitive("head", "(a)")
def head(CAVE, a):
   assert isinstance(a, Pair)
   return a.head

tests("""
> (head (' (a b c)))
a
> (head (: car cdr))
car
> (head ())
ERROR
""")

# if

@primitive("if", "(a 'b 'c)")
def conditional(CAVE, a, b, c):
   if a is TRUE:
      C_push(CAVE, [EVAL, b])
   elif a is FALSE:
      C_push(CAVE, [EVAL, c])
   else:
      raise ValueError("Expected bool")
   return NULL # discarded

tests("""
> (if true "yep" "nope")
"yep"
> (if false "yep" "nope")
"nope"
> (if (def? if) "yep" "nope")
"yep"
""")

def split1(si, sy):
   ar = s2a(si)
   pos = ar.index(sy)
   return a2s(ar[:pos]), a2s(ar[pos:])

# inner

k_counter = 0

@primitive("inner", "('name 'obj)")
def inner(CAVE, name, obj):
   global k_counter

   # Set name to the continuation
   # Get the continuation
   dcont, C = split1(CAVE["C"], sym("OUTER"))
   CAVE["C"] = C
   # print("DELCONT:", serialise(dcont))

   k_counter += 1
   pn = "c" + str(k_counter)
   @primitive(pn, "(obj)")
   def delcont(CAVE, obj):
      C_push(CAVE, [VALUE, obj] + s2a(dcont))
      return NULL # discarded
   env_set_global(CAVE["E"], name, sym(":" + pn))

   # Then eval obj
   C_push(CAVE, [EVAL, obj])
   return NULL # discarded

# loop

@primitive("loop", "('obj)")
def loop(CAVE, obj):
   C_push(CAVE, [EVAL, obj,
                 ARGS, EMPTY,
                 VALUE, obj, ARG,
                 APPLY, sym("PRIMITIVE"), sym(":loop")])
   return NULL # discarded

# outer

@primitive("outer", "('obj)")
def outer(CAVE, obj):
   # print("OUTER:", serialise(CAVE["C"]))
   C_push(CAVE, [EVAL, obj, sym("OUTER")])
   return NULL # discarded

tests("""
> (+ 1 (outer (+ 2 3)))
6
> (+ 1 (outer (+ 3 (inner c 5))))
6
> (+ 1 (outer (+ 3 (inner c (+ 5 (c 7))))))
16
> (+ 1 (outer (+ 2 (inner c (sum 3 (c 5) (c 2)))))) 
15
> (pair 1 (outer (pair 3 (inner c (pair 5 (c (c (pair 7 ()))))))))
(1 5 3 3 7)
> (+ 10 (outer (+ 2 (inner c (+ 100 (c (c 3)))))))
117
> (+ 5 (outer (+ 3 (inner c (+ (c 0) (c 1))))))
12
> (+ 1 (outer (+ 2 (+ (inner c 4) 8))))
5
> (outer (+ 1 (outer (+ 2 (outer (+ 4 (inner c1 (inner c2 8))))))))
11
> (outer (+ 1 (outer (+ 2 ((inner c1 (c1 (-> (x) x))) (inner c2 4))))))
5
> ((-> (f) (+ 1 (outer (+ 2 (f 3))))) (-> (x) (inner c x)))
4
""")

extra = """
> (outer (+ 10 (+ (inner c (+ 1 (c 1))) (inner c 2)))) 
2
EXPECTED: 3

((-> (f) (+ 1 (outer (+ 2 (f 3))))) (-> (x) (inner c x)))

(outer (+ 1 (outer (+ 2 ((inner c1 (c1 (-> (x) x))) (inner c2 4))))))

(def yield (x) ((inner c (pair x (c (void))))))

> (def yield (x) (inner c (pair x (c (void)))))
(-> (x) (:do inner c (pair x (c (void)))))
> (outer (do (yield 1) (yield 2) (yield 3) ()))
()
> (outer (do (yield 1)))
(1)
> (outer (do (yield 1) (yield 2)))
(2)
> (outer (do (inner c (pair 1 (c (void)))) (inner c (pair 2 (c (void))))))
(2 1)
> (outer (do (inner c (pair 1 (c (void)))) (inner c (pair 2 (c (void)))) (inner c (pair 3 (c (void)))) ()))
(3 2 1)
> (outer (do (inner c (pair 1 (c (void)))) ()))
(1)
> (outer (do (inner c (pair 1 (c (void)))) (' (2))))
(1 2)
> (outer (do (inner c (pair 1 (c (void)))) (inner c (pair 2 (c (void))))))
(2 1)
> (outer (do (inner c (pair 1 (c (void)))) (inner c (pair 2 (c (void)))) ()))
(2 1)
> (outer (do (pro 1) (inner c (do (c (void)) (c (void)) (pro 2))) (pro 3)))
1
3
3
2
"2"
"""

# pair

@primitive("pair", "(a b)")
def _pair(CAVE, a, b):
   return pair(a, b)

tests("""
> (pair (id 1) (id "b"))
(: 1 "b")
""")

# parse

@primitive("parse", "(a)")
def _parse(CAVE, a):
   return parse(a.decode("utf-8"))

tests("""
> (type (parse "()"))
Symbol
> (type (parse "1"))
Number
> (type (parse "(a)"))
Pair
> (parse "(+ 1 2)")
(+ 1 2)
> (eval (parse "(+ 1 2)"))
3
""")

# print

@primitive("print", "(*args)")
def _print(CAVE, args):
   s = " ".join(serialise(arg) for arg in s2a(args))
   print(s)
   return s.encode("utf-8")

# print-env

@primitive("print-env", "()")
def print_env(CAVE):
   print(CAVE["E"])
   return NULL

# pro

@primitive("pro", "(a)")
def pro(CAVE, a):
   s = serialise(a)
   print(s)
   return s.encode("utf-8")

# read

@primitive("read", "(*prompt)")
def read(CAVE, *prompts):
   if prompts:
      a = input(prompts[0].decode("utf-8"))
   else:
      a = input("$ ")
   return parse(a)

# serialise

@primitive("serialise", "(a)")
def _serialise(CAVE, a):
   return serialise(a).encode("utf-8")

tests("""
> (serialise 1)
"1"
> (serialise (' (a b c)))
"(a b c)"
""")

# strategy

@primitive("strategy", "(op)")
def strategy(CAVE, op):
   if isinstance(op, Symbol):
      return primitives[op].strategy
   elif isinstance(op, Pair):
      return closures[op][1]
   else:
      raise ValueError(str(op))

tests("""
> (strategy +)
(EVAL EVAL)
> (strategy if)
(EVAL VALUE VALUE)
> (strategy print)
()
""")

# sum

@primitive("sum", "(*args)")
def _sum(CAVE, *args):
   if args:
      return sum(s2a(args[0]))
   return 0

tests("""
> (sum)
0
> (sum 1 2 3)
6
""")

# symbols

@primitive("symbols", "()")
def symbols(CAVE):
   syms = set()
   for mapping in env_mappings(CAVE["E"]):
      syms.update(mapping.keys())
   return a2s(sorted(syms))

# tail

@primitive("tail", "(a)")
def tail(CAVE, a):
   assert isinstance(a, Pair)
   return a.tail

# type

@primitive("type", "(a)")
def _type(CAVE, a):
   if isinstance(a, Symbol):
      return sym("Symbol")
   elif isinstance(a, bytes):
      return sym("String")
   elif isinstance(a, int):
      return sym("Number")
   elif isinstance(a, Pair):
      return sym("Pair")
   else:
      return sym("Unknown")

tests("""
> (type (' a))
Symbol
> (type "text")
String
> (type 1)
Number
> (type (: a b))
Pair
> (type type)
Symbol
""")

# void

@primitive("void", "()")
def void(CAVE):
   return EMPTY

tests("""
> (outer (do (inner c (pair 1 (c (void)))) ()))
(1)
""")

extra = """
> (def yield (x) (inner c (pair x (c (void)))))
(-> (x) (:do inner c (pair x (c (void)))))
> (outer (do (yield 1) (yield 2) (yield 3) ()))
(1 2 3)
"""

# ->

@primitive("->", "('argspec 'body)")
def closure(CAVE, argspec, body):
   clo = a2s([sym("->"), argspec, body])
   strategy, named, variadic = argspec_instructions(argspec)
   closures[clo] = CAVE["E"], strategy, named, variadic
   return clo

def push(ar, si):
   while ar:
      si = pair(ar.pop(), si)
   return si

def A_push(CAVE, ar):
   while ar:
      CAVE["A"] = pair(ar.pop(), CAVE["A"])
   return CAVE

def A_pop(CAVE):
   obj = CAVE["A"].head
   CAVE["A"] = CAVE["A"].tail
   return obj

def C_push(CAVE, ar):
   while ar:
      CAVE["C"] = pair(ar.pop(), CAVE["C"])
   return CAVE

def C_pop(CAVE):
   obj = CAVE["C"].head
   CAVE["C"] = CAVE["C"].tail
   return obj

def evaluate_closure(CAVE, clo, args):
   # CAVE = the CAVE stacks
   # clo = a closed closure
   # args = a python array (list) of arguments to the closure
   bindings = {sy: obj for sy, obj in zip(s2a(closures[clo][2]), args)}

   # Update C before E!
   body = clo.tail.tail.head
   CAVE["C"] = push([EVAL, body, ENV, CAVE["E"]], CAVE["C"])
   CAVE["E"] = pair(bindings, closures[clo][0])

instructions = {}

def instruction(sy):
   def decorate(function):
      instructions[sy] = function
   return decorate

EVAL = sym("EVAL")

@instruction(EVAL)
def step_EVAL(CAVE):
   expr = C_pop(CAVE)

   if isinstance(expr, Symbol):
      if auto(expr) or (expr in auto_extra):
         CAVE["V"] = expr
      else:
         CAVE["V"] = env_get(CAVE["E"], expr)

   elif isinstance(expr, Pair):
      op = expr.head
      args = expr.tail
      C_push(CAVE, [EVAL, op, PARGS, args, ARGS, CAVE["A"]])
      CAVE["A"] = EMPTY

   else:
      CAVE["V"] = expr

   return CAVE

curries = {}

PARGS = sym("PARGS")

@instruction(PARGS)
def step_PARGS(CAVE):
   op = CAVE["V"]
   args = C_pop(CAVE)

   if isinstance(op, Symbol):
      kind = sym("PRIMITIVE")
      strategy = primitives[op].strategy
      variadic = primitives[op].variadic

   elif isinstance(op, Pair):
      if op in curries:
         command = a2s(s2a(op) + s2a(args))
         C_push(CAVE, [EVAL, command])
         return CAVE

      kind = sym("CLOSURE")
      strategy = closures[op][1]
      variadic = closures[op][3]

   else:
      raise Exception("Unknown op: %s" % op)

   processed = []
   top = args
   while True:
      if args is EMPTY:
         if strategy is not EMPTY:
            curry = a2s([op] + s2a(top))
            curries[curry] = True
            C_push(CAVE, [VALUE, curry])
            return CAVE
         break
      elif strategy is EMPTY:
         # args is not EMPTY
         if variadic is not FALSE:
            inst = variadic.head
            A_push(CAVE, [EMPTY])
            for arg in s2a(args):
               processed = processed + [inst, arg, VARG]
            break
         else:
            raise Exception("Too many args")

      arg = args.head
      inst = strategy.head

      processed = processed + [inst, arg, ARG]

      args = args.tail
      strategy = strategy.tail

   processed = processed + [APPLY, kind, op]
   C_push(CAVE, processed)

   return CAVE

tests("""
> (= swap (-> (a 'b 'c) (a c b)))
(-> (a 'b 'c) (a c b))
> (= .x (swap pair x))
((-> (a 'b 'c) (a c b)) pair x)
> (.x y)
(: y x)
""")

ARG = sym("ARG")

@instruction(ARG)
def step_ARG(CAVE):
   arg = CAVE["V"]
   A_push(CAVE, [arg])

   return CAVE

VARG = sym("VARG")

@instruction(VARG)
def step_VARG(CAVE):
   arg = CAVE["V"]
   first = A_pop(CAVE)
   first = a2s(s2a(first) + [arg])
   A_push(CAVE, [first])

   return CAVE

APPLY = sym("APPLY")

@instruction(APPLY)
def step_APPLY(CAVE):
   kind = C_pop(CAVE)
   op = C_pop(CAVE)
   args = list(reversed(s2a(CAVE["A"])))

   if kind is sym("PRIMITIVE"):
      CAVE["V"] = primitives[op](CAVE, *args)
   elif kind is sym("CLOSURE"):
      evaluate_closure(CAVE, op, args)
   else:
      raise Exception("Unknown kind: %s" % kind)

   return CAVE

# > (= swap (-> (a 'b c) (a c b)))
# (-> (a 'b c) (a c b))
# > (= .x (swap pair x))
# (-> (c) ((-> (a 'b c) (a c b)) pair x c))
# > (.x (' y))
# (: y x)
# — this should also work with (a 'b 'c) and then just (.x y)
# eval("a" + b) — a form of unquote
# (eval (a (unquote b)))
#
# (-> ('c) ((-> (a 'b 'c) (a c b)) pair x c))
# (-> (b) (:+ 1 b))
# (+ 1) => (:+ 1)
#
# (= swap (-> (a 'b 'c) (a c b)))

ENV = sym("ENV")

@instruction(ENV)
def step_ENV(CAVE):
   env = C_pop(CAVE)
   CAVE["E"] = env

   return CAVE

ARGS = sym("ARGS")

@instruction(ARGS)
def step_ARGS(CAVE):
   args = C_pop(CAVE)
   CAVE["A"] = args

   return CAVE

VALUE = sym("VALUE")

@instruction(VALUE)
def step_VALUE(CAVE):
   obj = C_pop(CAVE)
   CAVE["V"] = obj

   return CAVE

OUTER = sym("OUTER")
@instruction(OUTER)
def step_OUTER(CAVE):
   return CAVE # NOOP

STOP = sym("STOP")

@instruction(STOP)
def step_STOP(CAVE):
   CAVE["C"] = EMPTY

   return CAVE

def step(CAVE):
   C_instruction = C_pop(CAVE)
   function = instructions[C_instruction]
   CAVE = function(CAVE)
   return CAVE

debug = False

def evaluate(env, expr):
   CAVE = {"C": EMPTY, "A": EMPTY, "V": NULL, "E": NULL}
   CAVE["E"] = env
   C_push(CAVE, [EVAL, expr, STOP])

   while True:
      CAVE = step(CAVE)
      if CAVE["C"] is EMPTY:
         return CAVE["E"], CAVE["V"]

      if debug:
         print(serialise(CAVE["C"]))
         print(serialise(CAVE["A"]))
         print(serialise(CAVE["V"]))
         print(serialise(CAVE["E"]))
         print("--")

wat = """
> print(serialise(CAVE["E"]))
("dict(33)")
:print
"""

def test():
   env = a2s([basic])
   expr = parse("(add 1 (add 2 3))")
   res = evaluate(env, expr)
   print("RESULT:", res)

   env = a2s([basic])
   expr = parse("((-> ('a) a) oboe)")
   res = evaluate(env, expr)
   print("RESULT:", res)

tests("""
> (' sym)
sym
> (id 10)
10
> (id (' sym))
sym
> (self oboe)
oboe
> oboe
oboe
> (del oboe)
oboe
> (def? oboe)
false
""")

tests("""
> ()
()
> (type ())
Symbol
> +
:+
> (type +)
Symbol
> true
true
> false
false
> null
null
""")

r_prefix = re.compile(r"^((?:  )+)([^ ].*)")

# followed by a block?
# - put parens around it and the block
# indented section preceded by a colon?
# - put parens around the whole indented section
# not followed by an indented block, and has more than one object?
# - put parens around it

def window(seq, n=3):
   from itertools import islice
   it = iter(seq)
   result = tuple(islice(it, n))
   if len(result) == n:
      yield result    
   for elem in it:
      result = result[1:] + (elem,)
      yield result

alpha = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
alphanum = alpha + alpha.lower() + "0123456789"
op_token = lambda token: len(set(token) & set(alphanum)) == 0

def nudge(line):
   tokens = r_token.findall(line)
   previous = None
   stack = []

   # print("INPUT:", tokens)

   for i, token in enumerate(tokens):
      if token == "(":
         stack.append(i)

      elif token == ")":
         previous = stack.pop()

      else:
         if op_token(token):
            if previous is not None:
               token = tokens[i]
               del tokens[i:i+1]
               tokens[previous:previous] = [token]
         previous = i

   # print("OUTPUT:", tokens)

   text = " ".join(tokens)
   return serialise(parse("(" + text + ")"))[1:-1]

def melt(text):
   from functools import lru_cache
   from itertools import chain

   @lru_cache()
   def indentation(line):
      blocks = 0
      m = r_prefix.match(line)
      if m:
         blocks = len(m.group(1)) / 2
         line = m.group(2)
      return int(blocks), line

   lines = (line for line in text.splitlines() if line)
   grouped = []
   for (a, b, c) in window(chain([":"], lines, [":"]), 3):
      a_depth, a_line = indentation(a)
      b_depth, b_line = indentation(b)
      c_depth, c_line = indentation(c)

      if b_line.startswith(" "):
         raise IndentationError("Irregular indentation: %r" % b_line)

      change = "="
      if a_depth != b_depth:
         change = "<" if (b_depth < a_depth) else ">"
      head = c_depth > b_depth
      start = change == ">"
      end = c_depth < b_depth

      if change == ">":
         if b_depth > (a_depth + 1):
            raise IndentationError("Too much indentation: %r" % b_line)

      group = b_line.endswith(":")
      if head and group:
         b_line = b_line[:-1]

      # print("BEFORE:", b_line)
      b_line = nudge(b_line)
      # print("AFTER:", b_line)

      if head:
         b_line = "(" + b_line # 1(
         grouped.append(group)
      if start:
         if grouped[-1]:
            b_line = "(" + b_line # 2(
      if end:
         pop = b_depth - c_depth
         for i in range(pop):
            group = grouped.pop()
            if group:
               b_line = b_line + ")" # 2)         
         b_line = b_line + (")" * pop) # 1)
      if not head:
         if " " in b_line: # @@ parse
            b_line = "(" + b_line + ")" # 3( 3)
      yield ("  " * b_depth) + b_line
   # print(indentation.cache_info())

# Subset of SRFI 105/110 in a way
# Uses a trailing colon in place of an \\ block

def test_melt():
   print("\n".join(melt("""\
first = 1

name = (arg) ->
  do
    (symbols)
    pro 1
    true
  done

a
  b
    c
    d
      e

' = ('obj) -> obj
id = (obj) -> obj
self = ('sym) ->
  =v sym sym

= '
  -> ('obj) obj
= id
  -> (obj) obj
= self
  -> ('sym)
    =v sym sym

a = 5

' =
  -> ('obj) obj

id =
  -> (obj) obj

self =
  -> ('sym)
    =v sym sym

example1 =
  -> (a b)
    do
      print "a:" a
      print "b:" b
      a + b

example2 => (a b)
  print "a:" a
  print "b:" b
  a + b

' => ('obj)
  obj

id => (obj)
  obj

self => ('sym)
  =v sym sym
""")))

def standard_repl(env):
   while True:
      try: text = input("> ")
      except (EOFError, KeyboardInterrupt):
         print("")
         break

      if (not text) or text.startswith("#"):
         continue

      try: expr = parse(text)
      except Exception as err:
         print("Parse Error:", err)
         continue

      try: env, result = evaluate(env, expr)
      except Exception as err:
         print("Evaluation Error:", err)
         1/0
         continue

      print(serialise(result))

def melt_repl(env):
   mode = "multiple"
   lines = []
   run = False

   while True:
      prompt = {"single": "> ", "multiple": ""}[mode]
      try: text = input(prompt)
      except (EOFError, KeyboardInterrupt):
         print("")
         break

      if (not text) or text.startswith("#"):
         if text == "#":
            mode = "single"
         elif text == "##":
            mode = "multiple"

         if mode == "single":
            continue
         elif text.startswith("#"):
            continue

      if mode == "single":
         run = True
      elif mode == "multiple":
         if (not text):
            run = True
         else:
            run = False

      if run:
         if mode == "multiple":
            text = "\n".join(lines)
            lines = []

         # print("INPUT:", text)
         text = melt(text)
         # print("MELTED:", text)
         text = " ".join(text).replace("\n", "")
         # print("PROCESSED:", text)

         try: expr = parse(text)
         except Exception as err:
            print("Parse Error:", err)
            continue

         try: env, result = evaluate(env, expr)
         except Exception as err:
            print("Evaluation Error:", err)
            1/0
            continue

         if mode == "multiple":
            print("=", serialise(result))
            print("")
         else:
            print(serialise(result))
      elif mode == "multiple":
         lines.append(text)

def main():
   env = a2s([basic])

   # A standard library, of sorts
   env, _ = evaluate(env, parse("(= ' (-> ('obj) obj))"))
   env, _ = evaluate(env, parse("(= id (-> (obj) obj))"))
   env, _ = evaluate(env, parse("(= self (-> ('sym) (=v sym sym)))"))

   if True:
      for (test, expected) in tests.tests:
         try: env, output = evaluate(env, test)
         except Exception as err:
            output = sym("ERROR")
         out = serialise(output)
         assert (out == expected), "%s vs. %s" % (out, expected)

      print("%s tests passed" % len(tests.tests))

   melt_repl(env)

if __name__ == "__main__":
   main()