import pprint
import string
from typing import List
from unittest.mock import patch

import pytest
from hypothesis import given
from hypothesis.strategies import (
    lists,
    text,
    integers,
    floats,
    one_of,
    booleans,
    none,
    dictionaries,
)

import transformer.python as py
from transformer.builders_python import (
    indent_levels,
    ascii_inline_text,
    lines,
    opaque_blocks,
    statements,
    expressions,
    functions,
    decorations,
    classes,
    standalones,
    literals,
    function_calls,
    binary_ops,
    symbols,
    assignments,
    ifelses,
    imports,
)


class TestLine:
    def test_str_is_identity_with_indent_level_zero(self):
        assert str(py.Line("abc", 0)) == "abc"

    def test_default_indent_level_is_zero(self):
        assert str(py.Line("abc")) == "abc"

    def test_default_indent_unit_is_four_spaces(self):
        assert py.Line.INDENT_UNIT == "    "

    def test_str_indents_as_much_indent_units_as_provided_indent_level(self):
        assert (
            str(py.Line("abc", 1)) == py.Line.INDENT_UNIT + "abc"
        ), "indent level 1 means one time INDENT_UNIT"
        assert str(py.Line("x", 2)) == py.Line.INDENT_UNIT * 2 + "x"

    @given(lines, lines)
    def test_equal_iff_text_and_indent_are_equal(self, a: py.Line, b: py.Line):
        assert a != b or (a.text == b.text and a.indent_level == b.indent_level)

    def test_repr(self):
        line = py.Line(text="a'\" b", indent_level=3)
        assert repr(line) == "Line(text='a\\'\" b', indent_level=3)"


class TestStatement:
    def test_lines_must_be_implemented(self):
        with pytest.raises(NotImplementedError):
            py.Statement().lines()

    @given(statements, indent_levels, lists(ascii_inline_text(min_size=1), max_size=2))
    def test_comment_lines(self, stmt: py.Statement, level: int, comments: List[str]):
        stmt.comments = comments
        x = py.Line.INDENT_UNIT
        assert [str(l) for l in stmt.comment_lines(level)] == [
            x * level + f"# {line}" for line in comments
        ]

    @given(statements, lines)
    def test_attach_comment_without_comment_changes_nothing(
        self, stmt: py.Statement, line: py.Line
    ):
        clone = line.clone()
        stmt.comments.clear()  # Some comments may have been generated by Hypothesis.
        lines = stmt.attach_comment(line)
        assert len(lines) == 1, "no additional lines must be created"
        l = lines[0]
        assert l is line, "the same Line object must be returned"
        assert l.text == clone.text, "Line.text must not change"
        assert l.indent_level == clone.indent_level, "Line.indent_level must not change"

    @given(statements, lines, ascii_inline_text(min_size=1))
    def test_attach_comment_with_one_line_comment_changes_only_text(
        self, stmt: py.Statement, line: py.Line, comment: str
    ):
        clone = line.clone()
        stmt.comments = [comment]
        lines = stmt.attach_comment(line)
        assert len(lines) == 1, "no additional lines must be created"
        l = lines[0]
        assert l is line, "the same Line object must be returned"
        assert (
            l.text == f"{clone.text}  # {comment}"
        ), "the comment is appended to Line.text"
        assert l.indent_level == clone.indent_level, "Line.indent_level must not change"

    @given(statements, lines, lists(ascii_inline_text(min_size=1), min_size=2))
    def test_attach_comment_with_multiline_comment_adds_lines_above(
        self, stmt: py.Statement, line: py.Line, comments: List[str]
    ):
        clone = line.clone()
        stmt.comments = comments
        lines = stmt.attach_comment(line)
        assert len(lines) == 1 + len(comments)

        l = lines[-1]
        assert l is line, "the same Line object must be returned last"
        assert l.text == clone.text, "Line.text must not change"
        assert l.indent_level == clone.indent_level, "Line.indent_level must not change"

        assert lines[:-1] == [
            py.Line(f"# {s}", clone.indent_level) for s in comments
        ], "all lines but the last are standalone comment lines"


