"""py.py unit tests."""
import pytest
from pypyr.context import Context
from pypyr.dsl import PyString
from pypyr.errors import KeyInContextHasNoValueError, KeyNotInContextError
from pypyr.moduleloader import _ChainMapPretendDict
import pypyr.steps.py

# region py


def test_py_existing_key():
    """Py expression can update existing key."""
    context = Context({'x': 123, 'py': "x = abs(-42); save('x')"})
    assert context['x'] == 123
    pypyr.steps.py.run_step(context)
    # nothing added to context, but x is updated.
    assert context == {'x': 42, 'py': "x = abs(-42); save('x')"}


def test_py_with_import():
    """Py expression can use imports."""
    context = Context(
        {'y': 4, 'py': 'import math\nx = math.sqrt(y)\nsave("x")'})
    assert context['y'] == 4
    assert 'x' not in context
    pypyr.steps.py.run_step(context)
    assert context['x'] == 2
    assert context['y'] == 4


def test_py_single_code():
    """One word python function works."""
    context = Context({'py': 'abs(-1-2)'})
    pypyr.steps.py.run_step(context)


def test_py_sequence():
    """Sequence of py code works and touches context."""
    context = Context({'py': "test = 1; save('test')"})
    pypyr.steps.py.run_step(context)

    context.update({'py': "test += 2; save('test')"})
    pypyr.steps.py.run_step(context)

    context.update({'py': "test += 3; save('test')"})
    pypyr.steps.py.run_step(context)

    assert context['test'] == 6, "context should be 6 at this point"

    assert context == {'py': "test += 3; save('test')", 'test': 6}


def test_py_replace_objects():
    """Rebind mutable and immutable objects in exec scope."""
    pycode = """\
alist = [0, 1, 2]
adict = {'a': 'b', 'c': 'd'}
anint = 456
mutate_me.append(12)
astring = 'updated'
"""
    context = Context({'alist': [0, 1],
                       'adict': {'a': 'b'},
                       'anint': 123,
                       'mutate_me': [10, 11],
                       'astring': 'a string',
                       'py': pycode})

    pypyr.steps.py.run_step(context)

    # rebound vars not changed because no `save()`
    assert context == {'alist': [0, 1],
                       'adict': {'a': 'b'},
                       'anint': 123,
                       'mutate_me': [10, 11, 12],
                       'astring': 'a string',
                       'py': pycode}


def test_py_with_closure_scope():
    """Free variables scopes resolve to global."""
    pycode = """\
list1.append(1)
list2 = ['a', 'b', 'c']
out = [(x, y) for x in list1 for y in list2]
save('out')
"""
    context = Context({'list1': [0],
                       'list2': ['a', 'b'],
                       'py': pycode})

    pypyr.steps.py.run_step(context)

    assert context == {'list1': [0, 1],
                       'list2': ['a', 'b'],
                       'py': pycode,
                       'out': [(0, 'a'), (0, 'b'), (0, 'c'),
                               (1, 'a'), (1, 'b'), (1, 'c')]}


def test_py_save_inputs():
    """Special save takes args and kwargs."""
    pycode = """\
a = 1
b = 2
c = 3
d = 4
e = 5
f = 6
g = 7
h = 8
i = 9
j = 10
save('a')
save('b', 'c')
save(dd=d)
save(ee=e, ff=f+1)
save('g', 'h', ii=i, jj=j, k=5)
"""
    context = Context({'py': pycode})
    pypyr.steps.py.run_step(context)

    assert context == {'a': 1,
                       'b': 2,
                       'c': 3,
                       'dd': 4,
                       'ee': 5,
                       'ff': 7,
                       'g': 7,
                       'h': 8,
                       'ii': 9,
                       'jj': 10,
                       'k': 5,
                       'py': pycode}


def test_py_save_existing():
    """Save overwrites existing keys."""
    pycode = """\
alist = [0, 1, 2]
adict = {'a': 'b', 'c': 'd'}
anint = 456
mutate_me.append(12)
astring = 'updated'

save('alist', 'adict', 'anint', 'mutate_me', 'astring')
"""
    context = Context({'alist': [0, 1],
                       'adict': {'a': 'b'},
                       'anint': 123,
                       # strictly speaking not necessary to save mutate_me
                       # explicitly, coz it's mutated not rebound, but test
                       # that it works coz end-users mightn't see the
                       # distinction
                       'mutate_me': [10, 11],
                       'astring': 'a string',
                       'py': pycode})

    pypyr.steps.py.run_step(context)

    assert context == {'alist': [0, 1, 2],
                       'adict': {'a': 'b', 'c': 'd'},
                       'anint': 456,
                       'mutate_me': [10, 11, 12],
                       'astring': 'updated',
                       'py': pycode}


def test_py_save_key_error():
    """Save raises special message on KeyError."""
    pycode = "notfound='arb'; save(notfound)"
    context = Context({'py': pycode})

    with pytest.raises(KeyError) as err:
        pypyr.steps.py.run_step(context)

    # KeyError looks like it reprs on the __str__?
    expected = repr(
        "Trying to save 'arb', but can't find it in the py step scope. "
        "Remember it should be save('key'), not save(key) - mind the quotes.")

    assert str(err.value) == expected


