# %%
from typing import List, Tuple, Dict, Set
import logging
import sys
import os
import string
import keyword
import unicodedata
import json
import re

# Get the current recursion limit
# print(f"Current recursion limit: {sys.getrecursionlimit()}")

# Set a new recursion limit (e.g., to 2000)
sys.setrecursionlimit(1000000)
# print(f"New recursion limit: {sys.getrecursionlimit()}")

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Optional training-time early-stop (performance)
# ---------------------------------------------------------------------------
#
# The online PCFG generator can produce very long trace prompts. When training with
# `drop_truncated=True` and a fixed context window (e.g. 4096), any trace line whose
# prompt exceeds the window will be dropped after tokenization. Continuing to execute
# the program past that point wastes CPU and can cause frequent timeouts in DataLoader
# workers (GPU idle).
#
# We expose lightweight early-stop guards here so the runner can abort execution once
# the prompt/context becomes obviously too large, *without* changing semantics when
# disabled. These are intentionally opt-in via env vars.
#
# - MICROPY_MAX_PROMPT_CHARS: abort once `len(next_prompt)` reaches this many chars
#   (cheap O(1) check that avoids scanning huge strings).
# - MICROPY_MAX_CONTEXT_TOKENS: abort once `max_length` (whitespace-token estimate)
#   reaches this many tokens.
#
# The corresponding exceptions are caught by `pcfg_backend.ProgramLinesDataset` and
# treated as a normal "truncation" (like timeout/output-cap), so downstream parsing
# remains clean.
#
class MaxPromptExceeded(Exception):
    pass


class MaxContextExceeded(Exception):
    pass


try:
    MICROPY_MAX_PROMPT_CHARS = int(os.environ.get("MICROPY_MAX_PROMPT_CHARS", "0") or "0")
except Exception:
    MICROPY_MAX_PROMPT_CHARS = 0
try:
    MICROPY_MAX_CONTEXT_TOKENS = int(os.environ.get("MICROPY_MAX_CONTEXT_TOKENS", "0") or "0")
except Exception:
    MICROPY_MAX_CONTEXT_TOKENS = 0

def list_str(l:List):
    # NOTE: this used to be implemented recursively, which becomes extremely slow
    # (and can hit recursion depth) for large traces produced by the generator.
    # Keep identical formatting but do it iteratively.
    if not l:
        return ""
    return " ".join(str(x) for x in l)

def suffix_str(l:List):
    # Kept for backward-compat (used only by older code paths).
    if not l:
        return ""
    return " " + " ".join(str(x) for x in l)

def list_str_tokens(l: List):
    if len(l) == 0:
        return "empty"
    return list_str(l)

class Effect(): pass


# %%
def clear_state():
    global base_state
    global bit_index
    global max_length
    max_length = 0
    bit_index = 0
    base_state = Effect([])

def clear_funs():
    global base_funs
    global exp_index
    global exp_defs
    base_funs = []
    exp_index = 0
    exp_defs = []

def print_response(s:str):
    #This defines PENCIL scaffolding

    #Here the scaffolding is defined on strings but it should really be defined on tokens.
    #We assume a trivial tokenizer defined by space --- a token is any sequence of characters not containing a space.
    #We assume there are never two spaces in a row.
    #We satisfy the invariant that all expressions denoting strings do not start or end in space.
    #Concatenation must add one space character between the strings being concatenated.
    #We assume that [call], => and [ret] single tokens under this tokenizer.
    
    global next_prompt  # next_prompt is initialized in run_program
    global max_length

    # Fast early-stop: if the prompt prefix is already enormous, abort before doing any
    # expensive scanning/printing. This is useful when training drops truncated prompts.
    try:
        _max_prompt_chars = int(MICROPY_MAX_PROMPT_CHARS or 0)
    except Exception:
        _max_prompt_chars = 0
    if _max_prompt_chars > 0 and next_prompt and len(next_prompt) >= _max_prompt_chars:
        raise MaxPromptExceeded(f"prompt chars {len(next_prompt)} >= {_max_prompt_chars}")

    max_length = max(max_length, ntokens(s))
    try:
        _max_ctx = int(MICROPY_MAX_CONTEXT_TOKENS or 0)
    except Exception:
        _max_ctx = 0
    if _max_ctx > 0 and max_length >= _max_ctx:
        raise MaxContextExceeded(f"max context {max_length} >= {_max_ctx} tokens")

    #print(f"{s};input to print_response")
    #print(f"{s[:len(next_prompt)]}; prefix of input")
    #print(f"{next_prompt}; next prompt")
    
    print(f"{s[:len(next_prompt)]} [EOP] {s[len(next_prompt)+1:]}")

    if s[len(s)-5:] != "[ret]":
        print()
        print("response does not end in [ret]")
        raise Exception("Response does not end in [ret]")
    
    if s[:len(next_prompt)] != next_prompt:
        print()
        print("PENCIL Error")
        raise Exception(f"PENCIL Error")
      
    arrow_index = s.rfind(" => ")
    call_index = s[:arrow_index].rfind(" [call] ")
    response = s[arrow_index+4:len(s)-6] #begins with a non-space
    next_prompt = f"{s[:call_index]} {response}" #obeys space character invariants
    #print(f"{next_prompt}; next prompt as computed")
    
