from flask import Flask, request
import random
import os
import subprocess
import time
import traceback


app = Flask(__name__)
USER_NAME = "ubuntu"
TIMEOUT_S = 10
CGROUP_PATH = "/sys/fs/cgroup/rocker"
ROCKER_PATH = f"/home/{USER_NAME}/rocker/target/debug/rocker"
HANDLER_FILE = "handler.py"
HANDLER_CODE = ""
with open(f"./{HANDLER_FILE}") as f:
    HANDLER_CODE = f.read()


@app.route("/", methods=["POST"])
def main():
    code = request.json.get("code", "")
    language = request.json.get("language", "")
    _id = str(random.randint(8000_0000, 9000_0000))
    resp = {"msg": "", "code": code, "id": _id, "rt": 0.0}
    
    if len(code) > 1_0000:
        resp["msg"] = "代码过长"
        return resp

    user_path = f"/tmp/{_id}"
    os.mkdir(user_path)

    with open(f"{user_path}/code", "w") as f:
        f.write(code)
    with open(f"{user_path}/{HANDLER_FILE}", "w") as f:
        f.write(HANDLER_CODE)

    starat = time.time()
    try:
        _ = subprocess.run(["sudo", ROCKER_PATH, 
                        "--id", _id,
                        "--image", "base_ubuntu_image", 
                        "--run", f"python3 -u /tmp/handler.py {language} {TIMEOUT_S}",
                        "--env", f"/home/{USER_NAME}/rocker/quant.env", 
                        "--volume", f"{user_path}:/tmp",
                        "--log"])
    except Exception as e:
        resp["msg"] += f"{traceback.format_exc()}\n"
        print(e)
    
    # 20s 超时
    procs_path = f"{CGROUP_PATH}/{_id}/cgroup.procs"
    for _ in range(int(TIMEOUT_S * 2 / 0.01)):  # *2 留出冗余
        time.sleep(0.01)
        try:
            with open(procs_path) as f:
                procs = [proc for proc in f.read().split("\n") if proc != ""]
                if len(procs) == 0:
                    break
        except Exception as e:
            break
    resp["rt"] = time.time() - starat
    
    # 读取文件
    log_path = f"/home/rocker/containers/{_id}/upper/logs/log"
    try:
        with open(log_path) as f:
            log_str = f"{f.read()}\n"
            resp["msg"] += log_str
    except Exception as e:
        resp["msg"] += f"{traceback.format_exc()}\n"
    resp["msg"] = resp["msg"][:10000].rstrip()

    subprocess.run(["sudo", ROCKER_PATH, 
                    "--rm", _id])
    subprocess.run(["sudo", "rm", "-rf", user_path])

    return resp

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=8011, threaded=True)