class TestOpaqueBlock:
    @given(opaque_blocks, opaque_blocks)
    def test_equal_iff_components_are_equal(self, a: py.OpaqueBlock, b: py.OpaqueBlock):
        assert a != b or (a.block == b.block and a.comments == b.comments)

    @given(text(string.whitespace, max_size=5))
    def test_lines_raises_for_empty_input_block(self, block: str):
        with pytest.raises(ValueError):
            py.OpaqueBlock(block)

    def test_lines_returns_block_lines_if_top_and_bottom_are_not_empty(self):
        ob = py.OpaqueBlock(" a\n  b\n\n\n c")
        assert len(ob.lines()) == 5

    def test_lines_returns_block_lines_without_empty_top_and_bottom(self):
        ob = py.OpaqueBlock("\n\n a\n  b\n\n\n c\n\n\n")
        assert len(ob.lines()) == 5

    X = py.Line.INDENT_UNIT

    @pytest.mark.parametrize(
        "input_block, indent_level, expected",
        [
            ("x", 0, "x"),
            ("x", 1, " x"),
            (" x", 1, " x"),
            (" x", 2, "  x"),
            (" x", 0, "x"),
            ("x\nx", 0, "x\nx"),
            ("x\nx", 1, " x\n x"),
            (" x\n x", 1, " x\n x"),
            ("x\n x", 1, " x\n  x"),
            ("x\nx\n", 1, " x\n x"),
            (" x\n x\n", 1, " x\n x"),
            ("x\n x\n", 1, " x\n  x"),
            ("x\n x\n  x", 0, "x\n x\n  x"),
            ("x\n x\n  x", 1, " x\n  x\n   x"),
            ("\nx\n x", 0, "x\n x"),
            ("\nx\n x", 1, " x\n  x"),
            ("\n x\n  x", 0, "x\n x"),
            ("\n x\n  x", 1, " x\n  x"),
            ("\tx\n\t x", 1, " x\n  x"),
            ("\tx\n\t\tx", 1, f" x\n {' ' * py.OpaqueBlock.TAB_SIZE}x"),
        ],
    )
    def test_lines_indents_correctly(
        self, input_block: str, indent_level: int, expected: str
    ):
        lines = py.OpaqueBlock(input_block).lines(indent_level)
        print("lines =")
        pprint.pprint(lines)
        with patch("transformer.python.Line.INDENT_UNIT", " "):
            assert "\n".join(str(line) for line in lines) == expected

    @given(opaque_blocks, indent_levels, ascii_inline_text(min_size=1))
    def test_lines_displays_comment_always_above(
        self, ob: py.OpaqueBlock, level: int, comment: str
    ):
        x = py.Line.INDENT_UNIT
        ob.comments = [comment]
        assert [str(l) for l in ob.lines(level)] == [
            x * level + f"# {comment}",
            *[str(l) for l in ob.lines(level, comments=False)],
        ]

    @given(indent_levels)
    def test_lines_with_hidden_comments(self, level: int):
        x = py.Line.INDENT_UNIT
        ob = py.OpaqueBlock("hello", comments=["1", "2"])
        assert [str(l) for l in ob.lines(level, comments=False)] == [
            x * level + "hello"
        ]

    def test_repr(self):
        text = " a'\" b "
        assert (
            repr(py.OpaqueBlock(block=text, comments=["hi"]))
            == "OpaqueBlock(' a\\'\" b ', comments=['hi'])"
        )


