# -*- coding: future_annotations -*-
import functools
import logging

from nbsafety.data_model.code_cell import cells
from nbsafety.singletons import nbs
from test.utils import make_safety_fixture

logging.basicConfig(level=logging.ERROR)

# Reset dependency graph before each test
# _safety_fixture, run_cell_ = make_safety_fixture(trace_messages_enabled=True)
_safety_fixture, run_cell_ = make_safety_fixture()


def dynamic_or_static_only_test(test_f, dynamic_enabled, static_enabled):
    @functools.wraps(test_f)
    def dynamic_test_f(*args, **kwargs):
        orig_dynamic_enabled = nbs().mut_settings.dynamic_slicing_enabled
        orig_static_enabled = nbs().mut_settings.dynamic_slicing_enabled
        try:
            nbs().mut_settings.dynamic_slicing_enabled = dynamic_enabled
            nbs().mut_settings.static_slicing_enabled = static_enabled
            test_f(*args, **kwargs)
        finally:
            nbs().mut_settings.dynamic_slicing_enabled = orig_dynamic_enabled
            nbs().mut_settings.static_slicing_enabled = orig_static_enabled
    return dynamic_test_f


def dynamic_only_test(test_f):
    return dynamic_or_static_only_test(test_f, True, False)


def static_only_test(test_f):
    return dynamic_or_static_only_test(test_f, False, True)


def run_cell(cell):
    # print()
    # print('*******************************************')
    # print('running', cell)
    # print('*******************************************')
    # print()
    run_cell_(cell)


def compute_slice_stmts(cell_num):
    return cells().from_timestamp(cell_num).compute_slice_stmts()


def compute_slice(cell_num):
    return cells().from_timestamp(cell_num).compute_slice()


def num_stmts_in_slice(cell_num: int) -> int:
    return sum(len(stmts) for stmts in compute_slice_stmts(cell_num).values())


def test_simple():
    run_cell('a = 1')
    run_cell('b = 2')
    run_cell('logging.info(a)')
    run_cell('c = a + b')
    deps = set(compute_slice(4).keys())
    assert deps == {1, 2, 4}, 'got %s' % deps


def test_dynamic_symbol_usage():
    run_cell('x = 5')
    run_cell("""
class Foo:
    def foo(self):
        return x
""")
    run_cell("""
def foo():
    return Foo()
""")
    run_cell('logging.info(foo().foo())')
    deps = set(compute_slice(4).keys())
    assert deps == {1, 2, 3, 4}, 'got %s' % deps
    slice_size = num_stmts_in_slice(4)
    assert slice_size == 4, 'got %d' % slice_size


def test_nested_symbol_usage():
    run_cell('lst = [1, 2, 3, 4, 5]')
    run_cell('lst[1] = 3')
    run_cell('logging.info(lst[1])')
    deps = set(compute_slice(3).keys())
    assert deps == {1, 2, 3}, 'got %s' % deps
    slice_size = num_stmts_in_slice(3)
    assert slice_size == 3, 'got %d' % slice_size


def test_nested_symbol_usage_with_variable_subscript():
    for static_slicing_enabled in [True, False]:
        if not static_slicing_enabled and not nbs().mut_settings.dynamic_slicing_enabled:
            continue
        orig_enabled = nbs().mut_settings.static_slicing_enabled
        try:
            nbs().mut_settings.static_slicing_enabled = static_slicing_enabled
            run_cell('x = 1')
            run_cell('lst = [1, 2, 3, 4, 5]')
            run_cell('lst[x] = 3')
            run_cell('logging.info(lst[1])')
            deps = set(compute_slice(4).keys())
            assert deps == {1, 2, 3, 4}, 'got %s' % deps
            slice_size = num_stmts_in_slice(4)
            assert slice_size == 4, 'got %d' % slice_size
        finally:
            nbs().mut_settings.static_slicing_enabled = orig_enabled


@static_only_test
def test_liveness_timestamps():
    run_cell("""
x = 0
if True:
    y = 1
else:
    y = 2
""")
    run_cell('z = 42')
    run_cell('logging.info(x + 1)')
    deps = set(compute_slice(3).keys())
    assert deps == {1, 3}, 'got %s' % deps
    slice_size = num_stmts_in_slice(3)
    assert slice_size == 2, 'got %d' % slice_size