def ntokens(s:str):
    """
    Count "tokens" under the internal MicroPy whitespace-token assumption:
    a token is any sequence of characters not containing a space.

    This is on the hot path during generator execution (called from
    `print_response()` for every trace line). The prior recursive implementation
    was O(n^2) due to repeated slicing + rfind, and can stall training by
    triggering frequent `Timeout expired ...` in DataLoader workers.
    """
    if not s:
        return 0

    # Fast path: if the string obeys the (typical) invariant of single-space
    # separation and no leading/trailing spaces, token count is spaces + 1.
    #
    # We still handle rare leading/trailing spaces and double spaces robustly.
    start = 0
    end = len(s)
    while start < end and s[start] == " ":
        start += 1
    while end > start and s[end - 1] == " ":
        end -= 1
    if start >= end:
        return 0
    view = s if (start == 0 and end == len(s)) else s[start:end]
    if "  " not in view:
        return view.count(" ") + 1

    # Robust fallback for double-space cases (still O(n), no allocations).
    n = 0
    in_tok = False
    for ch in view:
        if ch == " ":
            in_tok = False
        else:
            if not in_tok:
                n += 1
                in_tok = True
    return n




# %%
class Val_Or_TailApply(): pass
    #this class is only used in type declarations as a comment.
 
class Value(Val_Or_TailApply): pass

class Obj(Value):
    def __init__(self,name:str):
        self.name = name
    def __str__(self):
        return f"O . {self.name}"

class Bool(Value):
    def __init__(self,name:bool):
        self.name = name
    def __str__(self):
        return f"{self.name}" # here we've removed the B . prefix for shorter token bool, making the model less likely to bias towards the wrong bool
    
class Unit(Value):
    #A pure effect expression has value "unit".
    def __init__(self):
        pass
    def __str__(self):
        return f"U . Unit"

class Failure(Value):
    #used to indicate a failure of a brachnin in Try
    def __init__(self):
        pass
    def __str__(self):
        return f"Fa . Failure"
    
true_ = Bool(True)
false_ = Bool(False)
unit_ = Unit()
fail_ = Failure()

class Attr():
    def __init__(self,name:str):
        self.name = name
    def __str__(self):
        return f"Att . {self.name}"
    
class Assertion():
    def __init__(self, obj:Obj, attr:Attr, val:Value):
        self.obj = obj
        self.attr = attr
        self.val = val
    def __str__(self):
        return f"Assertion ( {self.obj} {self.attr} {self.val} )"

class Effect():
    def __init__(self,assertions:List[Assertion]):
        self.assertions = assertions
    def __str__(self):
        return f"Eff ( {list_str_tokens(self.assertions)} )"

def add_effect(ef1:Effect,ef2:Effect) -> Effect:
    return Effect(ef1.assertions + ef2.assertions)

def lookup_attr(obj:Obj,attr:Attr,s:Effect) -> Value:
    for assertion in reversed(s.assertions):
        if obj.name == assertion.obj.name and attr.name == assertion.attr.name:
            return assertion.val
    raise Exception(f"no value fournd for {obj}.{attr}")

