import pytest

from synthegrator.lang_specs.lang_spec import (
    PythonFunctionParser,
)
from synthegrator.lang_specs.lang_spec_python import PythonLangSpec, LsFunctionDefPython
from synthegrator.lang_specs.lang_util import truncate_when_decreases_indent


@pytest.fixture
def example_dypy_method_nodocstring():
    return """def add_task(
        self,
        task: Task,
        queue: None | BaseTaskQueue = None,
        raise_error: bool = False,
    ) -> bool:
        if queue is None:
            queue = self.task_queue
        if task.priority is None or not task.priority_set_explicitly:
            task.priority = self.generate_task_priority()
            task.priority_set_explicitly = False
        else:
            task.priority_set_explicitly = True

        if not task.request.url or not task.request.url.startswith(
            ("http://", "https://", "ftp://", "file://", "feed://")
        ):
            self.collect_runtime_event("task-with-invalid-url", task.request.url)
            msg = "Invalid task URL: %s" % task.request.url
            if raise_error:
                raise SpiderError(msg)
            logger.error(
                "%s\nTraceback:\n%s",
                msg,
                "".join(format_stack()),
            )
            return False
        # TODO: keep original task priority if it was set explicitly
        # WTF the previous comment means?
        queue.put(task, priority=task.priority, schedule_time=task.schedule_time)
        return True"""


def test_get_prefixcomments_only_at_start():
    langspec_py = PythonLangSpec()
    comments = list(
        langspec_py.get_prefix_comments(
            "def foo():\n"
            "    this_isnot()\n"
            "    '''\n"
            "    This is not a docstring\n"
            "    '''\n"
            "    pass\n"
            "def bar():\n"
            "    this_isnot()\n"
            "    '''\n"
            "    This is not another docstring\n"
            "    '''\n"
            "    pass\n"
            "def boo():\n"
            "    '''\n"
            "    This is a docstring\n"
            "    '''\n"
            "    pass\n",
        ),
    )
    assert len(comments) == 1
    assert comments[0].doc_string_str == "'''\n    This is a docstring\n    '''"


def test_treesitter_getdocstring_only_start1(example_dypy_method_nodocstring):
    parser = PythonFunctionParser(example_dypy_method_nodocstring)
    assert parser.docstring_node is None


def test_treesitter_getdocstring_only_start2(example_dypy_method_nodocstring):
    parser = PythonFunctionParser(
        "def foo():\n"
        "    this_isnot()\n"
        "    '''\n"
        "    This is not a docstring\n"
        "    '''\n"
        "    pass\n",
    )
    assert parser.docstring_node is None


def test_treesitter_getdocstring_only_start3(example_dypy_method_nodocstring):
    parser = PythonFunctionParser(
        "def foo():\n    '''\n    This is a docstring\n    '''\n    pass\n",
    )
    assert (
        parser.docstring_node.text.decode() == "'''\n    This is a docstring\n    '''"
    )


def test_treesitter_getdocstring_only_start1():
    fixture1 = (
        "def foo():\n"
        "    this_isnot()\n"
        "    '''\n"
        "    This is a docstring\n"
        "    '''\n"
        "    pass\n"
    )
    parser = PythonFunctionParser(fixture1)
    assert parser.docstring_node is None


def test_treesitter_getdocstring_only_start2():
    fixture1 = "def foo():\n    '''\n    This is a docstring\n    '''\n    pass\n"
    parser = PythonFunctionParser(fixture1)
    assert parser.docstring_node is not None
    assert (
        parser.docstring_node.text.decode() == "'''\n    This is a docstring\n    '''"
    )


def test_get_docstrings():
    langspec_py = PythonLangSpec()
    comments = list(
        langspec_py.get_prefix_comments(
            "def foo():\n"
            "    '''\n"
            "    This is a docstring\n"
            "    '''\n"
            "    pass\n"
            "def bar():\n"
            "    '''\n"
            "    This is other docstring\n"
            "    '''\n"
            "    pass\n",
        ),
    )
    assert len(comments) == 2
    assert comments[0].doc_string_node.start_point == (1, 4)
    assert comments[0].doc_string_node.end_point == (3, 4 + 3)
    assert (
        comments[0].doc_string_node.text.decode("utf-8")
        == "'''\n    This is a docstring\n    '''"
    )
    assert comments[0].function_node.start_point == (0, 0)
    assert comments[0].function_node.end_point == (4, 4 + 4)


