"""Tests for call graph extraction."""

import pytest

from tokennuke.parser.extractor import parse_file
from tokennuke.parser.call_graph import extract_calls
from tokennuke.parser.languages import LANGUAGE_REGISTRY

try:
    from tree_sitter_language_pack import get_parser
    HAS_TREE_SITTER = True
except ImportError:
    HAS_TREE_SITTER = False


PYTHON_WITH_CALLS = '''
def authenticate(username, password):
    user = find_user(username)
    if user and verify_password(user, password):
        return create_token(user)
    return None

def main():
    token = authenticate("admin", "secret")
    print(token)
'''


@pytest.mark.skipif(not HAS_TREE_SITTER, reason='tree-sitter not installed')
class TestCallGraphExtraction:
    def test_extracts_calls_python(self):
        spec = LANGUAGE_REGISTRY['python']
        source_bytes = PYTHON_WITH_CALLS.encode('utf-8')
        parser = get_parser(spec.ts_language)
        tree = parser.parse(source_bytes)

        # Find the authenticate function node
        symbols = parse_file(PYTHON_WITH_CALLS, 'test.py', 'python')
        auth_sym = next(s for s in symbols if s.name == 'authenticate')

        # Find AST node for authenticate
        auth_node = _find_node(tree.root_node, auth_sym.byte_offset, auth_sym.byte_length)
        assert auth_node is not None

        calls = extract_calls(auth_node, spec, source_bytes)
        callee_names = {c.callee_name for c in calls}
        assert 'find_user' in callee_names
        assert 'verify_password' in callee_names
        assert 'create_token' in callee_names

    def test_extracts_calls_from_main(self):
        spec = LANGUAGE_REGISTRY['python']
        source_bytes = PYTHON_WITH_CALLS.encode('utf-8')
        parser = get_parser(spec.ts_language)
        tree = parser.parse(source_bytes)

        symbols = parse_file(PYTHON_WITH_CALLS, 'test.py', 'python')
        main_sym = next(s for s in symbols if s.name == 'main')

        main_node = _find_node(tree.root_node, main_sym.byte_offset, main_sym.byte_length)
        calls = extract_calls(main_node, spec, source_bytes)
        callee_names = {c.callee_name for c in calls}
        assert 'authenticate' in callee_names
        assert 'print' in callee_names

    def test_call_edge_has_line_number(self):
        spec = LANGUAGE_REGISTRY['python']
        source_bytes = PYTHON_WITH_CALLS.encode('utf-8')
        parser = get_parser(spec.ts_language)
        tree = parser.parse(source_bytes)

        symbols = parse_file(PYTHON_WITH_CALLS, 'test.py', 'python')
        main_sym = next(s for s in symbols if s.name == 'main')
        main_node = _find_node(tree.root_node, main_sym.byte_offset, main_sym.byte_length)
        calls = extract_calls(main_node, spec, source_bytes)

        for call in calls:
            assert call.line > 0

    def test_strips_self_prefix(self):
        code = '''
class Foo:
    def bar(self):
        self.baz()
        self.qux()
'''
        spec = LANGUAGE_REGISTRY['python']
        source_bytes = code.encode('utf-8')
        parser = get_parser(spec.ts_language)
        tree = parser.parse(source_bytes)

        symbols = parse_file(code, 'test.py', 'python')
        bar_sym = next(s for s in symbols if s.name == 'bar')
        bar_node = _find_node(tree.root_node, bar_sym.byte_offset, bar_sym.byte_length)
        calls = extract_calls(bar_node, spec, source_bytes)
        callee_names = {c.callee_name for c in calls}
        # self. should be stripped
        assert 'baz' in callee_names
        assert 'qux' in callee_names


def _find_node(root, start_byte, byte_length):
    """Find AST node by byte range."""
    end_byte = start_byte + byte_length
    if root.start_byte == start_byte and root.end_byte == end_byte:
        return root
    for child in root.children:
        if child.start_byte <= start_byte and child.end_byte >= end_byte:
            result = _find_node(child, start_byte, byte_length)
            if result:
                return result
    return None
