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
BASE_CODE_PATH = "/tmp/code"

def go_handler():
    source_path = f"{PROJECT_PATH}/go_project/main.go"
    subprocess.run(["cp", BASE_CODE_PATH, 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", BASE_CODE_PATH, 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", BASE_CODE_PATH, main_path], capture_output=True)
    return "", main_path

def rust_handler():
    lines = []
    is_have_main = False
    with open(BASE_CODE_PATH, "r") as f:
        for line in f.readlines():
            lines.append(line)
            if "fn" in line and "main" in line:
                is_have_main = True
    if is_have_main == False:
        lines.insert(0, "fn main() {\n")
        lines.append("}\n")
    
    source_path = f"{PROJECT_PATH}/rust_project/src/main.rs"
    with open(source_path, "w") as f:
        f.writelines(lines)

    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", BASE_CODE_PATH, 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()