class TestFunction:
    @given(indent_levels)
    def test_lines_with_no_params_and_no_body(self, level: int):
        f = py.Function("f", params=[], statements=[])
        x = py.Line.INDENT_UNIT
        assert [str(l) for l in f.lines(level)] == [
            x * level + "def f():",
            x * (level + 1) + "pass",
        ]

    @given(indent_levels)
    def test_lines_with_simple_body(self, level: int):
        f = py.Function(
            "f",
            params=[],
            statements=[py.OpaqueBlock("print('Hello!')"), py.OpaqueBlock("return")],
        )
        x = py.Line.INDENT_UNIT
        assert [str(l) for l in f.lines(level)] == [
            x * level + "def f():",
            x * (level + 1) + "print('Hello!')",
            x * (level + 1) + "return",
        ]

    def test_lines_with_simple_params(self):
        f = py.Function("f", params=["x", "y"], statements=[])
        assert [str(l) for l in f.lines()] == [
            "def f(x, y):",
            py.Line.INDENT_UNIT + "pass",
        ]

    def test_lines_with_complex_params(self):
        f = py.Function(
            "f", params=["x: int", "abc: bool = True", "*z: str"], statements=[]
        )
        assert [str(l) for l in f.lines()] == [
            "def f(x: int, abc: bool = True, *z: str):",
            py.Line.INDENT_UNIT + "pass",
        ]

    @given(indent_levels)
    def test_lines_with_nested_body(self, level: int):
        x = py.Line.INDENT_UNIT
        f = py.Function(
            "func",
            params=[],
            statements=[
                py.Assignment("a", py.Literal(2)),
                py.IfElse(
                    [(py.Literal(True), [py.Assignment("b", py.Literal(3))])],
                    [
                        py.Assignment("b", py.Literal(4)),
                        py.Assignment("c", py.Literal(1)),
                    ],
                ),
            ],
        )
        assert [str(l) for l in f.lines(level)] == [
            x * level + "def func():",
            x * (level + 1) + "a = 2",
            x * (level + 1) + "if True:",
            x * (level + 2) + "b = 3",
            x * (level + 1) + "else:",
            x * (level + 2) + "b = 4",
            x * (level + 2) + "c = 1",
        ]

    @given(indent_levels)
    def test_lines_with_comments(self, level: int):
        stmt = py.OpaqueBlock("foo", comments=["x"])
        f = py.Function("f", params=[], statements=[stmt], comments=["1", "2"])
        x = py.Line.INDENT_UNIT
        assert [str(l) for l in f.lines(level)] == [
            x * level + "# 1",
            x * level + "# 2",
            x * level + "def f():",
            *[str(l) for l in stmt.lines(level + 1)],
        ]

    @given(indent_levels)
    def test_lines_with_hidden_comments(self, level: int):
        stmt = py.OpaqueBlock("foo", comments=["x"])
        f = py.Function("f", params=[], statements=[stmt], comments=["1", "2"])
        x = py.Line.INDENT_UNIT
        assert [str(l) for l in f.lines(level, comments=False)] == [
            x * level + "def f():",
            *[str(l) for l in stmt.lines(level + 1, comments=False)],
        ]

    def test_repr(self):
        stmts = [py.OpaqueBlock("raise")]
        assert (
            repr(py.Function(name="f", params=["a"], statements=stmts, comments=["hi"]))
            == f"Function(name='f', params=['a'], statements={stmts!r}, comments=['hi'])"
        )

    @given(functions, functions)
    def test_equal_iff_components_are_equal(self, a: py.Function, b: py.Function):
        assert a != b or (
            a.name == b.name
            and a.params == b.params
            and a.statements == b.statements
            and a.comments == b.comments
        )


