# -*- coding: future_annotations -*-
from contextlib import contextmanager
import logging
from typing import TYPE_CHECKING

from nbsafety.safety import NotebookSafetySettings
from nbsafety.singletons import nbs
from test.utils import make_safety_fixture, skipif_known_failing

if TYPE_CHECKING:
    from typing import Dict

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()


@contextmanager
def override_settings(**kwargs):
    old_settings = nbs().settings
    new_settings = old_settings._asdict()
    new_settings.update(kwargs)
    new_settings = NotebookSafetySettings(**new_settings)
    try:
        nbs().settings = new_settings
        yield
    finally:
        nbs().settings = old_settings


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


def run_all_cells(cells: Dict[int, str], **kwargs):
    for cell_id in sorted(cells.keys()):
        run_cell(cells[cell_id], cell_id=cell_id, **kwargs)


def test_simple():
    cells = {
        0: 'x = 0',
        1: 'y = x + 1',
        2: 'logging.info(y)',
        3: 'x = 42',
    }
    run_all_cells(cells)
    response = nbs().check_and_link_multiple_cells()
    assert response.stale_cells == {2}
    assert response.fresh_cells == {1}
    assert response.stale_links == {2: {1}}
    assert response.refresher_links == {1: {2}}


def test_refresh_after_exception_fixed():
    cells = {
        0: 'x = 0',
        1: 'y = x + 1',
        2: 'logging.info(y)',
    }
    run_cell(cells[0], 0)
    run_cell(cells[2], 2, ignore_exceptions=True)
    run_cell(cells[1], 1)
    response = nbs().check_and_link_multiple_cells()
    assert response.fresh_cells == {2}


def test_refresh_after_val_changed():
    cells = {
        0: 'x = 0',
        1: 'y = x + 1',
        2: 'logging.info(y)',
        3: 'y = 42',
    }
    run_all_cells(cells)
    response = nbs().check_and_link_multiple_cells()
    assert response.fresh_cells == {2}


def test_inner_mutation_considered_fresh():
    cells = {
        0: 'lst_0 = [0,1,2]',
        1: 'lst_1 = [3,4,5]',
        2: 'lst = [lst_0, lst_1]',
        3: 'logging.info(lst)',
        4: 'lst_0.append(42)',
    }
    run_all_cells(cells)
    response = nbs().check_and_link_multiple_cells()
    assert response.stale_cells == set()
    assert response.fresh_cells == {2, 3}


# @pytest.mark.parametrize("force_subscript_symbol_creation", [True, False])
def test_update_list_elem():
    force_subscript_symbol_creation = True
    cells = {
        0: """
class Foo:
    def __init__(self):
        self.counter = 0
        self.dummy = 0
        
    def inc(self):
        self.counter += 1""",

        1: """
lst = []
for i in range(5):
    x = Foo()
    lst.append(x)""",

        2: """
for foo in lst:
    foo.inc()""",

        3: 'logging.info(lst)',
    }

    run_all_cells(cells)

    response = nbs().check_and_link_multiple_cells()
    assert response.stale_cells == set()
    assert response.fresh_cells == set(), 'got %s' % response.fresh_cells

    cells[4] = 'x.inc()'
    run_cell(cells[4], 4)

    response = nbs().check_and_link_multiple_cells()
    assert response.stale_cells == set(), 'got %s' % response.stale_cells
    assert response.fresh_cells == {2, 3}, 'got %s' % response.fresh_cells

    cells[5] = 'foo.inc()'
    run_cell(cells[5], 5)
    response = nbs().check_and_link_multiple_cells()
    assert response.stale_cells == set(), 'got %s' % response.stale_cells
    assert response.fresh_cells == {2, 3, 4}, 'got %s' % response.fresh_cells

    if force_subscript_symbol_creation:
        cells[6] = 'lst[-1]'
        run_cell(cells[6], 6)
        response = nbs().check_and_link_multiple_cells()
        assert response.stale_cells == set()
        assert response.fresh_cells == {2, 3, 4}

    run_cell(cells[4], 4)
    response = nbs().check_and_link_multiple_cells()
    assert response.stale_cells == set()
    assert response.fresh_cells == set([2, 3, 5] + ([6] if force_subscript_symbol_creation else []))


