chunk_python_code

import ast
from langchain.schema import Document 

def chunk_python_code_with_metadata(source_code, reference):
    """
    Entry point method to process the Python file.
    It invokes the iterate_ast function.
    """
    documents = []
    print(f"Processing file: {reference}")

    iterate_ast(source_code, documents, reference)

    # Determine usage based on the reference path
    if reference.startswith("kadi_apy/lib/"):
        usage = "library"
    elif reference.startswith("kadi_apy/cli/"):
        usage = "cli_library"
    elif reference.startswith("doc/"):
        usage = "doc"
    else:
        usage = "undefined" 

    # Add metadata for usage to all documents
    for doc in documents:
        doc.metadata["reference"] = reference
        doc.metadata["usage"] = usage  # Add the determined usage metadata

    return documents


def iterate_ast(source_code, documents, reference):
    """
    Parses the AST of the given Python file and delegates
    handling to specific methods based on node types.
    """
    # Parse the source code into an abstract syntax tree (AST)
    tree = ast.parse(source_code, filename=reference)

    # Gather all top-level imports for later use
    imports_dict = extract_imports(tree)

    first_level_nodes = list(ast.iter_child_nodes(tree))

    # Check if there are no first-level nodes
    if not first_level_nodes:
        handle_no_first_level_node_found(documents, source_code, imports_dict, reference)
        return


    all_imports = all(isinstance(node, (ast.Import, ast.ImportFrom)) for node in first_level_nodes)
    if all_imports:
        handle_first_level_imports_only(documents, source_code, imports_dict, reference)
    

    # Iterate over first-level nodes
    for first_level_node in ast.iter_child_nodes(tree):
        if isinstance(first_level_node, ast.ClassDef):
            handle_first_level_class(first_level_node, documents, source_code, imports_dict)
        elif isinstance(first_level_node, ast.FunctionDef):
            handle_first_level_func(first_level_node, documents, source_code, imports_dict)
        elif isinstance(first_level_node, ast.Assign):
            handle_first_level_assign(first_level_node, documents, source_code, imports_dict)

            

def  handle_first_level_imports_only(documents, source_code, imports_dict, reference):
        # Check if the file path before ".py" is "__init__"
    if reference.endswith("__init__.py"):
        type = "__init__-file"
    else:
        type = "undefined"

    # Create and store a Document with the full source code
    doc = Document(
        page_content=source_code,
        metadata={
            "type": type,
            "imports": imports_dict
        }
    )
    documents.append(doc)


def extract_imports(tree):
    """
    Extracts all import statements from the AST tree and organizes them
    into a dictionary keyed by their fully qualified names for later analysis.
    """
    imports_dict = {}
    for node in ast.walk(tree):
        if isinstance(node, ast.Import):
            for alias in node.names:
                imports_dict[alias.name] = alias.name
        elif isinstance(node, ast.ImportFrom):
            module = node.module if node.module else ""
            for alias in node.names:
                full_name = f"{module}.{alias.name}" if module else alias.name
                imports_dict[alias.name] = full_name
    return imports_dict

def analyze_imports(node, imports_dict):
    """
    Analyzes the given node's body and signature to find relevant imports.
    """
    relevant_imports = set()
    for sub_node in ast.walk(node):
        if isinstance(sub_node, ast.Name) and sub_node.id in imports_dict:
            relevant_imports.add(imports_dict[sub_node.id])
    return list(relevant_imports)

def handle_not_yet_defined_first_level_cases(documents, source_code, imports_dict):
    if source_code: 
        doc = Document(
            page_content=source_code,
            metadata={
                "type": "undefined",
                "imports": imports_dict
            }
        )
        documents.append(doc)



def handle_no_first_level_node_found(documents, source_code, imports_dict, reference):
    """
    Handles cases where no top-level nodes are found in the AST.
    Stores the full content (likely comments) in a Document object
    with metadata indicating type 'no code' or 'init' based on the reference.
    """
    # Check if the file path before ".py" is "__init__"
    if reference.endswith("__init__.py"):
        type = "__init__-file"
    else:
        type = "undefined"

    # Create and store a Document with the full source code
    doc = Document(
        page_content=source_code,
        metadata={
            "type": type,
            "imports": imports_dict
        }
    )
    documents.append(doc)


def handle_first_level_assign(assign_node, documents, source_code, imports_dict):
    """
    Handles assignment statements at the first level of the AST by storing them
    in a Document object with metadata, including relevant imports.
    """
    # Extract assignment source code
    assign_start_line = assign_node.lineno
    assign_end_line = assign_node.end_lineno
    assign_source = '\n'.join(source_code.splitlines()[assign_start_line-1:assign_end_line])

    # Extract relevant imports for this assignment
    assign_imports = analyze_imports(assign_node, imports_dict)
    
    # Create and store Document for the assignment
    doc = Document(
        page_content=assign_source,
        metadata={
            "type": "Assign",
            "imports": assign_imports
        }
    )
    documents.append(doc)