class TestDecoration:
    @given(decorations, decorations)
    def test_equal_iff_components_are_equal(self, a: py.Decoration, b: py.Decoration):
        assert a != b or (
            a.decorator == b.decorator
            and a.target == b.target
            and a.comments == b.comments
        )

    @given(indent_levels)
    def test_with_a_function(self, level: int):
        f = py.Function("f", params=[], statements=[py.Assignment("a", py.Symbol("f"))])
        d = py.Decoration("task(2)", f)
        x = py.Line.INDENT_UNIT
        assert [str(l) for l in d.lines(level)] == [
            x * level + "@task(2)",
            *[str(l) for l in f.lines(level)],
        ]

    @given(indent_levels)
    def test_with_a_class(self, level: int):
        c = py.Class(
            "C", superclasses=[], statements=[py.Assignment("a: int", py.Literal(1))]
        )
        d = py.Decoration("task", c)
        x = py.Line.INDENT_UNIT
        assert [str(l) for l in d.lines(level)] == [
            x * level + "@task",
            *[str(l) for l in c.lines(level)],
        ]

    @given(indent_levels)
    def test_nested_decorators(self, level: int):
        f = py.Function("f", params=[], statements=[py.Assignment("a", py.Symbol("f"))])
        first = py.Decoration("task(2)", f)
        second = py.Decoration("task_seq(1)", first)
        x = py.Line.INDENT_UNIT
        assert [str(l) for l in second.lines(level)] == [
            x * level + "@task_seq(1)",
            x * level + "@task(2)",
            *[str(l) for l in f.lines(level)],
        ]

    @given(indent_levels)
    def test_lines_with_comments(self, level: int):
        f = py.Function("f", params=[], statements=[], comments=["1", "2"])
        d = py.Decoration("task", f, comments=["x", "y"])
        x = py.Line.INDENT_UNIT
        assert [str(l) for l in d.lines(level)] == [
            x * level + "# x",
            x * level + "# y",
            x * level + "@task",
            x * level + "# 1",
            x * level + "# 2",
            x * level + "def f():",
            x * (level + 1) + "pass",
        ]

    @given(indent_levels)
    def test_lines_with_hidden_comments(self, level: int):
        f = py.Function("f", params=[], statements=[], comments=["1", "2"])
        d = py.Decoration("task", f, comments=["x", "y"])
        x = py.Line.INDENT_UNIT
        assert [str(l) for l in d.lines(level, comments=False)] == [
            x * level + "@task",
            x * level + "def f():",
            x * (level + 1) + "pass",
        ]

    def test_repr(self):
        f = py.Function("f", params=[], statements=[])
        assert (
            repr(py.Decoration("task", f, comments=["hi"]))
            == f"Decoration('task', {f!r}, comments=['hi'])"
        )


class TestClass:
    @given(classes, classes)
    def test_equal_iff_components_are_equal(self, a: py.Class, b: py.Class):
        assert a != b or (
            a.name == b.name
            and a.statements == b.statements
            and a.superclasses == b.superclasses
            and a.comments == b.comments
        )

    @given(indent_levels)
    def test_empty_class(self, level: int):
        c = py.Class("A", statements=[])
        x = py.Line.INDENT_UNIT
        assert [str(l) for l in c.lines(level)] == [
            x * level + "class A:",
            x * (level + 1) + "pass",
        ]

    @given(
        lists(
            text(string.ascii_letters, min_size=1, max_size=2), min_size=0, max_size=3
        )
    )
    def test_class_with_superclasses(self, names: List[str]):
        c = py.Class("A", statements=[], superclasses=names)
        x = py.Line.INDENT_UNIT

        if names:
            expected = "(" + ", ".join(names) + ")"
        else:
            expected = ""

        assert [str(l) for l in c.lines()] == [f"class A{expected}:", x + "pass"]

    @given(indent_levels)
    def test_class_with_fields(self, level: int):
        c = py.Class(
            "A",
            statements=[
                py.Assignment("a", py.Literal(2)),
                py.Assignment("b", py.Literal(3)),
            ],
        )
        x = py.Line.INDENT_UNIT
        assert [str(l) for l in c.lines(level)] == [
            x * level + "class A:",
            x * (level + 1) + "a = 2",
            x * (level + 1) + "b = 3",
        ]

    @given(indent_levels)
    def test_lines_with_comments(self, level: int):
        stmt = py.OpaqueBlock("foo", comments=["x"])
        c = py.Class("C", statements=[stmt], superclasses=[], comments=["1", "2"])
        x = py.Line.INDENT_UNIT
        assert [str(l) for l in c.lines(level)] == [
            x * level + "# 1",
            x * level + "# 2",
            x * level + "class C:",
            *[str(l) for l in stmt.lines(level + 1)],
        ]

    @given(indent_levels)
    def test_lines_with_hidden_comments(self, level: int):
        stmt = py.OpaqueBlock("foo", comments=["x"])
        c = py.Class("C", statements=[stmt], superclasses=[], comments=["1", "2"])
        x = py.Line.INDENT_UNIT
        assert [str(l) for l in c.lines(level, comments=False)] == [
            x * level + "class C:",
            *[str(l) for l in stmt.lines(level + 1, comments=False)],
        ]

    def test_repr(self):
        stmts = [py.OpaqueBlock("raise")]
        assert (
            repr(
                py.Class(
                    name="C", statements=stmts, superclasses=["A"], comments=["hi"]
                )
            )
            == f"Class(name='C', statements={stmts!r}, superclasses=['A'], comments=['hi'])"
        )


