from synthegrator.lang_specs.lang_spec import PythonFunctionParser
from synthegrator.lang_specs.lang_spec_python import PythonLangSpec


def test_get_function():
    src = "def foo():\n    pass\n"
    funcs = list(PythonLangSpec().find_functions(src))
    assert len(funcs) == 1
    func = funcs[0]
    assert func.get_function_name() == "foo"
    assert func.get_declaration_str() == "def foo():\n"
    assert func.get_declaration_str(include_body_indent=True) == "def foo():\n    "
    assert func.get_body_src() == "    pass\n"
    assert func.get_body_src(include_body_indent=False) == "pass\n"
    assert func.get_prefix_comment_str() is None


def test_get_function_multiline():
    src = "def foo():\n    x = 2\n    return x * 2\n"
    funcs = list(PythonLangSpec().find_functions(src))
    assert len(funcs) == 1
    func = funcs[0]
    assert func.get_function_name() == "foo"
    assert func.get_declaration_str() == "def foo():\n"
    assert func.get_body_src() == "    x = 2\n    return x * 2\n"
    assert func.get_prefix_comment_str() is None


def test_get_function_args():
    src = "def foo(a):\n    pass\n"
    funcs = list(PythonLangSpec().find_functions(src))
    assert len(funcs) == 1
    func = funcs[0]
    assert func.get_function_name() == "foo"
    assert func.get_declaration_str() == "def foo(a):\n"
    assert func.get_body_src() == "    pass\n"
    assert func.get_prefix_comment_str() is None


def test_get_function_args2():
    src = "def foo(a, b):\n    pass\n"
    funcs = list(PythonLangSpec().find_functions(src))
    assert len(funcs) == 1
    func = funcs[0]
    assert func.get_function_name() == "foo"
    assert func.get_declaration_str() == "def foo(a, b):\n"
    assert func.get_body_src() == "    pass\n"
    assert func.get_prefix_comment_str() is None


def test_get_function_args_type_hint():
    src = "def foo(a: int, b: list[tuple[str, ...]]):\n    pass\n"
    funcs = list(PythonLangSpec().find_functions(src))
    assert len(funcs) == 1
    func = funcs[0]
    assert func.get_function_name() == "foo"
    assert func.get_declaration_str() == src.splitlines(keepends=True)[0]
    assert func.get_body_src() == "    pass\n"
    assert func.get_prefix_comment_str() is None


def test_get_two_funcs():
    src = "def foo():\n    pass\n\ndef bar():\n    pass\n"
    funcs = list(PythonLangSpec().find_functions(src))
    assert len(funcs) == 2
    func = funcs[0]
    assert func.get_function_name() == "foo"
    assert func.get_declaration_str() == src.splitlines(keepends=True)[0]
    assert func.get_body_src() == "    pass\n"
    assert func.get_prefix_comment_str() is None
    func = funcs[1]
    assert func.get_function_name() == "bar"
    assert func.get_declaration_str() == src.splitlines(keepends=True)[3]
    assert func.get_body_src() == "    pass\n"
    assert func.get_prefix_comment_str() is None


def test_get_function_multiline_decl():
    src = "def foo(\n   a: int, \n   b: list[tuple[str, ...]]\n) -> str:\n    pass\n"
    funcs = list(PythonLangSpec().find_functions(src))
    assert len(funcs) == 1
    func = funcs[0]
    assert func.get_function_name() == "foo"
    assert func.get_declaration_str() == src[: src.index("pass") - 4]
    assert func.get_body_src() == "    pass\n"
    assert func.get_prefix_comment_str() is None


def test_line_numbers():
    doc = """
# hi

def h(x,
    y,
    z,
    k) -> None:
    \"\"\"hi\"\"\"
    logging.info(x, y)
    return None"""

    doc_f = PythonFunctionParser(doc)
    assert doc_f.body_span_begin == 8
    assert doc_f.body_span_end == 9

    short_function = PythonFunctionParser(
        """def hi():
    pass""",
    )
    assert short_function.body_span_begin == 1


def test_funtion_parser():
    definition = """def h(x,
    y) -> None:"""

    def_w_doc = """def h(x,
    y) -> None:
    \"\"\"hi\"\"\""""

    doc = """def h(x,
    y) -> None:
    \"\"\"hi\"\"\"
    logging.info(x, y)"""

    no_doc = """def h(x,
    y) -> None:
    logging.info(x, y)"""

    doc_f = PythonFunctionParser(doc)
    assert doc_f.has_docstring is True
    assert doc_f.function_without_doc == no_doc
    assert doc_f.function_definition_with_doc == def_w_doc
    assert doc_f.function_definition_without_doc == definition
    assert doc_f.body_span_begin == 3

    doc_nf = PythonFunctionParser(no_doc)
    assert doc_nf.has_docstring is False
    assert doc_nf.function_without_doc == no_doc
    assert doc_nf.function_definition_with_doc == definition
    assert doc_nf.function_definition_without_doc == definition
