import ast
import builtins
import json
import resource
import random
import sys
from pathlib import Path

ALLOWED_IMPORTS = {'math','cmath','re','string','collections','itertools','functools','heapq','bisect','typing',
                   'statistics','decimal','fractions','operator','random','copy','datetime','calendar','hashlib',
                   'json','base64','unicodedata','dataclasses','enum','array'}
BLOCKED_CALLS = {'open','exec','eval','compile','globals','locals','vars','getattr','setattr','delattr',
                 '__import__','breakpoint','help','input','exit','quit'}


def validate(source):
    tree = ast.parse(source)
    for node in ast.walk(tree):
        if isinstance(node,(ast.Import,ast.ImportFrom)):
            modules = [alias.name.split('.')[0] for alias in node.names] if isinstance(node,ast.Import) else [(node.module or '').split('.')[0]]
            if isinstance(node,ast.ImportFrom) and node.level:
                raise ValueError('unsupported_relative_import')
            if any(module not in ALLOWED_IMPORTS for module in modules):
                raise ValueError('unsupported_import:'+','.join(modules))
        if isinstance(node,ast.Name) and node.id.startswith('__') and node.id != '__name__':
            raise ValueError('unsupported_introspection')
        if isinstance(node,ast.Attribute) and (node.attr.startswith('_') or isinstance(node.ctx,ast.Store)):
            raise ValueError('unsupported_attribute_access')
        if isinstance(node,ast.Call):
            if isinstance(node.func,ast.Name) and node.func.id in BLOCKED_CALLS:
                raise ValueError('unsupported_call:'+node.func.id)
            if isinstance(node.func,ast.Attribute) and node.func.attr in {'attrgetter','methodcaller'}:
                raise ValueError('unsupported_introspection')
    return tree


def main():
    resource.setrlimit(resource.RLIMIT_CPU,(2,2))
    resource.setrlimit(resource.RLIMIT_AS,(384*1024**2,384*1024**2))
    resource.setrlimit(resource.RLIMIT_FSIZE,(65536,65536))
    resource.setrlimit(resource.RLIMIT_NOFILE,(64,64))
    resource.setrlimit(resource.RLIMIT_CORE,(0,0))
    payload = json.loads(Path('/work/payload.json').read_text())
    nonce = sys.argv[1]
    random.seed(0)
    result = {'passed':False}
    try:
        source = payload['code']
        tree = validate(source)
        namespace = {'__builtins__':dict(vars(builtins)),'__name__':'candidate'}
        exec(compile(tree,'<candidate>','exec'),namespace)
        candidate = namespace.get(payload['entry_point'])
        if not callable(candidate):
            raise ValueError('missing_entry_point')
        test_namespace = {'__builtins__':dict(vars(builtins)),'__name__':'tests'}
        exec(compile(payload['test_code'],'<tests>','exec'),test_namespace)
        test_namespace['check'](candidate)
        result = {'passed':True}
    except BaseException as error:
        result = {'passed':False,'error_type':type(error).__name__,'error':str(error)[:300]}
    print('__WSP_RESULT__'+nonce+json.dumps(result),flush=True)


if __name__ == '__main__':
    main()