def handle_first_level_class(class_node, documents, source_code, imports_dict):

    """
    Handles classes at the first level of the AST by storing them
    in a Document object with metadata, including relevant imports.
    """
    # Extract class source code
    class_start_line = class_node.lineno

    # Find the line where the first function (def) starts or the next top-level node
    class_body_lines = [child.lineno for child in class_node.body if isinstance(child, ast.FunctionDef)]
    class_end_line = min(class_body_lines, default=class_node.end_lineno) - 1  # Use `-1` to exclude the next node

    # Generate the class source code
    class_source = '\n'.join(source_code.splitlines()[class_start_line-1:class_end_line])
    
    # Extract relevant imports for this class
    class_imports = analyze_imports(class_node, imports_dict)
    
    # Create and store Document for the class
    doc = Document(
        page_content=class_source,
        metadata={
            "type": "class",
            "class": class_node.name,
            "visibility": "public",
            "imports": class_imports,  
        }
    )
    documents.append(doc)


    # Handle methods within the class
    for second_level_node in ast.iter_child_nodes(class_node):
        if isinstance(second_level_node, ast.FunctionDef):
            # Extract method source code
            method_start_line = (
                second_level_node.decorator_list[0].lineno
                if second_level_node.decorator_list else second_level_node.lineno
            )
            method_end_line = second_level_node.end_lineno
            method_source = '\n'.join(source_code.splitlines()[method_start_line-1:method_end_line])
            
            # Determine visibility metadata
            visibility = "internal" if second_level_node.name.startswith("_") else "public"
            # Extract relevant imports for this method
            method_imports = analyze_imports(second_level_node, imports_dict)
            
            # Create and store Document

            doc = Document(
                page_content=method_source,
                metadata={
                    "type": "method",
                    "method": second_level_node.name,
                    "visibility": "visibility",
                    "imports": method_imports,  
                    "class": class_node.name
                }
            )
            documents.append(doc)

def handle_first_level_func(function_node, documents, source_code, imports_dict):
    """
    Handles functions at the first level of the AST by storing them
    in a Document object with metadata, including relevant imports.
    """
    # Extract function source code
    function_start_line = (
        function_node.decorator_list[0].lineno
        if function_node.decorator_list else function_node.lineno
    )
    function_end_line = function_node.end_lineno
    function_source = '\n'.join(source_code.splitlines()[function_start_line-1:function_end_line])
    
    # Determine visibility metadata
    visibility = "internal" if function_node.name.startswith("_") else "public"
    
    # Check if the function is a CLI command (e.g., decorated with `@apy_command`)
    is_command = any(
        decorator.id == "apy_command"  # Check decorator name
        for decorator in function_node.decorator_list
        if hasattr(decorator, "id")  # Ensure the decorator has an identifier
    )
    
    # Extract relevant imports for this function
    function_imports = analyze_imports(function_node, imports_dict)
    
    # Create and store Document
    
    if is_command:
        doc = Document(
            page_content=function_source,
            metadata={
                "type": "command",
                "command": function_node.name,
                "visibility": "public",
                "imports": function_imports
            }
        )
    else:
          doc = Document(
            page_content=function_source,
            metadata={
                "type": "function",
                "method": function_node.name,
                "visibility": visibility,
                "imports": function_imports
            }
        )          
    documents.append(doc)


# Example usage
#file_path = r"C:\Users\Anwender\Downloads\exampleScript.py"

#with open(file_path, "r", encoding="utf-8") as file:
#        source_code = file.read()
#chunkPythonFiles(source_code, file_path)


import os

def process_folder(folder_path):
    # Initialize a counter for the number of Python files
    python_file_count = 0
    docsT = []
    # Walk through all subdirectories and files in the folder
    for root, _, files in os.walk(folder_path):
        for file_name in files:
            # Create the full file path
            file_path = os.path.join(root, file_name)
            #print(file_path)

            # Ensure it's a Python file
            if file_name.endswith(".py"):
                python_file_count += 1  # Increment the counter
                with open(file_path, "r", encoding="utf-8") as file:
                    source_code = file.read()
                    print(file_name)
                    
                # Call your function
                docs = chunkPythonFiles(source_code, file_path)

                print("HWHWHWWHWHWHWH!:" ,len(docs))
                docsT.extend(docs)
    # Print the total number of Python files processed
    print(f"Total Python files processed: {python_file_count}")
    print(f"Total docs files processed: {len(docsT)}")