def test_no_freshness_for_alias_assignment_post_mutation():
    cells = {
        0: 'x = []',
        1: 'y = x',
        2: 'x.append(5)',
    }
    run_all_cells(cells)
    response = nbs().check_and_link_multiple_cells()
    assert response.stale_cells == set()
    assert response.fresh_cells == set()


def test_fresh_after_import():
    cells = {
        0: 'x = np.random.random(10)',
        1: 'import numpy as np'
    }
    run_all_cells(cells, ignore_exceptions=True)
    response = nbs().check_and_link_multiple_cells()
    assert response.stale_cells == set()
    assert response.fresh_cells == {0}


def test_external_object_update_propagates_to_stale_namespace_symbols():
    cells = {
        0: 'import fakelib',
        1: 'foo = fakelib.Foo()',
        2: 'logging.info(foo.x)',
        3: 'x = 42',
        4: 'foo.x = x + 1',
        5: 'x = 43',
        6: 'foo = foo.set_x(10)',
    }
    with override_settings(mark_stale_symbol_usages_unsafe=False):
        run_all_cells(cells)
        response = nbs().check_and_link_multiple_cells()
        assert response.stale_cells == set(), 'got %s' % response.stale_cells
        assert response.fresh_cells == {2, 4}


def test_symbol_on_both_sides_of_assignment():
    cells = {
        0: 'x = 0',
        1: 'y = x + 1',
        2: 'y += 7',
        3: 'x = 42',
    }
    run_all_cells(cells)
    response = nbs().check_and_link_multiple_cells()
    assert response.stale_cells == {2}
    assert response.fresh_cells == {1}
    assert list(response.refresher_links.keys()) == [1]


def test_updated_namespace_after_subscript_dep_removed():
    cells = {
        0: 'x = 5',
        1: 'd = {x: 5}',
        2: 'logging.info(d[5])',
        3: 'x = 9',
    }
    run_all_cells(cells)
    response = nbs().check_and_link_multiple_cells()
    assert response.stale_cells == {2}
    assert response.fresh_cells == {1}
    cells[1] = 'd = {5: 6}'
    run_cell(cells[1], 1)
    response = nbs().check_and_link_multiple_cells()
    assert response.stale_cells == set()
    assert response.fresh_cells == {2}
    run_cell(cells[2], 2)
    response = nbs().check_and_link_multiple_cells()
    assert response.stale_cells == set()
    assert response.fresh_cells == set()
    run_cell(cells[0], 0)
    response = nbs().check_and_link_multiple_cells()
    assert response.stale_cells == set()
    assert response.fresh_cells == set(), 'got %s' % response.fresh_cells


def test_equal_list_update_does_induce_fresh_cell():
    cells = {
        0: 'x = ["f"] + ["o"] * 10',
        1: 'y = x + list("bar")',
        2: 'logging.info(y)',
        3: 'y = list("".join(y))',
    }
    run_all_cells(cells)
    response = nbs().check_and_link_multiple_cells()
    assert response.stale_cells == set()
    assert response.fresh_cells == {2}
    run_cell('y = ("f",)', 4)
    response = nbs().check_and_link_multiple_cells()
    assert response.stale_cells == set()
    assert response.fresh_cells == {2, 3}