def MPhasattr(obj:Obj,attr:Attr,s:Effect) -> bool:
    for assertion in s.assertions:
        if obj.name == assertion.obj.name and attr.name == assertion.attr.name:
            return true_
    return false_

class Binding():
    #for local variables.
    def __init__(self,name:str,val:Value):
        self.name = name
        self.val = val
    def __str__(self):
        return f"Bind ( {self.name} {self.val} )"

class Env():
    def __init__(self,bindings:List[Binding]):
        self.bindings = bindings
    def __str__(self):
        return f"Env ( {list_str(self.bindings)} )"

def lookup_var(name:str,env:Env) -> Value:
    for binding in env.bindings:
        if binding.name == name:
            return binding.val
    raise Exception(f"no value found for variable {name}")

class Exp(): pass
    #There wil be various subclasses of Exp (expressions)
    #every Exp will have a name such as Exp32 or Exp157
    #the expression self will print using the names of it subexpressions
    #for example If(Exp57,Exp78,Exp103)

#Readers are advised to first skip ahead and read some expression class definitions and example runs

def set_name(e:Exp):
    global exp_index
    global exp_defs
    exp_index += 1
    name = f"Exp{exp_index}"
    exp_defs = exp_defs + [f"D . {name} = {e}"]
    e.name = name

class TailApply():
    #in a TailApply the arguments are values
    #this is different from an App where the arguments are expressions needing evaluation
    #A TailApply is closed --- there are no free variables and hence no dependence on a binding environment
    def __init__(self,f,args: List[Value]):
        self.f = f
        self.args = args
    def __str__(self):
        return f"TailApp ( {self.f.name} Args ( {list_str(self.args)} ) )"

# %%
def name(e):
    return e.name if isinstance(e,Exp) else e

def run_program(e:Exp):
    global fun_defs
    global exp_defs
    global base_state
    def MicroPy_def(fun):
        return f"FD {fun.name} = lambda ( {list_str(fun.varnames)} ) {name(fun.body())} ;"
    fundefs = list_str([MicroPy_def(f) for f in base_funs])
    base_stack = f"FunDefs ( {fundefs} ) ********* ExpDefs ( {list_str(exp_defs)} ) ****** {base_state} ********* Env ( )"
    global next_prompt
    next_prompt = (f"{base_stack} [call] [call] {name(e)}")
    ef,val = eval(e,base_stack,base_state,Env([]))
    base_state = add_effect(base_state,ef)
    print(f"{ef} {val}      max context = {max_length} tokens")

def eval(e, stack: str, s: Effect, env:Env)->Tuple[Effect,Value]:
    #e is a Value or Exp, s is the state
    #the stack is extended by [call] e.
    #most rewrites replace a call with a call (computation within a stack frame)
    #the [call] token is removed when the stack frame returns.
    #eval must return a value (not a tail call).

    if isinstance(e,(Value,int)):
        print_response (f"{stack} [call] {e} => {Effect([])} {e} [ret]") #the stack frame returns
        return Effect([]),e

    ef,cval = tailval(e, f"{stack} [call]", s, env)
    #tail calls allow for loops
    #the new [call] token creates a "loop frame" for accumulating effects.

    while isinstance(cval,TailApply): #Tail Call Loop
        print_response(f"{stack} [call] {ef} {cval} => [call] {ef} [call] {cval} [ret]")
        ef2, next_cval = compute_TailApply(cval,f"{stack} [call] {ef}",add_effect(s,ef))
        next_ef = add_effect(ef,ef2)
        print_response(f"{stack} [call] {ef} {ef2} {next_cval} => [call] {next_ef} {next_cval} [ret]")
        cval = next_cval
        ef = next_ef
    print_response(f"{stack} [call] {ef} {cval} => {ef} {cval} [ret]") #return from the loop
    return ef, cval
    
def tailval(e, stack: str, s: Effect, env:Env) -> Tuple[Effect,Val_Or_TailApply]:
    #e is a Value or Exp
    if isinstance(e,(Value,int)):
        print_response (f"{stack} [call] {e} => {Effect([])} {e} [ret]")
        return Effect([]),e
    
    #e must be an Exp
    print_response(f"{stack} [call] {name(e)} => [call] {e} [ret]")
    return e.tailval(stack,s,env)

