from itertools import chain

from core.languages.python.testgen.preprocessing.dataclass_constructor_notes import get_dataclass_constructor_notes
from core.languages.python.testgen.preprocessing.torch_tensor_limit import get_tensor_size_note

# Preprocessing functions that analyze code context and return notes
# Each function takes test_context (str) and returns a list of notes (list[str])
_PREPROCESSING_FUNCTIONS = [get_tensor_size_note, get_dataclass_constructor_notes]


def preprocessing_testgen_pipeline(test_context: str) -> list[str]:
    """Apply preprocessing functions in a pipeline.

    Analyzes the code context and generates notes to help the LLM understand
    the code better. For example, adding a note to limit tensor sizes when
    torch is detected, or providing explicit dataclass constructor signatures.

    Args:
        test_context: The source code context to analyze.

    Returns:
        List of notes to include in the LLM prompt.

    """
    return list(chain.from_iterable(func(test_context) for func in _PREPROCESSING_FUNCTIONS))