def test_list_mutations():
    run_cell('dummy = 0')
    run_cell('lst = []')
    run_cell('lst.append(1)')
    run_cell('lst.append(2)')
    run_cell('lst.append(3); lst.append(4)')
    run_cell('logging.info(lst)')
    deps = set(compute_slice(5).keys())
    assert deps == {2, 3, 4, 5}, 'got %s' % deps
    slice_size = num_stmts_in_slice(5)
    assert slice_size == 5, 'got %d' % slice_size


def test_imports():
    run_cell('import numpy as np')
    run_cell('arr = np.zeros((5,))')
    run_cell('logging.info(arr * 3)')
    deps = set(compute_slice(3).keys())
    assert deps == {1, 2, 3}, 'got %s' % deps
    slice_size = num_stmts_in_slice(3)
    assert slice_size == 3, 'got %d' % slice_size


def test_handle_stale():
    run_cell('a = 1')
    run_cell('b = 2 * a')
    run_cell('a = 2')
    run_cell('logging.info(b)')
    run_cell('logging.info(b)')
    deps = set(compute_slice(4).keys())
    assert deps == {1, 2, 4}, 'got %s' % deps
    slice_size = num_stmts_in_slice(4)
    assert slice_size == 3, 'got %d' % slice_size


@dynamic_only_test
def test_multiple_versions_captured():
    run_cell('x = 0')
    run_cell('logging.info(x); y = 7')
    run_cell('x = 5')
    run_cell('logging.info(x + y)')
    deps = set(compute_slice(4).keys())
    assert deps == {1, 2, 3, 4}, 'got %s' % deps
    slice_size = num_stmts_in_slice(4)
    assert slice_size == 3, 'got %d' % slice_size


def test_version_used_when_live():
    run_cell('x = 0')
    run_cell("""
if True:
    y = 7
else:
    # even though this branch is not taken,
    # liveness-based usage should detect the
    # version of `x` used at the time it was
    # live, meaning cell 1 should get included
    # in the slice
    logging.info(x)
""")
    run_cell('x = 5')
    run_cell('logging.info(x + y)')
    deps = set(compute_slice(4).keys())
    assert deps == {1, 2, 3, 4}, 'got %s' % deps
    slice_size = num_stmts_in_slice(4)
    assert slice_size == 4, 'got %d' % slice_size


def test_no_definitely_spurious_cells():
    run_cell('x = 0')
    run_cell("""
if True:
    y = 7
else:
    # even though this branch is not taken,
    # liveness-based usage should detect the
    # version of `x` used at the time it was
    # live, meaning cell 1 should get included
    # in the slice
    logging.info(x)
""")
    run_cell('x = 5')
    run_cell('logging.info(y)')
    deps = set(compute_slice(4).keys())
    assert deps == {1, 2, 4}, 'got %s' % deps
    slice_size = num_stmts_in_slice(4)
    assert slice_size == 3, 'got %d' % slice_size


def test_parent_usage_includes_child_update():
    run_cell('lst = [3]')
    run_cell('lst[0] += 1')
    run_cell('lst2 = lst + [5]')
    deps = set(compute_slice(3).keys())
    assert deps == {1, 2, 3}, 'got %s' % deps
    slice_size = num_stmts_in_slice(3)
    assert slice_size == 3, 'got %d' % slice_size


def test_object_subscripting():
    run_cell("""
class Foo:
    def __init__(self):
        self.counter = 0
        
    def inc(self):
        self.counter += 1

    def bar(self):
        return {'baz': 0}
""")
    run_cell("obj = Foo()")
    run_cell("name = 'baz'")
    run_cell("obj.bar()[name]")
    run_cell("something_i_dont_care_about = obj.bar()[name]")
    run_cell("something_i_care_about = obj.bar()[name]")
    deps = set(compute_slice(6).keys())
    assert deps == {1, 2, 3, 6}, 'got %s' % deps
    slice_size = num_stmts_in_slice(6)
    assert slice_size == 4, 'got %d' % slice_size


def test_complicated_subscripting():
    run_cell("""
class Foo:
    def __init__(self):
        self.counter = 0
        
    def inc(self):
        self.counter += 1

    def bar(self, indicator):
        if indicator == 1:
            return the_dictionary
        else:
            return x + 5
    
    def new(self, indicator):
        if indicator == 1:
            return Foo()
        
        return Bar()
""")
    run_cell("""
class Bar:
    def __init__(self):
        self.counter = 0
    
    def foo(self, indicator):
        if indicator == 1:
            return the_dictionary
        
        return Foo()
""")
    run_cell("x = 1")
    run_cell("the_dictionary = {'something': 1}")
    run_cell("Foo().new(0).foo(1)")
    deps = set(compute_slice(5).keys())
    assert deps == {1, 2, 4, 5}, 'got %s' % deps
    slice_size = num_stmts_in_slice(5)
    assert slice_size == 4, 'got %d' % slice_size