def compute_TailApply(app:TailApply,stack:str,s:Effect) -> Tuple[Effect,Val_Or_TailApply]:
    #app is an application of a defined function to values (not expressions)
    body = app.f.body()
    print_response(f"{stack} [call] {app} => [call] let ( {list_str_tokens(app.f.varnames)} ) be ( {list_str_tokens(app.args)} ) in {name(body)} [ret]")
    env = Env([Binding(name,val) for name,val in zip(app.f.varnames,app.args)])
    print_response(f"{stack} [call] let ( {list_str_tokens(app.f.varnames)} ) be ( {list_str_tokens(app.args)} ) in {name(body)} => [call] {env} [call] {name(body)} [ret]")
    ef,val = tailval(body,f"{stack} [call] {env}",s,env)
    print_response(f"{stack} [call] {env} {ef} {val} => {ef} {val} [ret]") #return
    return ef,val



# %%
class Equal(Exp):
    # a simple case
    # here tailval always returns a value (a Boolean)
    def __init__(self,exp1:Exp,exp2:Exp):
        self.exp1 = exp1
        self.exp2 = exp2
        set_name(self)
    def __str__(self):
        return f"Equal ( {name(self.exp1)} {name(self.exp2)} )"
    
    def tailval(self,stack:str,s:Effect,env:Env) -> Tuple[Effect,Bool]:
        #in any tailval method the context is f"{stack}[call]{self}
        #The tailval method must rewrite [call]{self} to a (possibly empty) effect and a VAl_OR_TailApply
        #The EQUAL expressions always evaluate to a Boolean.
        ef1,val1 = eval(self.exp1,f"{stack} [call] {self} exp1 =",s,env)
        ef2,val2 = eval(self.exp2,f"{stack} [call] {self} exp1 = {ef1} {val1} exp2 =",add_effect(s,ef1),env)
        combined_ef = add_effect(ef1,ef2)
        rval = true_ if str(val1) == str(val2) else false_
        print_response(f"{stack} [call] {self} exp1 = {ef1} {val1} exp2 = {ef2} {val2} => {combined_ef} {rval} [ret]") #return to stack
        return combined_ef,rval

# %%
clear_funs()
clear_state()
# run_program(Equal(Obj("A"),Obj("B")))

# print()
clear_funs()
clear_state()
# run_program(Equal(Obj("A"),Obj("A")))

# %%
class App(Exp):
    #an application of a function to argument expressions
    #this is different from TailApply where the arguments are values

    #the tailval method first pops the stack frame that computes the values of the argument expressions
    #and then returns a TailApply (tail call optimization)

    def __init__(self,f,args:List[Exp]):
        self.f = f
        self.args = args
        set_name(self)

    def __str__(self):
        return f"App ( {self.f.name} ( {list_str([name(e) for e in self.args])} ) )"
    
    def tailval(self, stack:str, s:Effect, env:Env) -> Tuple[Effect,Val_Or_TailApply]:
        argvals = []
        args_effect = Effect([])
        args_state=s
        args_stack = f"{stack} [call] {self}"
        for arg,i in zip(self.args,range(len(self.args))):
            ef,val = eval(arg,f"{args_stack} arg{i} =",args_state,env)
            argvals = argvals + [val]
            args_effect = add_effect(args_effect,ef)
            args_state = add_effect(args_state,ef)
            args_stack = f"{args_stack} arg{i} = {ef} {val}"
        
        app = TailApply(self.f,argvals)
        print_response(f"{args_stack} => {args_effect} {app} [ret]") #return to call in (original) stack.
        return args_effect, app


# %%
class LookupVar(Exp):
    #lookup a local variable (register value)
    def __init__(self, var:str):
        self.var = var;
        set_name(self)
    
    def __str__(self):
        return f"LookupVar ( {self.var} )"
    
    def tailval(self, stack:str, s:Effect,env:Env) -> Tuple[Effect,Value]:
        val = lookup_var(self.var,env)
        ef = Effect([])
        print_response(f"{stack} [call] {self} => {ef} {val} [ret]")
        return ef,val

