import ast
import pytest

from docugenius.core import (
    extract_definitions,
    get_function_metadata,
    instrument_tree,
    generate_instrumented_code,
    validate_file,
)


# ---------------------------------------------------
# HELPER
# ---------------------------------------------------

def _parse_and_instrument(source: str):
    tree = ast.parse(source)
    instrumented_tree = instrument_tree(tree, style="google")
    return generate_instrumented_code(instrumented_tree)


# ---------------------------------------------------
# TESTS
# ---------------------------------------------------

def test_empty_python_file():
    """Empty file should be handled safely."""
    source = ""
    tree = ast.parse(source)
    instrumented_tree = instrument_tree(tree, style="google")
    code = generate_instrumented_code(instrumented_tree)

    # Should at least be valid Python
    ast.parse(code)


def test_module_docstring_added_if_missing():
    """Missing module docstring should be inserted."""
    source = "x = 10"
    instrumented = _parse_and_instrument(source)

    assert 'Module description.' in instrumented


def test_nested_function_only_outer_documented():
    """Nested functions should not necessarily be documented separately."""
    source = """
def outer():
    def inner():
        return 5
    return inner()
"""
    instrumented = _parse_and_instrument(source)
    tree = ast.parse(instrumented)

    functions, _ = extract_definitions(tree)

    outer = next(f for f in functions if f.name == "outer")
    inner = next(f for f in functions if f.name == "inner")

    assert ast.get_docstring(outer) is not None
    # Depending on your logic, inner may or may not be documented.
    # If your instrumentor documents all, change this to assert is not None.
    assert ast.get_docstring(inner) is not None


def test_class_without_methods_gets_docstring():
    """Classes without methods should receive docstrings."""
    source = """
class Empty:
    pass
"""
    instrumented = _parse_and_instrument(source)
    tree = ast.parse(instrumented)
    _, classes = extract_definitions(tree)

    empty_class = next(c for c in classes if c.name == "Empty")
    assert ast.get_docstring(empty_class) is not None


def test_existing_docstring_not_overwritten():
    """Existing function docstrings should not be replaced."""
    source = '''
def multiply(x, y):
    """Existing docstring."""
    return x * y
'''
    instrumented = _parse_and_instrument(source)
    tree = ast.parse(instrumented)
    functions, _ = extract_definitions(tree)

    multiply = next(f for f in functions if f.name == "multiply")
    doc = ast.get_docstring(multiply)

    assert "Existing docstring." in doc


def test_syntax_error_raises():
    """Invalid Python input should raise SyntaxError."""
    bad_source = "def broken(:"
    with pytest.raises(SyntaxError):
        ast.parse(bad_source)


def test_generated_code_removes_missing_docstring_violation():
    """Instrumentation should remove missing-docstring violations."""
    source = """
def add(x, y):
    return x + y
"""

    before_violations = validate_file(source)
    instrumented = _parse_and_instrument(source)
    after_violations = validate_file(instrumented)

    before_codes = {v["code"] for v in before_violations}
    after_codes = {v["code"] for v in after_violations}

    # Ensure D103 (missing function docstring) is removed
    assert "D103" in before_codes
    assert "D103" not in after_codes

