from flask import Flask, request, jsonify
from jupyter_client import KernelManager
import sys, io, time, queue, os


app = Flask(__name__)
TIMEOUT = os.getenv("TIMEOUT", 60)

km = KernelManager(kernel_name='python3')

km.start_kernel()
kc = km.client()

@app.route('/execute', methods=['POST'])
def execute_code():
    global kc
    code = request.json.get('code')
    if not code:
        return jsonify({'error': 'No code provided'}), 400
    
    outputs = []
    start_time = time.time()

    _ = kc.execute(code)

    success = True
    error = None

    while True:
        try:
            # if we timed out we interrupt the kernel so we can send the next request
            if time.time() - start_time > TIMEOUT:
                outputs.append({"text": f"Code execution has timed out (max {TIMEOUT}sec)."})
                success = False
                error = "TimeOut"
                km.interrupt_kernel()
                break
            
            # get message from output buffer
            msg = kc.get_iopub_msg(timeout=1)

            # if the kernel is idle again we can wrap up
            if msg['header']['msg_type'] == 'status' and msg['content']['execution_state'] == 'idle':
                break
            
            # save output messages
            if msg['header']['msg_type'] == 'stream':
                outputs.append({"text": msg['content']['text']})

            # handle error messages
            elif msg['header']['msg_type'] == 'error':
                outputs.append({"text": "\n".join(msg['content']['traceback'])})
                success = False
                error = msg["content"]["ename"]

            # handle figures
            elif msg['header']['msg_type'] == 'display_data':
                msg["content"]["data"]["text"] = msg["content"]["data"].pop("text/plain")
                outputs.append(msg["content"]["data"])
        
        # if queue is empty we try again, unless the kernel is dead
        except queue.Empty:
            print("No message received (timeout)", )
            if not kc.is_alive():
                outputs.append({"text": "Kernel has died!"})
                success = False
                error = "KernelDied"
                km.restart_kernel()
                break

    return jsonify({
        'result': outputs,
        'success': success,
        'error': error
    })

@app.route('/restart', methods=['POST'])
def restart_kernel():
    global kc
    km.restart_kernel()
    kc = km.client()
    return jsonify(status="Kernel has been restarted."), 200

@app.route('/health', methods=['GET'])
def health_check():
    return jsonify(status="healthy"), 200

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