"""Tests for Data Transformation Density (Python)."""

from vibelexity.analyzers.python import analyze_source


def _dtd(code: str) -> int:
    """Return DTD of the first function in code."""
    result = analyze_source(code)
    assert result.functions, f"No functions found in:\n{code}"
    dtd_val = result.functions[0].dtd
    assert dtd_val is not None
    return dtd_val


def _dtds(code: str) -> dict[str, int]:
    """Return {name: dtd} for all functions in code."""
    result = analyze_source(code)
    return {f.name: f.dtd for f in result.functions if f.dtd is not None}


# --- Empty / no data ---


def test_empty_function():
    assert _dtd("def f(): pass") == 0


def test_no_data_structures():
    code = """
def f(x):
    return x + 1
"""
    assert _dtd(code) == 0


# --- Flat structures ---


def test_flat_dict():
    code = """
def f():
    return {"a": 1, "b": 2, "c": 3}
"""
    # dict at depth 1, 3 pairs x 1 = 3
    assert _dtd(code) == 3


def test_flat_list():
    code = """
def f():
    return [1, 2, 3]
"""
    # list at depth 1, 3 elements x 1 = 3
    assert _dtd(code) == 3


# --- Nested structures ---


def test_nested_dict_in_list():
    code = """
def f():
    return [{"a": 1}, {"b": 2}]
"""
    # list at depth 1: 2 elements x 1 = 2
    # each dict at depth 2: 1 pair x 2 = 2 each -> 4
    # total = 2 + 4 = 6
    assert _dtd(code) == 6


def test_deeply_nested():
    code = """
def f():
    return {"outer": {"inner": [1, 2]}}
"""
    # outer dict depth 1: 1 pair x 1 = 1
    # inner dict depth 2: 1 pair x 2 = 2
    # list depth 3: 2 elements x 3 = 6
    # total = 1 + 2 + 6 = 9
    assert _dtd(code) == 9


# --- Nested function skipped ---


def test_nested_function_skipped():
    code = """
def outer():
    x = {"a": 1}
    def inner():
        y = {"b": 2, "c": 3}
    return x
"""
    dtds = _dtds(code)
    assert dtds["outer"] == 1  # only {"a": 1}
    assert dtds["inner"] == 2  # {"b": 2, "c": 3} -> 2 pairs x 1


# --- Comprehensions ---


def test_dict_comprehension():
    code = """
def f(items):
    return {k: v for k, v in items}
"""
    # dict_comprehension is a container at depth 1
    # pair inside at depth 1: 1 pair x 1 = 1
    assert _dtd(code) == 1


def test_list_comprehension():
    code = """
def f(items):
    return [x for x in items]
"""
    # list_comprehension is a container at depth 1
    # no explicit payload children (the inner expression isn't a pair/keyword_arg)
    assert _dtd(code) == 0
