# parser.py
import ast

def get_category(node):
    """Determine the category of an AST node."""
    if isinstance(node, (ast.Import, ast.ImportFrom)):
        return 'import'
    elif isinstance(node, (ast.Assign, ast.AnnAssign, ast.AugAssign)):
        return 'assignment'
    elif isinstance(node, ast.FunctionDef):
        return 'function'
    elif isinstance(node, ast.AsyncFunctionDef):
        return 'async_function'
    elif isinstance(node, ast.ClassDef):
        return 'class'
    elif isinstance(node, ast.If):
        return 'if'
    elif isinstance(node, ast.While):
        return 'while'
    elif isinstance(node, ast.For):
        return 'for'
    elif isinstance(node, ast.Expr):
        return 'expression'
    else:
        return 'other'

def create_vector(category, level, location, total_lines, parent_path):
    """Create a vector representation including hierarchy info."""
    category_map = {
        'import': 1, 'assignment': 2, 'function': 3, 'async_function': 4, 'class': 5,
        'if': 6, 'while': 7, 'for': 8, 'expression': 9, 'spacer': 10, 'other': 11
    }
    category_id = category_map.get(category, 0)
    start_line, end_line = location
    # Extend vector with parent path length and a hash of the path for uniqueness
    parent_depth = len(parent_path)
    parent_hash = hash(' -> '.join(parent_path)) % 1000  # Simple hash for vector
    return [
        category_id,
        level,
        start_line / total_lines,
        end_line / total_lines,
        parent_depth,
        parent_hash / 1000  # Normalized hash (0.0 to 1.0)
    ]

def parse_node(node, lines, prev_end, level=0, total_lines=None, parent_path=None, counters=None):
    """Recursively parse an AST node with full hierarchy tracking."""
    if total_lines is None:
        total_lines = len(lines)
    if parent_path is None:
        parent_path = []
    if counters is None:
        counters = {'if': 0, 'while': 0, 'for': 0, 'function': 0, 'class': 0}  # Track counts per type

    parts = []
    start_line = getattr(node, 'lineno', prev_end + 1)
    end_line = getattr(node, 'end_lineno', start_line)
    category = get_category(node)

    # Assign a node ID for relevant categories
    node_id = ''
    if category in counters:
        counters[category] += 1
        node_id = f"{category.capitalize()}[{counters[category]}]"

    # Handle spacers before the node
    if start_line > prev_end + 1:
        spacer_lines = lines[prev_end:start_line - 1]
        spacer_vector = create_vector('spacer', level, (prev_end + 1, start_line - 1), total_lines, parent_path)
        parts.append({
            'category': 'spacer',
            'source': ''.join(spacer_lines),
            'location': (prev_end + 1, start_line - 1),
            'level': level,
            'vector': spacer_vector,
            'parent_path': ' -> '.join(parent_path) if parent_path else 'Top-Level',
            'node_id': ''
        })

    # Capture the node's source
    stmt_lines = lines[start_line - 1:end_line]
    current_path = parent_path + ([node_id] if node_id else [])
    node_vector = create_vector(category, level, (start_line, end_line), total_lines, current_path)
    parts.append({
        'category': category,
        'source': ''.join(stmt_lines),
        'location': (start_line, end_line),
        'level': level,
        'vector': node_vector,
        'parent_path': ' -> '.join(parent_path) if parent_path else 'Top-Level',
        'node_id': node_id
    })

    # Process nested nodes (body and orelse)
    nested_prev_end = end_line - 1
    for attr in ('body', 'orelse'):
        if hasattr(node, attr) and getattr(node, attr):
            for child in getattr(node, attr):
                child_parts = parse_node(child, lines, nested_prev_end, level + 1, total_lines, current_path, counters)
                parts.extend(child_parts)
                nested_prev_end = child_parts[-1]['location'][1]

    return parts

def parse_python_code(code):
    """Parse Python code string and return parts with hierarchy and vectors."""
    lines = code.splitlines(keepends=True)
    total_lines = len(lines)
    try:
        tree = ast.parse(code)
    except SyntaxError:
        return [{'category': 'error', 'source': 'Invalid Python code', 'location': (1, 1), 'level': 0, 'vector': [0, 0, 1.0, 1.0, 0, 0], 'parent_path': 'Top-Level', 'node_id': ''}]

    parts = []
    prev_end = 0

    for stmt in tree.body:
        stmt_parts = parse_node(stmt, lines, prev_end, total_lines=total_lines)
        parts.extend(stmt_parts)
        prev_end = stmt_parts[-1]['location'][1]

    if prev_end < total_lines:
        remaining_lines = lines[prev_end:]
        spacer_vector = create_vector('spacer', 0, (prev_end + 1, total_lines + 1), total_lines, [])
        parts.append({
            'category': 'spacer',
            'source': ''.join(remaining_lines),
            'location': (prev_end + 1, total_lines + 1),
            'level': 0,
            'vector': spacer_vector,
            'parent_path': 'Top-Level',
            'node_id': ''
        })

    return parts