def test_py_save_keyword_exists():
    """Save does not overwrite existing user 'save' key."""
    pycode = "k2='v2'; save('k2')"
    context = Context({'save': 123, 'py': pycode})

    pypyr.steps.py.run_step(context)

    assert context == {'save': 123, 'py': pycode, 'k2': 'v2'}


def test_py_imports_not_in_context():
    """Imported modules + non saves don't end up in context."""
    pycode = """\
from math import sqrt
b = sqrt(a)
"""
    context = Context({'a': 4,
                       'py': pycode})

    pypyr.steps.py.run_step(context)

    assert context == {'a': 4,
                       'py': pycode}


def test_py_sequence_with_semicolons():
    """Single py code string with semi - colons works."""
    context = Context({'py':
                       'x = abs(-1); x+= abs(-2); x += abs(-3); save("x")'})
    pypyr.steps.py.run_step(context)

    assert context == Context({
        'py': 'x = abs(-1); x+= abs(-2); x += abs(-3); save("x")',
        'x': 6})


def test_py_sequence_with_linefeeds():
    """Single py code string with linefeeds works."""
    context = Context({'py':
                       'abs(-1)\nabs(-2)\nabs(-3)'})
    pypyr.steps.py.run_step(context)


def test_py_error_throws():
    """Input pycode error should raise up to caller."""
    with pytest.raises(AssertionError):
        context = Context({'py': 'assert False'})
        pypyr.steps.py.run_step(context)


def test_py_no_context_throw():
    """No pycode in context should throw assert error."""
    with pytest.raises(KeyNotInContextError) as err_info:
        context = Context({'blah': 'blah blah'})
        pypyr.steps.py.run_step(context)

    assert str(err_info.value) == ("context['py'] "
                                   "doesn't exist. It must exist for "
                                   "pypyr.steps.py.")


def test_py_none_context_throw():
    """None pycode in context should throw assert error."""
    with pytest.raises(KeyInContextHasNoValueError):
        context = Context({'py': None})
        pypyr.steps.py.run_step(context)
# endregion py

# region pycode


def test_pycode_single_code():
    """One word python function works."""
    context = Context({'pycode': 'abs(-1-2)'})
    pypyr.steps.py.run_step(context)


def test_pycode_sequence():
    """Sequence of py code works and touches context."""
    context = Context({'pycode': "context['test'] = 1;"})
    pypyr.steps.py.run_step(context)

    context.update({'pycode': "context['test'] += 2"})
    pypyr.steps.py.run_step(context)

    context.update({'pycode': "context['test'] += 3"})
    pypyr.steps.py.run_step(context)

    assert context['test'] == 6, "context should be 6 at this point"

    assert context == {'pycode': "context['test'] += 3", 'test': 6}


def test_pycode_replace_objects():
    """Rebind mutable and immutable objects in eval scope."""
    pycode = """\
context['alist'] = [0, 1, 2]
context['adict'] = {'a': 'b', 'c': 'd'}
context['anint'] = 456
context['mutate_me'].append(12)
context['astring'] = 'updated'
"""
    context = Context({'alist': [0, 1],
                       'adict': {'a': 'b'},
                       'anint': 123,
                       'mutate_me': [10, 11],
                       'astring': 'a string',
                       'pycode': pycode})

    pypyr.steps.py.run_step(context)

    assert context == {'alist': [0, 1, 2],
                       'adict': {'a': 'b', 'c': 'd'},
                       'anint': 456,
                       'mutate_me': [10, 11, 12],
                       'astring': 'updated',
                       'pycode': pycode}


def test_pycode_with_closure_scope():
    """Free variables in closure scopes resolve to global."""
    pycode = """\
context['list1'].append(1)
context['list2'] = ['a', 'b', 'c']
context['out'] = [(x, y) for x in context['list1'] for y in context['list2']]
"""
    context = Context({'list1': [0],
                       'list2': ['a', 'b'],
                       'pycode': pycode})

    pypyr.steps.py.run_step(context)

    assert context == {'list1': [0, 1],
                       'list2': ['a', 'b', 'c'],
                       'pycode': pycode,
                       'out': [(0, 'a'), (0, 'b'), (0, 'c'),
                               (1, 'a'), (1, 'b'), (1, 'c')]}


def test_pycode_imports_not_in_context():
    """Imported modules don't end up in context."""
    pycode = """\
from math import sqrt
context['b'] = sqrt(context['a'])
"""
    context = Context({'a': 4,
                       'pycode': pycode})

    pypyr.steps.py.run_step(context)

    assert context == {'a': 4,
                       'b': 2,
                       'pycode': pycode}


def test_pycode_sequence_with_semicolons():
    """Single py code string with semi - colons works."""
    context = Context({'pycode':
                       'abs(-1); abs(-2); abs(-3);'})
    pypyr.steps.py.run_step(context)

    assert context == {'pycode':
                       'abs(-1); abs(-2); abs(-3);'}, ("context in and out "
                                                       "the same")


def test_pycode_sequence_with_linefeeds():