import docker
from func_timeout import func_set_timeout, exceptions
from docker.errors import ImageNotFound
from common import *

logger = logging.getLogger(__name__)
logger.setLevel(level=logging.DEBUG)

client = docker.from_env()


def get_and_set_image(name):
    if Languages[name].image:
        return Languages[name].image
    try:
        image = client.images.get(TAG_STR.format(name))
        Languages[name].set_image(image)
        return image
    except ImageNotFound:
        pass
    return None


def load():
    for name, value in Languages.items():
        if get_and_set_image(name) is None:
            value.build_image()
            get_and_set_image(name)


def execute(code, _type, stdin=""):
    if _type not in Languages or Languages[_type].image is None:
        raise Exception("找不到运行环境，请等待后台创建")
    result = ""
    command = ""
    try:
        # step1 生成代码文件
        if callable(Languages[_type].file_way):
            # 支持特殊文件名的语言
            command += Languages[_type].file_way(code)
        else:
            command = Languages[_type].file_way + code + "\nEOF\n"

        # step2 生成stdin，通过管道给到代码
        if stdin != "":
            command += "cat > stdin << EOF\n" + stdin + "\nEOF\n"
            command += "cat stdin | "

        # step3 执行代码
        if callable(Languages[_type].run_way):
            command += Languages[_type].run_way(code)
        else:
            command += Languages[_type].run_way

        # print(command)
        container = client.containers.run(Languages[_type].image, stderr=True, stdout=True, tty=True,
                                          network_disabled=True,
                                          detach=True,
                                          command=["/bin/sh", "-c", command])

        @func_set_timeout(10)
        def execute_withtimeout(container):
            container.wait()

        execute_withtimeout(container)

    except exceptions.FunctionTimedOut as e:
        logger.error("执行超时 %s", e)
        result = bytes.decode(container.logs())
        result += "\n代码执行超时"
        container.kill()
    finally:
        if not result:
            result = bytes.decode(container.logs())
        container.stop()
        container.remove()
    return result


if __name__ == "__main__":
    load()
    res = execute("a=input()\nprint(a=='1234')\n", "python", "1234")
    print(res)