import json
import os
from flask import Flask, request, jsonify
from multiprocessing import Pool, TimeoutError
import time
import requests
from bs4 import BeautifulSoup

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
        },
        'request': request_obj,
        'environment': environment,
        'requests': requests,
        'BeautifulSoup': BeautifulSoup,
        'json': json
    }

    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)