def test_truncate_when_decrease_indent():
    result = truncate_when_decreases_indent("    return 2\nassert False")
    assert result == "    return 2"


def test_truncate_when_decrease_indent2():
    result = truncate_when_decreases_indent(  # fmt: off
        "    x = 2\n"
        "    y = x * 2\n"
        "    for foo in bar:\n"
        "        y += 1\n"
        "    return y\n"
        "assert False",
    )
    assert (
        result
        == "    x = 2\n    y = x * 2\n    for foo in bar:\n        y += 1\n    return y"
    )  # fmt: on


def test_get_body_src_include_prefix_comment_triple_double_quotes():
    """Test get_body_src with include_prefix_comment parameter using triple double quotes"""
    source = '''def foo():
    """This is a docstring"""
    x = 1
    return x'''
    
    langspec = PythonLangSpec()
    tree = langspec.get_tree_sitter_tree(source)
    func_defs = list(PythonFunctionParser.get_function_definitions(tree.root_node))
    func_def = LsFunctionDefPython(func_defs[0], tree)
    
    # With prefix comment (default behavior)
    with_prefix = func_def.get_body_src(include_trailing_new_line=False, include_prefix_comment=True)
    assert with_prefix == '    """This is a docstring"""\n    x = 1\n    return x'
    
    # Without prefix comment
    without_prefix = func_def.get_body_src(include_trailing_new_line=False, include_prefix_comment=False)
    assert without_prefix == '    x = 1\n    return x'


def test_get_body_src_include_prefix_comment_triple_single_quotes():
    """Test get_body_src with include_prefix_comment parameter using triple single quotes"""
    source = '''def bar():
    \'\'\'This is a docstring\'\'\'
    x = 1
    return x'''
    
    langspec = PythonLangSpec()
    tree = langspec.get_tree_sitter_tree(source)
    func_defs = list(PythonFunctionParser.get_function_definitions(tree.root_node))
    func_def = LsFunctionDefPython(func_defs[0], tree)
    
    # With prefix comment (default behavior)
    with_prefix = func_def.get_body_src(include_trailing_new_line=False, include_prefix_comment=True)
    assert with_prefix == "    '''This is a docstring'''\n    x = 1\n    return x"
    
    # Without prefix comment
    without_prefix = func_def.get_body_src(include_trailing_new_line=False, include_prefix_comment=False)
    assert without_prefix == '    x = 1\n    return x'


def test_get_body_src_include_prefix_comment_multiline_docstring():
    """Test get_body_src with multiline docstrings using both quote styles"""
    # Multi-line triple double quotes
    source_double = '''def baz():
    """
    This is a 
    multi-line docstring
    """
    x = 1
    return x'''
    
    # Multi-line triple single quotes
    source_single = '''def qux():
    \'\'\'
    This is a 
    multi-line docstring
    \'\'\'
    x = 1
    return x'''
    
    langspec = PythonLangSpec()
    
    for source in [source_double, source_single]:
        tree = langspec.get_tree_sitter_tree(source)
        func_defs = list(PythonFunctionParser.get_function_definitions(tree.root_node))
        func_def = LsFunctionDefPython(func_defs[0], tree)
        
        # With prefix comment should include docstring
        with_prefix = func_def.get_body_src(include_trailing_new_line=False, include_prefix_comment=True)
        assert 'multi-line docstring' in with_prefix
        assert 'x = 1' in with_prefix
        
        # Without prefix comment should exclude docstring
        without_prefix = func_def.get_body_src(include_trailing_new_line=False, include_prefix_comment=False)
        assert 'multi-line docstring' not in without_prefix
        assert without_prefix == '    x = 1\n    return x'


def test_get_body_src_include_prefix_comment_no_docstring():
    """Test get_body_src with include_prefix_comment when there's no docstring"""
    source = '''def nope():
    x = 1
    return x'''
    
    langspec = PythonLangSpec()
    tree = langspec.get_tree_sitter_tree(source)
    func_defs = list(PythonFunctionParser.get_function_definitions(tree.root_node))
    func_def = LsFunctionDefPython(func_defs[0], tree)
    
    # Both should return the same result when there's no docstring
    with_prefix = func_def.get_body_src(include_trailing_new_line=False, include_prefix_comment=True)
    without_prefix = func_def.get_body_src(include_trailing_new_line=False, include_prefix_comment=False)
    
    assert with_prefix == without_prefix == '    x = 1\n    return x'