# %%
#First look at the examples in the test cells.

class MicroPyFun():
    def __init__(self,func):
        self.func = func
        self.name = func.__name__
        self.varnames = func.__code__.co_varnames
    def body(self):
        if hasattr(self,"body_cache"): return self.body_cache
        self.body_cache = self.func(*[LookupVar(name) for name in self.varnames])
        return self.body_cache
    def __str__(self):
        return f"F.{self.name}"
    
# For a recursive function f we must wait until the name f is redefined by the decorator MicroPy before computing the body.
# Without this delay (for a recursive function) the body becomes an infinte expression.

def make_symbolic(fun):
    return lambda *args: App(fun,args)

def MicroPy(func):
    global base_funs
    f = MicroPyFun(func)
    base_funs += [f]
    return make_symbolic(f)



# %%
clear_funs()
clear_state()

@MicroPy
def f(x,y):
    return Equal(x,y)

# run_program(f(Obj("A"),Obj("B")))

# %%
class Seq(Exp):
    #It simpler, and ok, to use @microPy def seq(x,y) return y.

    def __init__(self,x:Exp,y:Exp):
        self.x = x
        self.y = y
        set_name(self)

    def __str__(self):
        return f"Seq ( {name(self.x)} , {name(self.y)} )"
    
    def tailval(self, stack:str, s:Effect, env:Env) -> Tuple[Effect,Val_Or_TailApply]:
        xef,xval = eval(self.x,f"{stack} [call] {self} x =",s,env)
        print_response(f"{stack} [call] {self} x = {xef} {xval} => [call] {xef} [call] {name(self.y)} [ret]")
        yef,yval = tailval(self.y,f"{stack} [call] {xef}",add_effect(s,xef),env)
        combined_ef = add_effect(xef,yef)
        print_response(f"{stack} [call] {xef} {yef} {yval} => {combined_ef} {yval} [ret]")
        return combined_ef,yval

# %%
clear_funs()
clear_state()

# run_program(Seq(Obj("A"),Obj("B")))

# %%
class If(Exp):
    #this is similar to seq
    def __init__(self, test:Exp, then:Exp, els:Exp):
        self.test = test
        self.then = then
        self.els = els
        set_name(self)

    def __str__(self):
        return f"If ( {name(self.test)} {name(self.then)} {name(self.els)} )"
    
    def tailval(self,stack:str,s:Effect,env:Env) -> Tuple[Effect,Val_Or_TailApply]:
        testef,testval = eval(self.test,f"{stack} [call] {self} test =",s,env)
        case = self.then if testval == true_ else self.els
        print_response(f"{stack} [call] {self} test = {testef} {testval} => [call] {testef} [call] {name(case)} [ret]")
        caseef,caseval = tailval(case,f"{stack} [call] {testef}",add_effect(s,testef),env)
        combined_ef = add_effect(testef,caseef)
        print_response(f"{stack} [call] {testef} {caseef} {caseval} => {combined_ef} {caseval} [ret]")
        return combined_ef,caseval

# %%
clear_funs()
clear_state()

# run_program(If(true_,Obj("a"),Obj("b")))
# print("")

# clear_funs()
# clear_state()

# run_program(If(false_,Obj("a"),Obj("b")))
# print("")

# %%
class Assert(Exp):
    def __init__(self, objexp:Exp, attr:Attr, valexp:Exp):
        self.objexp = objexp
        self.attr = attr
        self.valexp = valexp
        set_name(self)

    def __str__(self):
        return f"Assert ( {name(self.objexp)} {self.attr} {name(self.valexp)} )"
    
    def tailval(self,stack,s:Effect,env:Env) -> Tuple[Effect,Unit]:
        ef1,obj = eval(self.objexp,f"{stack} [call] {self} obj =",s, env)
        ef2,val = eval(self.valexp,f"{stack} [call] {self} obj = {ef1} {obj} val =", add_effect(s,ef1), env)
        total_effect = add_effect(ef1, add_effect(ef2, Effect([Assertion(obj,self.attr,val)])))
        print_response(f"{stack} [call] {self} obj = {ef1} {obj} val = {ef2} {val} => {total_effect} {unit_} [ret]")
        return total_effect,unit_