class TestExpression:
    def test_str_must_be_implemented(self):
        with pytest.raises(NotImplementedError):
            str(py.Expression())


class TestStandalone:
    @given(standalones, standalones)
    def test_equal_iff_components_are_equal(self, a: py.Standalone, b: py.Standalone):
        assert a != b or (a.expr == b.expr and a.comments == b.comments)

    @given(expressions, indent_levels)
    def test_lines_returns_the_expression_as_single_line(
        self, e: py.Expression, level: int
    ):
        stmt = py.Standalone(e)
        lines = stmt.lines(level)
        assert len(lines) == 1
        line = lines[0]
        assert line.indent_level == level
        assert line.text == str(e)

    @given(indent_levels)
    def test_lines_with_comments(self, level: int):
        x = py.Line.INDENT_UNIT
        ob = py.Standalone(py.Symbol("a"), comments=["1", "2"])
        assert [str(l) for l in ob.lines(level)] == [
            x * level + "# 1",
            x * level + "# 2",
            x * level + "a",
        ]

    @given(indent_levels)
    def test_lines_with_hidden_comments(self, level: int):
        x = py.Line.INDENT_UNIT
        ob = py.Standalone(py.Symbol("a"), comments=["1", "2"])
        assert [str(l) for l in ob.lines(level, comments=False)] == [x * level + "a"]

    def test_repr(self):
        expr = py.Symbol("a")
        assert (
            repr(py.Standalone(expr, comments=["hi"]))
            == f"Standalone({expr!r}, comments=['hi'])"
        )


class TestLiteral:
    @given(literals, literals)
    def test_equal_iff_components_are_equal(self, a: py.Literal, b: py.Literal):
        assert a != b or a.value == b.value

    scalars = one_of(none(), booleans(), text(max_size=5), integers(), floats())

    @given(scalars)
    def test_literal_scalar_uses_repr(self, x):
        assert str(py.Literal(x)) == repr(x)

    @given(lists(scalars))
    def test_literal_list_of_scalars_uses_repr(self, x: list):
        assert str(py.Literal(x)) == repr(x)

    @given(dictionaries(scalars, scalars))
    def test_literal_dict_of_scalars_uses_repr(self, x: dict):
        assert str(py.Literal(x)) == repr(x)

    def test_literal_composites_with_expr_use_repr_except_for_expr(self):
        lit = py.Literal("b")
        assert repr(lit) == "Literal('b')", "repr works on literal"
        assert str(py.Literal([1, {"a": lit}])) == "[1, {'a': 'b'}]"
        assert repr(lit) == "Literal('b')", "repr is still working on literal"


class TestFString:
    def test_strings_appear_as_fstrings(self):
        assert str(py.FString("")) == "f''"
        assert str(py.FString("ab")) == "f'ab'"
        assert str(py.FString("a'b")) == """ f"a'b" """.strip()
        assert str(py.FString('a"b')) == """ f'a"b' """.strip()

    def test_non_strings_raise_error(self):
        with pytest.raises(TypeError):
            assert str(py.FString(24))

    def test_format_template_is_not_replaced(self):
        a = 2
        assert str(py.FString("a {a} {} {a!r}")) == "f'a {a} {} {a!r}'"


