"""
Copyright 2020 The OneFlow Authors. All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""
import functools


class SessionStatus:
    OPEN = "OPEN"

    RUNNING = "RUNNING"

    CLOSED = "CLOSED"


def GetDefaultSession():
    assert _default_session is not None
    return _default_session


def OpenDefaultSession(sess):
    global _default_session
    assert _default_session is None
    _default_session = sess


def TryCloseDefaultSession():
    global _default_session
    assert _default_session is not None
    if _default_session is not None:
        _default_session.TryClose()
    _default_session = None


def try_init_default_session(func):
    @functools.wraps(func)
    def Func(*args, **kwargs):
        GetDefaultSession().TryInit()
        return func(*args, **kwargs)

    return Func


_default_session = None
