import json
import os
from flask import Flask, request, jsonify
from multiprocessing import Pool, TimeoutError
import time
import requests
from bs4 import BeautifulSoup
import math
import random
import hashlib
from html import escape
import base64, hmac
from urllib.parse import urlparse
import urllib.request, urllib.error

app = Flask(__name__)

# ────── 1. one-time constants ──────────────────────────────────────────
PORT = int(os.environ.get('PORT', 5001))
TIMEOUT_SECONDS = 5
HTTP_STATUS_CODES = {
    "OK": 200,
    "BAD_REQUEST": 400,
    "INTERNAL_SERVER_ERROR": 500,
}

States = {
    "SUCCESS": "SUCCESS",
    "NOT_A_FUNCTION": "NOT_A_FUNCTION",
    "SCRIPT_ERROR": "ERROR",
    "TIMEOUT": "TIMEOUT",
}

# ────── 2. execution worker ───────────────────────────────────────────
def execute_code(code, request_obj, environment):
    """
    Executes the user's code in a restricted environment.
    """
    logs = []
    start_time = time.time()

    def custom_print(*args, **kwargs):
        # Concatenate all arguments into a single string, separated by spaces
        output = ' '.join(map(str, args))
        logs.append(output)

    restricted_globals = {
        '__builtins__': {
            'print': custom_print,
            # Whitelist safe builtins
            'abs': abs, 'all': all, 'any': any, 'ascii': ascii, 'bin': bin,
            'bool': bool, 'bytearray': bytearray, 'bytes': bytes, 'callable': callable,
            'chr': chr, 'complex': complex, 'dict': dict, 'divmod': divmod,
            'enumerate': enumerate, 'filter': filter, 'float': float, 'format': format,
            'frozenset': frozenset, 'getattr': getattr, 'hasattr': hasattr, 'hash': hash,
            'hex': hex, 'int': int, 'isinstance': isinstance, 'issubclass': issubclass,
            'iter': iter, 'len': len, 'list': list, 'map': map, 'max': max,
            'min': min, 'next': next, 'object': object, 'oct': oct, 'ord': ord,
            'pow': pow, 'range': range, 'repr': repr, 'reversed': reversed,
            'round': round, 'set': set, 'slice': slice, 'sorted': sorted,
            'str': str, 'sum': sum, 'super': super, 'tuple': tuple, 'type': type,
            'zip': zip, 'base64': base64,
            'BaseException': BaseException,
            'Exception': Exception,
            'ArithmeticError': ArithmeticError,
            'AssertionError': AssertionError,
            'AttributeError': AttributeError,
            'EOFError': EOFError,
            'FloatingPointError': FloatingPointError,
            'GeneratorExit': GeneratorExit,
            'ImportError': ImportError,
            'IndexError': IndexError,
            'KeyError': KeyError,
            'KeyboardInterrupt': KeyboardInterrupt,
            'LookupError': LookupError,
            'MemoryError': MemoryError,
            'NameError': NameError,
            'NotImplementedError': NotImplementedError,
            'OSError': OSError,
            'OverflowError': OverflowError,
            'RecursionError': RecursionError,
            'ReferenceError': ReferenceError,
            'RuntimeError': RuntimeError,
            'StopIteration': StopIteration,
            'StopAsyncIteration': StopAsyncIteration,
            'SyntaxError': SyntaxError,
            'SystemError': SystemError,
            'TypeError': TypeError,
            'UnboundLocalError': UnboundLocalError,
            'UnicodeError': UnicodeError,
            'ValueError': ValueError,
            'ZeroDivisionError': ZeroDivisionError
        },
        'request': request_obj,
        'environment': environment,
        'requests': requests,
        'BeautifulSoup': BeautifulSoup,
        'json': json,
        'escape': escape,
        'math': math,
        'random': random,
        'hashlib': hashlib,
        'urlparse': urlparse,
        'urllib.request': urllib.request,
        'urllib.error': urllib.error
    }

    try:
        # Execute the code
        exec(code, restricted_globals)

        # Check if a function (e.g., main) is defined and call it
        if 'main' in restricted_globals and callable(restricted_globals['main']):
            result = restricted_globals['main'](request_obj, environment)
        else:
            # If no main function, maybe the script just runs top-level
            result = None 

        execution_time = (time.time() - start_time) * 1000  # in milliseconds
        return {
            'status': States['SUCCESS'],
            'result': result,
            'logs': logs,
            'environment': environment,
            'execution_time': execution_time,
        }
    except Exception as e:
        execution_time = (time.time() - start_time) * 1000  # in milliseconds
        return {
            'status': States['SCRIPT_ERROR'],
            'result': str(e),
            'logs': logs,
            'environment': environment,
            'execution_time': execution_time,
        }

# ────── 3. process pool ────────────────────────────────────────────────
# It's generally better to initialize the pool once.
# For simplicity in this example, we'll create it on demand, but a real app should manage it.
# pool = Pool(processes=4) 

# ────── 4. API surface ────────────────────────────────────────────────
@app.route("/execute", methods=['POST'])
def execute():
    body = request.json
    code = body.get('code', '')
    request_obj = body.get('request', {})
    environment = body.get('environment', {})

    with Pool(processes=1) as pool:
        async_result = pool.apply_async(execute_code, (code, request_obj, environment))
        try:
            payload = async_result.get(timeout=TIMEOUT_SECONDS)
        except TimeoutError:
            payload = {
                'status': States['TIMEOUT'],
                'result': f'Execution timed out after {TIMEOUT_SECONDS} seconds.',
                'logs': [],
                'environment': environment,
                'execution_time': TIMEOUT_SECONDS * 1000,
            }
        except Exception as e:
             payload = {
                'status': States['SCRIPT_ERROR'],
                'result': str(e),
                'logs': [],
                'environment': environment,
            }

    return jsonify(payload)

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=PORT)