def test_complicated_subscripting_use_conditional():
    run_cell("""
class Foo:
    def __init__(self):
        self.counter = 0
        
    def inc(self):
        self.counter += 1

    def bar(self, indicator):
        if indicator == 1:
            return the_dictionary
        else:
            return x + 5
    
    def new(self, indicator):
        if indicator == 1:
            return Foo()
        
        return Bar()
""")
    run_cell("""
class Bar:
    def __init__(self):
        self.counter = 0
    
    def foo(self, indicator):
        if indicator == 1:
            return the_dictionary
        
        return Foo()
""")
    run_cell("x = 1")
    run_cell("the_dictionary = {'something': 1}")
    run_cell("""
y = Foo().new(0).foo(1)
z = Foo().new(1).bar(0)
""")
    run_cell("logging.info(z)")
    run_cell("logging.info(y)")
    deps = set(compute_slice(6).keys())
    assert deps == {1, 2, 3, 4, 5, 6}, 'got %s' % deps
    slice_size = num_stmts_in_slice(6)
    assert slice_size == 4, 'got %d' % slice_size
    slice_size = num_stmts_in_slice(7)
    assert slice_size == 5, 'got %d' % slice_size


def test_non_relevant_child_symbol_modified():
    run_cell('lst = [0, 1, 2]')
    run_cell('lst[0] += 1')
    run_cell('lst[0] += 1')
    run_cell('lst[1] += 1')
    run_cell('lst[0] += 1')
    run_cell('lst[0] += 1')
    run_cell('logging.info(lst[1])')
    run_cell('lst[1] = 42')
    deps = set(compute_slice(7).keys())
    assert deps == {1, 4, 7}, 'got %s' % deps
    deps = set(compute_slice(8).keys())
    assert deps == {1, 8}, 'got %s' % deps
    slice_size = num_stmts_in_slice(7)
    assert slice_size == 3, 'got %d' % slice_size


@dynamic_only_test
def test_dynamic_only_increment():
    nbs().mut_settings.static_slicing_enabled = False
    run_cell('x = 0')
    run_cell('x += 1')
    run_cell('logging.info(x)')
    deps = set(compute_slice(3).keys())
    assert deps == {1, 2, 3}, 'got %s' % deps
    slice_size = num_stmts_in_slice(3)
    assert slice_size == 3, 'got %d' % slice_size


@dynamic_only_test
def test_dynamic_only_variable_subscript():
    nbs().mut_settings.static_slicing_enabled = False
    run_cell('lst = [0, 1, 2]')
    run_cell('x = 0')
    run_cell('lst[x] += 1')
    run_cell('lst[x] += 1')
    run_cell('x += 1')
    run_cell('lst[x] += 1')
    run_cell('x -= 1')
    run_cell('lst[x] += 1')
    run_cell('lst[x] += 1')
    run_cell('x += 1')
    run_cell('logging.info(lst[x])')
    deps = set(compute_slice(11).keys())
    assert deps == {1, 2, 5, 6, 7, 10, 11}, 'got %s' % deps
    slice_size = num_stmts_in_slice(11)
    assert slice_size == len(deps), 'got %d' % slice_size


@dynamic_only_test
def test_handler():
    run_cell("""
try:
    r = map()
except:
    success = False""")
    run_cell("""
try:
    logging.info("%d %d", (1, 2))
except:
    success = False""")
    deps = set(compute_slice(2).keys())
    assert deps == {2}, 'got %s' % deps
    slice_size = num_stmts_in_slice(2)
    assert slice_size == 1, 'got %d' % slice_size


@dynamic_only_test
def test_list_delete():
    run_cell('lst = [0, 1, 2, 3, 4, 5, 6]')
    run_cell('lst.append(7)')
    run_cell('del lst[0]')
    run_cell('del lst[2]')
    run_cell('logging.info(lst[1])')
    deps = set(compute_slice(5).keys())
    assert deps == {1, 2, 3, 5}, 'got %s' % deps
    slice_size = num_stmts_in_slice(5)
    assert slice_size == len(deps), 'got %d' % slice_size
