import os
import sys
import subprocess
import traceback

TARGET_LANGUAGE = set(["python", "cpp", "go", "rust", "node", "javascript"])
PROJECT_PATH = "/home/rocker/project"
BUILD_TIMEOUT_S = 3

def go_handler():
    source_path = f"{PROJECT_PATH}/go_project/main.go"
    subprocess.run(["cp", "/tmp/code", source_path], capture_output=True)
    main_path = f"{PROJECT_PATH}/go_project/main"
    ret = subprocess.run(["/usr/bin/timeout", f"{BUILD_TIMEOUT_S}s", "go", "build", "-o", main_path], cwd=f"{PROJECT_PATH}/go_project", capture_output=True)
    std_err = ret.stderr.decode("utf-8")
    if len(std_err) > 0:
        return std_err, ""
    return "", main_path

def cpp_handler():
    source_path = f"{PROJECT_PATH}/cpp_project/main.cpp"
    subprocess.run(["cp", "/tmp/code", source_path], capture_output=True)
    main_path = f"{PROJECT_PATH}/cpp_project/main"
    ret = subprocess.run(["/usr/bin/timeout", f"{BUILD_TIMEOUT_S}s", "g++", source_path, "-o", main_path], capture_output=True)
    std_err = ret.stderr.decode("utf-8")
    if len(std_err) > 0:
        return std_err, ""
    return "", main_path

def node_handler():
    main_path = f"{PROJECT_PATH}/node_project/main.js"
    subprocess.run(["cp", "/tmp/code", main_path], capture_output=True)
    return "", main_path

def rust_handler():
    source_path = f"{PROJECT_PATH}/rust_project/src/main.rs"
    subprocess.run(["cp", "/tmp/code", source_path], capture_output=True)
    ret = subprocess.run(["/usr/bin/timeout", f"{BUILD_TIMEOUT_S}s", "cargo", "build"], cwd=f"{PROJECT_PATH}/rust_project", capture_output=True)
    std_err = ret.stderr.decode("utf-8")
    if len(std_err) > 0 and ret.returncode != 0:
        return std_err, ""
    main_path = f"{PROJECT_PATH}/rust_project/target/debug/main"
    subprocess.run(["mv", f"{PROJECT_PATH}/rust_project/target/debug/rust_project", main_path], capture_output=True)
    return "", main_path

def python_handler():
    main_path = f"{PROJECT_PATH}/python_project/main.py"
    subprocess.run(["cp", "/tmp/code", main_path], capture_output=True)
    return "", main_path


def main():
    language = sys.argv[1]
    timeout_s = sys.argv[2]
    if language not in TARGET_LANGUAGE:
        print(f"Unsupported programming language: {language}")
        return
    
    # 编译和执行
    cmd = ["/usr/bin/time", "-f", "'%M %U %S %e'", "/usr/bin/timeout", f"{timeout_s}s"]
    err = ""
    if language == "python":
        err, main_path = python_handler()
        cmd.extend(["python", "-u", main_path])
    elif language == "cpp":
        err, main_path = cpp_handler()
        cmd.extend([main_path])
    elif language == "go":
        err, main_path = go_handler()
        cmd.extend([main_path])
    elif language == "node":
        err, main_path = node_handler()
        cmd.extend(["node", main_path])
    elif language == "rust":
        err, main_path = rust_handler()
        cmd.extend([main_path])
    elif language == "javascript":
        err, main_path = node_handler()
        cmd.extend(["node", main_path])
    else:
        ...

    # 阻塞执行
    if err != "":
        print(err)
    else:
        ret = subprocess.run(cmd, capture_output=True)
        out_lines = f"{ret.stdout[:1_0000].decode('utf-8')} \n {ret.stderr[:1_0000].decode('utf-8')}".rstrip().split("\n")
        try:
            end_line_lis = out_lines[-1].replace("'", "").strip().split(" ")
            end_line = f"MaxRSS: {int(end_line_lis[0])//1024}mb, User: {end_line_lis[1]}s, Sys: {end_line_lis[2]}s, Elapsed: {end_line_lis[3]}s"
            out_lines[-1] = end_line
        except Exception as e:
            print(traceback.format_exc())
        out = "\n".join(out_lines)
        print(out)


if __name__ == "__main__":
    main()