class TestSymbol:
    @given(symbols, symbols)
    def test_equal_iff_components_are_equal(self, a: py.Symbol, b: py.Symbol):
        assert a != b or a.name == b.name

    @given(text(string.ascii_letters))
    def test_strings_appear_unchanged(self, s: str):
        assert str(py.Symbol(s)) == s

    def test_non_strings_raise_error(self):
        with pytest.raises(TypeError):
            assert str(py.Symbol(True))

    def test_repr(self):
        assert repr(py.Symbol(" x'\" y ")) == "Symbol(' x\\'\" y ')"


class TestFunctionCall:
    @given(function_calls, function_calls)
    def test_equal_iff_components_are_equal(
        self, a: py.FunctionCall, b: py.FunctionCall
    ):
        assert a != b or (
            a.name == b.name
            and a.positional_args == b.positional_args
            and a.named_args == b.named_args
        )

    def test_with_no_args(self):
        assert str(py.FunctionCall("f")) == "f()"

    def test_with_positional_args(self):
        assert str(py.FunctionCall("f", [py.Literal(2)])) == "f(2)"

    def test_with_kwargs(self):
        assert (
            str(
                py.FunctionCall(
                    "f", named_args={"a": py.Literal(2), "bc": py.Literal("x")}
                )
            )
            == "f(a=2, bc='x')"
        )

    def test_with_positional_and_kwargs(self):
        assert (
            str(
                py.FunctionCall(
                    "m.f",
                    [py.Literal(True), py.FunctionCall("g", [py.Symbol("f")])],
                    {"a": py.Literal(2), "bc": py.Literal("x")},
                )
            )
            == "m.f(True, g(f), a=2, bc='x')"
        )

    def test_repr(self):
        arg = py.Symbol("a")
        kwarg = py.Symbol("v")
        assert (
            repr(py.FunctionCall("f", [arg], {"k": kwarg}))
            == f"FunctionCall('f', [{arg!r}], {{'k': {kwarg!r}}})"
        )


class TestBinaryOp:
    @given(binary_ops, binary_ops)
    def test_equal_iff_components_are_equal(self, a: py.BinaryOp, b: py.BinaryOp):
        assert a != b or (a.op == b.op and a.lhs == b.lhs and a.rhs == b.rhs)

    def test_simple(self):
        assert str(py.BinaryOp(py.Literal(2), "**", py.Literal(10))) == "2 ** 10"

    def test_nested(self):
        assert (
            str(
                py.BinaryOp(
                    py.Literal(2),
                    "+",
                    py.BinaryOp(
                        py.BinaryOp(py.Literal(3), "-", py.Literal(4)),
                        "*",
                        py.Literal(5),
                    ),
                )
            )
            == "2 + ((3 - 4) * 5)"
        )


class TestAssignment:
    @given(assignments, assignments)
    def test_equal_iff_components_are_equal(self, a: py.Assignment, b: py.Assignment):
        assert a != b or (
            a.lhs == b.lhs and a.rhs == b.rhs and a.comments == b.comments
        )

    @given(indent_levels)
    def test_simple(self, level: int):
        x = py.Line.INDENT_UNIT
        assert [str(l) for l in py.Assignment("foo", py.Literal(3)).lines(level)] == [
            x * level + "foo = 3"
        ]

    @given(indent_levels)
    def test_lines_with_comments(self, level: int):
        x = py.Line.INDENT_UNIT
        stmt = py.Assignment("x", py.Symbol("a"), comments=["1", "2"])
        assert [str(l) for l in stmt.lines(level)] == [
            x * level + "# 1",
            x * level + "# 2",
            x * level + "x = a",
        ]

    @given(indent_levels)
    def test_lines_with_hidden_comments(self, level: int):
        x = py.Line.INDENT_UNIT
        stmt = py.Assignment("x", py.Symbol("a"), comments=["1", "2"])
        assert [str(l) for l in stmt.lines(level, comments=False)] == [
            x * level + "x = a"
        ]

    def test_repr(self):
        rhs = py.Symbol("a")
        assert (
            repr(py.Assignment("x", rhs, comments=["hi"]))
            == f"Assignment(lhs='x', rhs={rhs!r}, comments=['hi'])"
        )