def test_equal_list_update_does_induce_fresh_cell_LITERAL_WITH_F_IS_REUSED_ON_UBUNTU_20_04_2_PYTHON_3_8_11():
    cells = {
        0: 'x = ["f"] + ["o"] * 10',
        1: 'y = x + list("bar")',
        2: 'logging.info(y)',
        3: 'y = list("".join(y))',
    }
    run_all_cells(cells)
    response = nbs().check_and_link_multiple_cells()
    assert response.stale_cells == set()
    assert response.fresh_cells == {2}
    run_cell('y = ["f"]', 4)
    response = nbs().check_and_link_multiple_cells()
    assert response.stale_cells == set()
    assert response.fresh_cells == {2, 3}


def test_equal_dict_update_does_induce_fresh_cell():
    cells = {
        0: 'x = {"foo": 42, "bar": 43}',
        1: 'y = dict(set(x.items()) | set({"baz": 44}.items()))',
        2: 'logging.info(y)',
        3: 'y = dict(y.items())',
    }
    run_all_cells(cells)
    response = nbs().check_and_link_multiple_cells()
    assert response.stale_cells == set()
    assert response.fresh_cells == {2}, 'got %s' % response.fresh_cells
    run_cell('y = {"foo": 99}', 4)
    response = nbs().check_and_link_multiple_cells()
    assert response.stale_cells == set()
    assert response.fresh_cells == {2, 3}


def test_list_append():
    cells = {
        0: 'lst = [0, 1]',
        1: 'x = lst[1] + 1',
        2: 'logging.info(x)',
        3: 'lst.append(2)',
    }
    run_all_cells(cells)
    response = nbs().check_and_link_multiple_cells()
    assert response.stale_cells == set()
    assert response.fresh_cells == set()
    run_cell('lst[1] += 42', 4)
    response = nbs().check_and_link_multiple_cells()
    assert response.stale_cells == {2}
    assert response.fresh_cells == {1}


def test_list_extend():
    cells = {
        0: 'lst = [0, 1]',
        1: 'x = lst[1] + 1',
        2: 'logging.info(x)',
        3: 'lst.extend([2, 3, 4])',
    }
    run_all_cells(cells)
    response = nbs().check_and_link_multiple_cells()
    assert response.stale_cells == set()
    assert response.fresh_cells == set()
    run_cell('lst[1] += 42', 4)
    response = nbs().check_and_link_multiple_cells()
    assert response.stale_cells == {2}
    assert response.fresh_cells == {1}


def test_implicit_subscript_symbol_does_not_bump_ts():
    cells = {
        0: 'lst = [] + [0, 1]',
        1: 'logging.info(lst)',
        2: 'logging.info(lst[0])',
    }
    run_all_cells(cells)
    response = nbs().check_and_link_multiple_cells()
    assert response.stale_cells == set()
    assert response.fresh_cells == set()


def test_liveness_skipped_for_simple_assignment_involving_aliases():
    cells = {
        0: 'lst = [1, 2, 3]',
        1: 'lst2 = lst',
        2: 'lst.append(4)',
    }
    run_all_cells(cells)
    response = nbs().check_and_link_multiple_cells()
    assert response.stale_cells == set()
    assert response.fresh_cells == set()
    run_cell('lst = [1, 2, 3, 4]', 3)
    response = nbs().check_and_link_multiple_cells()
    assert response.stale_cells == set()
    assert response.fresh_cells == {1, 2}, 'got %s' % response.fresh_cells


def test_incorrect_object_not_used_for_argument_symbols():
    cells = {
        0: 'import numpy as np',
        1: 'arr = np.random.randn(100)',
        2: 'def f(x): return []',
        # when obj for `np` was passed to `x`, this created issues
        3: 'f(np.arange(10))',
    }
    run_all_cells(cells)
    response = nbs().check_and_link_multiple_cells()
    assert response.stale_cells == set()
    assert response.fresh_cells == set(), 'got %s' % response.fresh_cells