def test_get_body_src_include_prefix_comment_invalid_docstring():
    """Test get_body_src with include_prefix_comment when docstring is not at start"""
    source = '''def invalid():
    x = 1
    """This is not a docstring"""
    return x'''
    
    langspec = PythonLangSpec()
    tree = langspec.get_tree_sitter_tree(source)
    func_defs = list(PythonFunctionParser.get_function_definitions(tree.root_node))
    func_def = LsFunctionDefPython(func_defs[0], tree)
    
    # Both should return the same result when there's no valid prefix docstring
    with_prefix = func_def.get_body_src(include_trailing_new_line=False, include_prefix_comment=True)
    without_prefix = func_def.get_body_src(include_trailing_new_line=False, include_prefix_comment=False)
    
    assert with_prefix == without_prefix == '    x = 1\n    """This is not a docstring"""\n    return x'


def test_get_body_src_include_prefix_comment_only_docstring():
    """Test get_body_src with include_prefix_comment when function has only a docstring"""
    source = '''def empty():
    """Just a docstring"""'''
    
    langspec = PythonLangSpec()
    tree = langspec.get_tree_sitter_tree(source)
    func_defs = list(PythonFunctionParser.get_function_definitions(tree.root_node))
    func_def = LsFunctionDefPython(func_defs[0], tree)
    
    # With prefix comment should include docstring
    with_prefix = func_def.get_body_src(include_trailing_new_line=False, include_prefix_comment=True)
    assert with_prefix == '    """Just a docstring"""'
    
    # Without prefix comment should be empty
    without_prefix = func_def.get_body_src(include_trailing_new_line=False, include_prefix_comment=False)
    assert without_prefix == ''


def test_get_body_src_include_prefix_comment_with_other_params():
    """Test get_body_src with include_prefix_comment combined with other parameters"""
    source = '''def foo():
    """This is a docstring"""
    x = 1
    return x'''
    
    langspec = PythonLangSpec()
    tree = langspec.get_tree_sitter_tree(source)
    func_defs = list(PythonFunctionParser.get_function_definitions(tree.root_node))
    func_def = LsFunctionDefPython(func_defs[0], tree)
    
    # Test with trailing newline
    with_trailing = func_def.get_body_src(include_trailing_new_line=True, include_prefix_comment=False)
    assert with_trailing == '    x = 1\n    return x\n'
    
    # Test without trailing newline
    without_trailing = func_def.get_body_src(include_trailing_new_line=False, include_prefix_comment=False)
    assert without_trailing == '    x = 1\n    return x'


def test_get_body_src_include_prefix_comment_nested_function():
    """Test get_body_src with include_prefix_comment for nested functions"""
    source = '''def outer():
    """Outer function docstring"""
    x = 1
    
    def inner():
        """Inner function docstring"""
        y = 2
        return y
    
    return inner() + x'''
    
    langspec = PythonLangSpec()
    tree = langspec.get_tree_sitter_tree(source)
    func_defs = list(PythonFunctionParser.get_function_definitions(tree.root_node))
    
    # Test outer function
    outer_func = LsFunctionDefPython(func_defs[0], tree)
    
    # With prefix comment should include outer docstring
    with_prefix = outer_func.get_body_src(include_trailing_new_line=False, include_prefix_comment=True)
    assert '"""Outer function docstring"""' in with_prefix
    assert 'def inner():' in with_prefix
    assert '"""Inner function docstring"""' in with_prefix
    
    # Without prefix comment should exclude outer docstring but keep inner
    without_prefix = outer_func.get_body_src(include_trailing_new_line=False, include_prefix_comment=False)
    assert '"""Outer function docstring"""' not in without_prefix
    assert 'def inner():' in without_prefix
    assert '"""Inner function docstring"""' in without_prefix  # Inner docstring should remain
    assert 'x = 1' in without_prefix
    assert 'return inner() + x' in without_prefix