class TestIfElse:
    @given(ifelses, ifelses)
    def test_equal_iff_components_are_equal(self, a: py.IfElse, b: py.IfElse):
        assert a != b or (
            a.condition_blocks == b.condition_blocks
            and a.else_block == b.else_block
            and a.comments == b.comments
        )

    def test_init_with_no_condition_raises_error(self):
        with pytest.raises(ValueError):
            py.IfElse([])
        with pytest.raises(ValueError):
            py.IfElse([], else_block=[])
        with pytest.raises(ValueError):
            py.IfElse([], else_block=[py.Assignment("a", py.Literal(2))])

    @given(indent_levels)
    def test_lines_for_single_if(self, level: int):
        x = py.Line.INDENT_UNIT
        assert [
            str(l)
            for l in py.IfElse(
                [
                    (
                        py.BinaryOp(py.Symbol("t"), "is", py.Literal(None)),
                        [py.Assignment("t", py.Literal(1))],
                    )
                ]
            ).lines(level)
        ] == [x * level + "if t is None:", x * (level + 1) + "t = 1"]

    @given(indent_levels)
    def test_lines_for_if_else(self, level: int):
        x = py.Line.INDENT_UNIT
        assert [
            str(l)
            for l in py.IfElse(
                [
                    (
                        py.BinaryOp(py.Symbol("t"), "is", py.Literal(None)),
                        [py.Assignment("t", py.Literal(1))],
                    )
                ],
                [py.Assignment("t", py.Literal(2))],
            ).lines(level)
        ] == [
            x * level + "if t is None:",
            x * (level + 1) + "t = 1",
            x * level + "else:",
            x * (level + 1) + "t = 2",
        ]

    @given(indent_levels)
    def test_lines_for_if_elif(self, level: int):
        x = py.Line.INDENT_UNIT
        assert [
            str(l)
            for l in py.IfElse(
                [
                    (
                        py.BinaryOp(py.Symbol("t"), "is", py.Literal(None)),
                        [py.Assignment("t", py.Literal(1))],
                    ),
                    (
                        py.Literal(False),
                        [
                            py.Assignment("t", py.Literal(2)),
                            py.Standalone(py.FunctionCall("f", [py.Symbol("t")])),
                        ],
                    ),
                ]
            ).lines(level)
        ] == [
            x * level + "if t is None:",
            x * (level + 1) + "t = 1",
            x * level + "elif False:",
            x * (level + 1) + "t = 2",
            x * (level + 1) + "f(t)",
        ]

    @given(indent_levels)
    def test_lines_for_if_elif_else_with_no_statements(self, level: int):
        x = py.Line.INDENT_UNIT
        assert [
            str(l)
            for l in py.IfElse(
                [
                    (py.BinaryOp(py.Symbol("t"), "is", py.Literal(None)), []),
                    (py.Literal(False), []),
                    (py.Literal(True), []),
                ],
                [],
            ).lines(level)
        ] == [
            x * level + "if t is None:",
            x * (level + 1) + "pass",
            x * level + "elif False:",
            x * (level + 1) + "pass",
            x * level + "elif True:",
            x * (level + 1) + "pass",
        ]

    @given(indent_levels)
    def test_lines_with_comments(self, level: int):
        x = py.Line.INDENT_UNIT
        cond = py.Literal(True)
        if_true = py.Assignment("x", py.Symbol("a"), comments=["tx", "ty"])
        if_false = py.Assignment("x", py.Symbol("b"), comments=["fx", "fy"])
        stmt = py.IfElse([(cond, [if_true])], [if_false], comments=["1", "2"])
        assert [str(l) for l in stmt.lines(level)] == [
            x * level + "# 1",
            x * level + "# 2",
            x * level + "if True:",
            *[str(l) for l in if_true.lines(level + 1)],
            x * level + "else:",
            *[str(l) for l in if_false.lines(level + 1)],
        ]

    @given(indent_levels)
    def test_lines_with_hidden_comments(self, level: int):
        x = py.Line.INDENT_UNIT
        cond = py.Literal(True)
        if_true = py.Assignment("x", py.Symbol("a"), comments=["tx", "ty"])
        if_false = py.Assignment("x", py.Symbol("b"), comments=["fx", "fy"])
        stmt = py.IfElse([(cond, [if_true])], [if_false], comments=["1", "2"])
        assert [str(l) for l in stmt.lines(level, comments=False)] == [
            x * level + "if True:",
            *[str(l) for l in if_true.lines(level + 1, comments=False)],
            x * level + "else:",
            *[str(l) for l in if_false.lines(level + 1, comments=False)],
        ]

    def test_repr(self):
        cond = py.Literal(True)
        if_true = [(cond, [py.Assignment("x", py.Symbol("a"))])]
        stmt = py.IfElse(condition_blocks=if_true, comments=["hi"])
        assert (
            repr(stmt)
            == f"IfElse(condition_blocks={if_true!r}, else_block=None, comments=['hi'])"
        )


