from __future__ import annotations

import ast
from collections.abc import Iterable
from dataclasses import dataclass


@dataclass(frozen=True)
class NotebookSignals:
    imports: set[str]
    functions: set[str]
    has_bigquery: bool
    has_sql_strings: bool
    has_dataframe_ops: bool


def extract_signals(cells: Iterable[str]) -> NotebookSignals:
    imports: set[str] = set()
    functions: set[str] = set()
    has_bigquery = False
    has_sql_strings = False
    has_dataframe_ops = False

    for cell in cells:
        try:
            tree = ast.parse(cell)
        except SyntaxError:
            # Skip cells that aren't parseable (magics, partials, etc.)
            continue

        for node in ast.walk(tree):
            if isinstance(node, ast.Import):
                for n in node.names:
                    imports.add(n.name.split(".")[0])
            elif isinstance(node, ast.ImportFrom):
                if node.module:
                    imports.add(node.module.split(".")[0])
            elif isinstance(node, ast.FunctionDef):
                functions.add(node.name)
            elif isinstance(node, ast.Attribute):
                if node.attr in {"to_dataframe", "query", "cut", "rename", "map"}:
                    has_dataframe_ops = True
            elif isinstance(node, ast.Constant) and isinstance(node.value, str):
                s = node.value
                if "SELECT" in s.upper() and "FROM" in s.upper():
                    has_sql_strings = True

        if ("bigquery" in imports or "google" in imports) and "bigquery" in cell:
            has_bigquery = True

    return NotebookSignals(
        imports=imports,
        functions=functions,
        has_bigquery=has_bigquery,
        has_sql_strings=has_sql_strings,
        has_dataframe_ops=has_dataframe_ops,
    )
