import os
import pytest
import subprocess
from io import StringIO
import sys

from fastNLP.core.utils.cache_results import cache_results
from fastNLP.core import rank_zero_rm


def get_subprocess_results(cmd):
    output = subprocess.check_output(cmd, shell=True)
    return output.decode('utf8')


class Capturing(list):
    # 用来捕获当前环境中的stdout和stderr，会将其中stderr的输出拼接在stdout的输出后面
    def __enter__(self):
        self._stdout = sys.stdout
        self._stderr = sys.stderr
        sys.stdout = self._stringio = StringIO()
        sys.stderr = self._stringioerr = StringIO()
        return self

    def __exit__(self, *args):
        self.append(self._stringio.getvalue() + self._stringioerr.getvalue())
        del self._stringio, self._stringioerr    # free up some memory
        sys.stdout = self._stdout
        sys.stderr = self._stderr


class TestCacheResults:
    def test_cache_save(self):
        cache_fp = 'demo.pkl'
        try:
            @cache_results(cache_fp)
            def demo():
                print("¥")
                return 1
            res = demo()
            with Capturing() as output:
                res = demo()
            assert '¥' not in output[0]

        finally:
            rank_zero_rm(cache_fp)

    def test_cache_save_refresh(self):
        cache_fp = 'demo.pkl'
        try:
            @cache_results(cache_fp, _refresh=True)
            def demo():
                print("¥")
                return 1

            res = demo()
            with Capturing() as output:
                res = demo()
            assert '¥' in output[0]
        finally:
            rank_zero_rm(cache_fp)

    def test_cache_no_func_change(self):
        cache_fp = os.path.abspath('demo.pkl')
        try:
            @cache_results(cache_fp)
            def demo():
                print('¥')
                return 1

            with Capturing() as output:
                res = demo()
            assert '¥' in output[0]

            @cache_results(cache_fp)
            def demo():
                print('¥')
                return 1

            with Capturing() as output:
                res = demo()
            assert '¥' not in output[0]
        finally:
            rank_zero_rm('demo.pkl')

    def test_cache_func_change(self, capsys):
        cache_fp = 'demo.pkl'
        try:
            @cache_results(cache_fp)
            def demo():
                print('¥')
                return 1

            with Capturing() as output:
                res = demo()
            assert '¥' in output[0]

            @cache_results(cache_fp)
            def demo():
                print('¥¥')
                return 1

            with Capturing() as output:
                res = demo()
            assert 'different' in output[0]
            assert '¥' not in output[0]

            # 关闭check_hash应该不warning的
            with Capturing() as output:
                res = demo(_check_hash=0)
            assert 'different' not in output[0]
            assert '¥' not in output[0]

        finally:
            rank_zero_rm('demo.pkl')

    def test_cache_check_hash(self):
        cache_fp = 'demo.pkl'
        try:
            @cache_results(cache_fp, _check_hash=False)
            def demo():
                print('¥')
                return 1

            with Capturing() as output:
                res = demo(_check_hash=0)
            assert '¥' in output[0]

            @cache_results(cache_fp, _check_hash=False)
            def demo():
                print('¥¥')
                return 1

            # 默认不会check
            with Capturing() as output:
                res = demo()
            assert 'different' not in output[0]
            assert '¥' not in output[0]

            # check也可以
            with Capturing() as output:
                res = demo(_check_hash=True)
            assert 'different' in output[0]
            assert '¥' not in output[0]

        finally:
            rank_zero_rm('demo.pkl')

    # 外部 function 改变也会 导致改变
    def test_refer_fun_change(self):
        cache_fp = 'demo.pkl'
        test_type = 'func_refer_fun_change'
        try:
            cmd = f'python {__file__} --cache_fp {cache_fp} --test_type {test_type} --turn 0'
            res = get_subprocess_results(cmd)
            assert "¥" in res
            # 引用的function没有变化
            cmd = f'python {__file__} --cache_fp {cache_fp} --test_type {test_type} --turn 0'
            res = get_subprocess_results(cmd)
            assert "¥" not in res

            assert 'Read' in res
            assert 'different' not in res

            # 引用的function有变化
            cmd = f'python {__file__} --cache_fp {cache_fp} --test_type {test_type} --turn 1'
            res = get_subprocess_results(cmd)
            assert "¥" not in res
            assert 'different' in res

        finally:
            rank_zero_rm(cache_fp)

    # 外部 method 改变也会 导致改变
    def test_refer_class_method_change(self):
        cache_fp = 'demo.pkl'
        test_type = 'refer_class_method_change'
        try:
            cmd = f'python {__file__} --cache_fp {cache_fp} --test_type {test_type} --turn 0'
            res = get_subprocess_results(cmd)
            assert "¥" in res

            # 引用的class没有变化
            cmd = f'python {__file__} --cache_fp {cache_fp} --test_type {test_type} --turn 0'
            res = get_subprocess_results(cmd)
            assert 'Read' in res
            assert 'different' not in res
            assert "¥" not in res

            cmd = f'python {__file__} --cache_fp {cache_fp} --test_type {test_type} --turn 1'
            res = get_subprocess_results(cmd)
            assert 'different' in res
            assert "¥" not in res

        finally:
            rank_zero_rm(cache_fp)

    def test_duplicate_keyword(self):
        with pytest.raises(RuntimeError):
            @cache_results(None)
            def func_verbose(a, _verbose):
                pass

            func_verbose(0, 1)
        with pytest.raises(RuntimeError):
            @cache_results(None)
            def func_cache(a, _cache_fp):
                pass

            func_cache(1, 2)
        with pytest.raises(RuntimeError):
            @cache_results(None)
            def func_refresh(a, _refresh):
                pass

            func_refresh(1, 2)

        with pytest.raises(RuntimeError):
            @cache_results(None)
            def func_refresh(a, _check_hash):
                pass

            func_refresh(1, 2)

    def test_create_cache_dir(self):
        @cache_results('demo/demo.pkl')
        def cache():
            return 1, 2

        try:
            results = cache()
            assert (1, 2) == results
        finally:
            rank_zero_rm('demo/')

    def test_result_none_error(self):
        @cache_results('demo.pkl')
        def cache():
            pass

        try:
            with pytest.raises(RuntimeError):
                results = cache()
        finally:
            rank_zero_rm('demo.pkl')


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--test_type', type=str, default='refer_class_method_change')
    parser.add_argument('--turn', type=int, default=1)
    parser.add_argument('--cache_fp', type=str, default='demo.pkl')
    args = parser.parse_args()

    test_type = args.test_type
    cache_fp = args.cache_fp
    turn = args.turn

    if test_type == 'func_refer_fun_change':
        if turn == 0:
            def demo():
                b = 1
                return b
        else:
            def demo():
                b = 2
                return b

        @cache_results(cache_fp)
        def demo_refer_other_func():
            b = demo()
            print("¥")
            return b

        res = demo_refer_other_func()

    if test_type == 'refer_class_method_change':
        print(f"Turn:{turn}")
        if turn == 0:
            from helper_for_cache_results_1 import Demo
        else:
            from helper_for_cache_results_2 import Demo

        demo = Demo()
        # import pdb
        # pdb.set_trace()
        @cache_results(cache_fp)
        def demo_func():
            print("¥")
            b = demo.demo()
            return b

        res = demo_func()