class TestImport:
    @given(imports, imports)
    def test_equal_iff_components_are_equal(self, a: py.Import, b: py.Import):
        assert a != b or (
            a.targets == b.targets
            and a.source == b.source
            and a.alias == b.alias
            and a.comments == b.comments
        )

    def test_init_without_targets_raises_error(self):
        with pytest.raises(ValueError):
            py.Import([])

    @given(indent_levels)
    def test_lines_without_targets_raises_error(self, level: int):
        i = py.Import(["safe"])
        i.targets.clear()
        with pytest.raises(ValueError):
            i.lines(level)

    @given(indent_levels)
    def test_lines_with_single_target(self, level: int):
        x = py.Line.INDENT_UNIT
        name = "locust.http"
        assert [str(l) for l in py.Import([name]).lines(level)] == [
            x * level + f"import {name}"
        ]

    @given(indent_levels)
    def test_lines_with_multiple_targets(self, level: int):
        x = py.Line.INDENT_UNIT
        names = ["locust.http", "math", "a.b.c"]
        assert [str(l) for l in py.Import(names).lines(level)] == [
            x * level + f"import {name}" for name in names
        ]

    @given(indent_levels)
    def test_lines_with_single_target_and_alias(self, level: int):
        x = py.Line.INDENT_UNIT
        name = "transformer.python"
        alias = "py"
        assert [str(l) for l in py.Import([name], alias=alias).lines(level)] == [
            x * level + f"import {name} as {alias}"
        ]

    def test_init_with_multiple_targets_and_alias_raises_error(self):
        with pytest.raises(ValueError):
            py.Import(["a", "b"], alias="c")

    @given(indent_levels)
    def test_lines_with_multiple_targets_and_alias_raises_error(self, level: int):
        i = py.Import(["safe"], alias="A")
        i.targets.append("oops")
        with pytest.raises(ValueError):
            i.lines(level)

    @given(indent_levels, lists(text(string.ascii_letters, min_size=1), min_size=1))
    def test_lines_with_targets_and_source(self, level: int, targets: List[str]):
        x = py.Line.INDENT_UNIT
        source = "bar"
        assert [str(l) for l in py.Import(targets, source).lines(level)] == [
            x * level + f"from {source} import {target}" for target in targets
        ]

    @given(indent_levels)
    def test_lines_with_comments(self, level: int):
        x = py.Line.INDENT_UNIT
        stmt = py.Import(["a", "b", "c"], comments=["1", "2"])
        assert [str(l) for l in stmt.lines(level)] == [
            x * level + "# 1",
            x * level + "# 2",
            x * level + "import a",
            x * level + "import b",
            x * level + "import c",
        ]

    @given(indent_levels)
    def test_lines_with_hidden_comments(self, level: int):
        x = py.Line.INDENT_UNIT
        stmt = py.Import(["a", "b", "c"], comments=["1", "2"])
        assert [str(l) for l in stmt.lines(level, comments=False)] == [
            x * level + "import a",
            x * level + "import b",
            x * level + "import c",
        ]

    def test_repr(self):
        stmt = py.Import(targets=["a", "b"], comments=["hi"])
        assert (
            repr(stmt)
            == f"Import(targets=['a', 'b'], source=None, alias=None, comments=['hi'])"
        )