def test_increment_by_same_amount():
    cells = {
        0: 'x = 2',
        1: 'y = x + 1',
        2: 'logging.info(y)',
    }
    run_all_cells(cells)
    response = nbs().check_and_link_multiple_cells()
    assert response.stale_cells == set()
    assert response.fresh_cells == set()
    run_cell('x = 3', 0)
    response = nbs().check_and_link_multiple_cells()
    assert response.stale_cells == {2}
    assert response.fresh_cells == {1}


def test_list_insert():
    cells = {
        0: 'lst = [0, 1, 2, 4, 5, 6]',
        1: 'logging.info(lst[0])',
        2: 'logging.info(lst[1])',
        3: 'logging.info(lst[2])',
        4: 'logging.info(lst[3])',
        5: 'logging.info(lst[4])',
        6: 'logging.info(lst[5])',
        7: 'x = lst[5] + 42',
        8: 'logging.info(x)',
        9: 'lst.insert(3, 3)',
    }
    run_all_cells(cells)
    response = nbs().check_and_link_multiple_cells()
    assert response.stale_cells == {8}, 'got %s' % response.stale_cells
    assert response.fresh_cells == {4, 5, 6, 7}, 'got %s' % response.fresh_cells


def _test_list_delete_helper(last_cell):
    cells = {
        0: 'lst = [0, 1, 2, 3, 3, 4, 5, 6]',
        1: 'logging.info(lst[0])',
        2: 'logging.info(lst[1])',
        3: 'logging.info(lst[2])',
        4: 'logging.info(lst[3])',
        5: 'logging.info(lst[4])',
        6: 'logging.info(lst[5])',
        7: 'logging.info(lst[6])',
        8: 'logging.info(lst[7])',
        9: 'x = lst[6] + 42',
        10: 'logging.info(x)',
        11: last_cell,
    }
    run_all_cells(cells)
    response = nbs().check_and_link_multiple_cells()
    assert response.stale_cells == {10}
    # TODO: ideally we would detect that lst[3] is the same after deleting
    #  and not consider cell 4 to be fresh
    # TODO: ideally cell 8 would be considered unsafe since the last entry
    #  no longer exists
    assert response.fresh_cells == {4, 5, 6, 7, 9}, 'got %s' % response.fresh_cells


def test_list_delete():
    _test_list_delete_helper('del lst[3]')


def test_list_pop():
    _test_list_delete_helper('lst.pop(3)')


def test_list_remove():
    _test_list_delete_helper('lst.remove(3)')


def test_list_clear():
    cells = {
        0: 'lst = [0]',
        1: 'logging.info(lst[0])',
        2: 'logging.info(lst)',
    }
    run_all_cells(cells)
    response = nbs().check_and_link_multiple_cells()
    assert response.stale_cells == set()
    assert response.fresh_cells == set()
    run_cell('lst.clear()', 3)
    response = nbs().check_and_link_multiple_cells()
    assert response.stale_cells == set()
    assert response.fresh_cells == {2}


def test_dict_clear():
    cells = {
        0: 'd = {0: 0}',
        1: 'logging.info(d[0])',
        2: 'logging.info(d)',
    }
    run_all_cells(cells)
    response = nbs().check_and_link_multiple_cells()
    assert response.stale_cells == set()
    assert response.fresh_cells == set()
    run_cell('d.clear()', 3)
    response = nbs().check_and_link_multiple_cells()
    assert response.stale_cells == set()
    assert response.fresh_cells == {2}


def test_adhoc_pandas_series_update():
    cells = {
        0: 'import pandas as pd',
        1: 'df = pd.DataFrame({1: [2, 3], 4: [5, 7]})',
        2: 'df["foo"] = [8, 9]',
        3: 'df.foo.dropna(inplace=True)',
    }
    run_all_cells(cells)
    response = nbs().check_and_link_multiple_cells()
    assert response.stale_cells == set()
    assert response.fresh_cells == set()
    run_cell('df["foo"] = [8, 9]', 4)
    response = nbs().check_and_link_multiple_cells()
    assert response.stale_cells == set()
    assert response.fresh_cells == {3}