# %%
clear_funs()
clear_state()

# run_program(Assert(Obj("A"), Attr("B"), Obj("C")))

# %%
class HasAttr(Exp):
    def __init__(self, objexp:Exp, attr:Attr):
        self.objexp = objexp
        self.attr = attr
        set_name(self)

    def __str__(self):
        return f"HasAttr ( {name(self.objexp)} {self.attr} )"
    
    def tailval(self,stack,s:Effect,env:Env) -> Tuple[Effect,Value]:
        ef,obj = eval(self.objexp,f"{stack} [call] {self} obj =",s, env)
        val = MPhasattr(obj,self.attr,add_effect(s,ef))
        print_response(f"{stack} [call] {self} obj = {ef} {obj} => {ef} {val} [ret]")
        return ef,val

# %%
clear_funs()
clear_state()

# run_program(Assert(Obj("A"), Attr("B"), Obj("C")))
# print("")

# run_program(HasAttr(Obj("A"),Attr("B")))
# print("")

# run_program(HasAttr(Obj("A"),Attr("D")))

# %%
class LookupAttr(Exp):
    def __init__(self, objexp:Exp, attr:Attr, return_obj: bool = False):
        self.objexp = objexp
        self.attr = attr
        # When True, still *perform* the lookup (so missing attrs still error),
        # but return the evaluated object instead of the attribute value.
        #
        # This is useful for "spine" usages where we want a LookupAttr frame on the
        # stack while evaluating `objexp`, but we don't want the (often non-object)
        # attribute value to become the value threaded upward in the plan compiler.
        self.return_obj = bool(return_obj)
        set_name(self)

    def __str__(self):
        return f"LookupAttr ( {name(self.objexp)} {self.attr} )"
    
    def tailval(self,stack,s:Effect,env:Env) -> Tuple[Effect,Value]:
        ef,obj = eval(self.objexp,f"{stack} [call] {self} obj =",s, env)
        val = lookup_attr(obj,self.attr,add_effect(s,ef))
        out = obj if bool(self.return_obj) else val
        print_response(f"{stack} [call] {self} obj = {ef} {obj} => {ef} {out} [ret]")
        return ef,out

# %%
clear_funs()
clear_state()

# run_program(Assert(Obj("A"), Attr("B"), Obj("C")))
# print("")
# run_program(LookupAttr(Obj("A"),Attr("B")))

# %%
class Try(Exp):
    def __init__(self,first_attempt:'Exp',second_attempt:'Exp'):
        self.first_attempt = first_attempt
        self.second_attempt = second_attempt
        set_name(self)

    def __str__(self):
        return f"Try ( {name(self.first_attempt)} {name(self.second_attempt)} )"
    
    def tailval(self,stack:str,s:Effect,env:Env) -> Tuple[Effect,Val_Or_TailApply]:
        ef,v = eval(self.first_attempt,f"{stack} [call] {self} [call] first =",s,env)
        if v != fail_:
            print_response(f"{stack} [call] {self} [call] first = {ef} {v} => {ef} {v} [ret]")
            print_response(f"{stack} [call] {self} {ef} {v} => {ef} {v} [ret]")
            return ef,v
        print_response(f"{stack} [call] {self} [call] first = {ef} {v} => second = [ret]")
        ef,v = eval(self.second_attempt,f"{stack} [call] {self} second =",s,env)
        if v != fail_:
            print_response(f"{stack} [call] {self} second = {ef} {v} => {ef} {v} [ret]")
            return ef,v
        noef = Effect([])
        print_response(f"{stack} [call] {self} second = {ef} {v} => {noef} {fail_} [ret]")
        return noef,fail_

# %%
clear_funs()
clear_state()

# run_program(Try(Seq(Assert(Obj("a"),Attr("b"),Obj("c")),Obj("success")),Seq(Assert(Obj("a"),Attr("b"),Obj("d")),Obj("success"))))
# print("")

# clear_state()
# run_program(Try(Seq(Assert(Obj("a"),Attr("b"),Obj("c")),fail_),Seq(Assert(Obj("a"),Attr("b"),Obj("d")),Obj("success"))))


# %%