from itertools import chain

from aiservice.common.markdown_utils import split_markdown_code
from core.languages.python.optimizer.context_utils.context_helpers import is_multi_context
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]


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))


def format_notes_markdown(notes: list[str] | set[str]) -> str:
    return f"Notes:\n{'\n'.join(notes)}" if notes else ""


def generate_notes(source_code: str) -> str:
    if is_multi_context(source_code):
        all_notes: set[str] = set()
        for scoped_code in split_markdown_code(source_code).values():
            all_notes.update(preprocessing_testgen_pipeline(scoped_code))
        return format_notes_markdown(all_notes)
    return format_notes_markdown(preprocessing_testgen_pipeline(source_code))
