stack/dev/scripts/generate_expressions.py

"""Introspect DuckDB views and generate narwhals expression modules.

Each view's SQL is parsed into its fundamental operations (select, filter,
join, group_by, window, union) and emitted as a narwhals-native function
in src/aco/express/<schema>.py.

Usage:
    python generate_expressions.py [--db PATH] [--out DIR]
"""

from __future__ import annotations

import argparse
import re
from collections import defaultdict
from pathlib import Path

import duckdb

# ── SQL parsing helpers ─────────────────────────────────────────────────────

REF_RE = re.compile(r"aco\.(\w+)\.(\w+)", re.IGNORECASE)
ALIAS_RE = re.compile(r"\bAS\s+(\w+)\b", re.IGNORECASE)
FROM_TABLE_RE = re.compile(r"FROM\s+aco\.(\w+)\.(\w+)(?:\s+AS\s+(\w+))?", re.IGNORECASE)
JOIN_RE = re.compile(
    r"((?:INNER|LEFT|RIGHT|FULL|CROSS)\s+)?JOIN\s+aco\.(\w+)\.(\w+)"
    r"(?:\s+AS\s+(\w+))?\s+ON\s*\(\((.+?)\)\)",
    re.IGNORECASE | re.DOTALL,
)
# Fallback: looser pattern for JOINs with single-paren ON
JOIN_LOOSE_RE = re.compile(
    r"((?:INNER|LEFT|RIGHT|FULL|CROSS)\s+)?JOIN\s+aco\.(\w+)\.(\w+)"
    r"(?:\s+AS\s+(\w+))?\s+ON\s*\((.+?)\)\)",
    re.IGNORECASE | re.DOTALL,
)
WHERE_RE = re.compile(
    r"\bWHERE\s+\((.+?)\)(?:\s+GROUP|\s+ORDER|\s*$)", re.IGNORECASE | re.DOTALL
)


def _extract_where_clause(sql: str) -> str | None:
    """Extract the WHERE clause content using paren-balanced parsing.

    DuckDB wraps WHERE conditions in parens: WHERE (condition);
    We track paren depth to find the matching close-paren.
    """
    upper = sql.upper()
    # Find WHERE keyword at top level (not inside parens)
    idx = 0
    depth = 0
    where_pos = -1
    while idx < len(upper):
        if upper[idx] == "(":
            depth += 1
        elif upper[idx] == ")":
            depth -= 1
        elif depth == 0 and upper[idx : idx + 6] == "WHERE ":
            where_pos = idx
        idx += 1

    if where_pos == -1:
        return None

    # Find the opening paren after WHERE
    rest = sql[where_pos + 6 :].lstrip()
    if not rest.startswith("("):
        return None

    # Track parens to find the matching close
    depth = 0
    for i, ch in enumerate(rest):
        if ch == "(":
            depth += 1
        elif ch == ")":
            depth -= 1
            if depth == 0:
                # Extract content between outermost parens
                return rest[1:i]
    return None


GROUP_BY_RE = re.compile(
    r"GROUP\s+BY\s+(.+?)(?:\s+HAVING|\s+ORDER|\s*;|\s*\)|\s*$)", re.IGNORECASE
)
SELECT_COLS_RE = re.compile(
    r"SELECT\s+(DISTINCT\s+)?(.+?)\s+FROM\s+", re.IGNORECASE | re.DOTALL
)
WINDOW_RE = re.compile(r"(\w+)\s*\(\s*\)\s+OVER\s*\((.+?)\)", re.IGNORECASE)
UNION_RE = re.compile(r"\bUNION\s+ALL\b", re.IGNORECASE)
CTE_RE = re.compile(r"WITH\s+(.+?)\)\s*SELECT", re.IGNORECASE | re.DOTALL)


def ref_to_param(schema: str, table: str) -> str:
    """Convert schema.table to a function parameter name."""
    return f"{schema}__{table}".replace("-", "_")


def ref_to_table_import(schema: str, table: str) -> str:
    """Return the aco.table import class name."""
    parts = f"{schema}__{table}".split("_")
    return "".join(p.capitalize() for p in parts if p)


def to_func_name(table: str) -> str:
    """Convert a view name to a snake_case function name."""
    return table.replace("-", "_").lstrip("_")


def extract_deps(sql: str) -> list[tuple[str, str]]:
    """Extract all aco.schema.table references from SQL."""
    refs = REF_RE.findall(sql)
    return sorted(set(refs))


def outer_query(sql: str) -> str:
    """Extract the outermost SELECT from a CREATE VIEW ... AS [WITH ...] SELECT ...

    Strips the CREATE VIEW prefix and any CTE block, returning only the
    final SELECT statement that defines the view's output.
    """
    # Strip CREATE VIEW ... AS
    m = re.search(r"\bAS\s+", sql, re.IGNORECASE)
    if m:
        sql = sql[m.end() :]

    # Strip CTE: WITH name AS (...), name AS (...) SELECT ...
    # Find the last top-level SELECT (not inside parens)
    depth = 0
    last_select = 0
    i = 0
    upper = sql.upper()
    while i < len(upper):
        if upper[i] == "(":
            depth += 1
        elif upper[i] == ")":
            depth -= 1
        elif depth == 0 and upper[i : i + 7] == "SELECT ":
            last_select = i
        i += 1
    return sql[last_select:]


def _strip_subqueries(sql: str) -> str:
    """Remove parenthesized subqueries to expose only top-level SQL keywords."""
    result = []
    depth = 0
    for ch in sql:
        if ch == "(":
            depth += 1
        elif ch == ")":
            depth = max(0, depth - 1)
        elif depth == 0:
            result.append(ch)
    return "".join(result)


def classify_sql(sql: str) -> list[str]:
    """Classify a SQL statement into operation types based on the outer query."""
    ops = []
    oq = outer_query(sql)
    top = _strip_subqueries(oq).upper()
    full = sql.upper()
    oq_upper = oq.upper()
    if (
        "SELECT *" in oq_upper
        and "JOIN" not in top
        and "WHERE" not in top
        and "UNION" not in full
    ):
        return ["passthrough"]
    if "UNION ALL" in full:
        ops.append("union")
    if "JOIN" in top:
        ops.append("join")
    if "WHERE" in top:
        ops.append("filter")
    if "GROUP BY" in top:
        ops.append("group_by")
    if "OVER (" in oq_upper:
        ops.append("window")
    if not ops:
        ops.append("select")
    return ops


def extract_select_columns(sql: str) -> list[str]:
    """Extract the column expressions from a SELECT clause.

    Uses paren-balanced parsing to find the top-level FROM keyword,
    handling subqueries like CAST((SELECT id FROM ...) AS VARCHAR).
    """
    upper = sql.upper()
    # Find SELECT keyword
    sel_m = re.search(r"\bSELECT\s+(DISTINCT\s+)?", sql, re.IGNORECASE)
    if not sel_m:
        return []
    start = sel_m.end()

    # Find top-level FROM (not inside parens)
    depth = 0
    end = len(sql)
    i = start
    while i < len(sql):
        if sql[i] == "(":
            depth += 1
        elif sql[i] == ")":
            depth -= 1
        elif (
            depth == 0
            and upper[i : i + 5] == "FROM "
            and (i == 0 or not sql[i - 1].isalpha())
        ):
            end = i
            break
        i += 1

    cols_str = sql[start:end].strip()
    if not cols_str:
        return []

    # Split on commas not inside parens
    depth = 0
    parts = []
    current = []
    for ch in cols_str:
        if ch == "(":
            depth += 1
        elif ch == ")":
            depth -= 1
        elif ch == "," and depth == 0:
            parts.append("".join(current).strip())
            current = []
            continue
        current.append(ch)
    if current:
        parts.append("".join(current).strip())
    return parts


def col_expr_to_nw(expr: str) -> str:
    """Convert a SQL column expression to a narwhals-ish string representation."""
    expr = expr.strip()
    # Handle alias: ... AS name
    alias_m = re.search(r'\bAS\s+(?:"(\w+)"|(\w+))\s*$', expr, re.IGNORECASE)
    alias = (alias_m.group(1) or alias_m.group(2)) if alias_m else None
    if alias_m:
        expr = expr[: alias_m.start()].strip()

    # CAST((SELECT col FROM table WHERE (id = 'val')) AS TYPE) — scalar subquery constant
    scalar_m = re.match(
        r"(?:TRY_)?CAST\(\(SELECT\s+.+?\s+FROM\s+.+?\s+WHERE\s+\(.+?=\s*'([^']+)'\)\)\s+AS\s+\w+\)",
        expr,
        re.IGNORECASE,
    )
    if scalar_m:
        val = scalar_m.group(1)
        if alias:
            return f'nw.lit("{val}").alias("{alias}")'
        return f'nw.lit("{val}")'

    # CAST(x AS TYPE) / TRY_CAST(x AS TYPE)
    cast_m = re.match(r"(?:TRY_)?CAST\((.+?)\s+AS\s+.+\)", expr, re.IGNORECASE)
    if cast_m:
        inner = cast_m.group(1).strip()
        expr = inner  # drop the cast for narwhals, typing handled by schema

    # date_diff('unit', a, b)
    dd_m = re.match(r"date_diff\('(\w+)',\s*(.+?),\s*(.+)\)", expr, re.IGNORECASE)
    if dd_m:
        _unit, a, b = dd_m.group(1), dd_m.group(2).strip(), dd_m.group(3).strip()
        if alias:
            return f'({_col_ref(b)} - {_col_ref(a)}).alias("{alias}")'
        return f"{_col_ref(b)} - {_col_ref(a)}"

    # substring(col, start, len)
    sub_m = re.match(
        r'(?:main\.)?"?substring"?\((.+?),\s*(\d+),\s*(\d+)\)', expr, re.IGNORECASE
    )
    if sub_m:
        col = sub_m.group(1).strip()
        start = int(sub_m.group(2)) - 1  # 0-indexed
        length = int(sub_m.group(3))
        ref = _col_ref(col)
        result = f"{ref}.str.slice({start}, {length})"
        if alias:
            return f'{result}.alias("{alias}")'
        return result

    # min/max/count/sum(col)
    agg_m = re.match(r"(min|max|count|sum)\((.+?)\)", expr, re.IGNORECASE)
    if agg_m:
        func = agg_m.group(1).lower()
        col = agg_m.group(2).strip()
        result = f"{_col_ref(col)}.{func}()"
        if alias:
            return f'{result}.alias("{alias}")'
        return result

    # row_number() OVER (...)
    if "row_number" in expr.lower():
        # Parse PARTITION BY and ORDER BY from the OVER clause
        over_m = re.search(r"OVER\s*\((.+)\)", expr, re.IGNORECASE)
        if over_m:
            over_body = over_m.group(1)
            part_m = re.search(
                r"PARTITION\s+BY\s+(.+?)(?:\s+ORDER|\s*$)", over_body, re.IGNORECASE
            )
            order_m = re.search(r"ORDER\s+BY\s+(.+?)$", over_body, re.IGNORECASE)
            comment_parts = []
            if part_m:
                comment_parts.append(f"partition_by={part_m.group(1).strip()}")
            if order_m:
                comment_parts.append(f"order_by={order_m.group(1).strip()}")
            comment = ", ".join(comment_parts)
            if alias:
                return f'nw.lit(1).alias("{alias}")  # row_number: {comment}'
            return f"nw.lit(1)  # row_number: {comment}"
        if alias:
            return f'nw.lit(1).alias("{alias}")  # TODO: row_number window'
        return "# TODO: row_number window"

    # dense_rank() OVER (...)
    if "dense_rank" in expr.lower():
        over_m = re.search(r"OVER\s*\((.+)\)", expr, re.IGNORECASE)
        comment = over_m.group(1).strip() if over_m else ""
        if alias:
            return f'nw.lit(1).alias("{alias}")  # dense_rank: {comment}'
        return f"nw.lit(1)  # dense_rank: {comment}"

    # lower(COALESCE(a, b)) or lower(col)
    lower_m = re.match(r"lower\((.+)\)", expr, re.IGNORECASE)
    if lower_m:
        inner = lower_m.group(1).strip()
        # Check if inner is COALESCE
        coal_inner = re.match(r"COALESCE\((.+)\)", inner, re.IGNORECASE)
        if coal_inner:
            args = [a.strip() for a in coal_inner.group(1).split(",")]
            ref = _col_ref(args[0])
            for a in args[1:]:
                ref = f"{ref}.fill_null({_col_ref(a)})"
            result = f"{ref}.str.to_lowercase()"
        else:
            result = f"{_col_ref(inner)}.str.to_lowercase()"
        if alias:
            return f'{result}.alias("{alias}")'
        return result

    # COALESCE(a, b, ...)
    coal_m = re.match(r"COALESCE\((.+)\)", expr, re.IGNORECASE)
    if coal_m:
        args = [a.strip() for a in coal_m.group(1).split(",")]
        if len(args) == 2:
            a_ref = _col_ref(args[0])
            b_ref = _col_ref(args[1])
            result = f"{a_ref}.fill_null({b_ref})"
        else:
            result = _col_ref(args[0])
            for a in args[1:]:
                result = f"{result}.fill_null({_col_ref(a)})"
        if alias:
            return f'{result}.alias("{alias}")'
        return result

    # CASE WHEN ... THEN ... ELSE ... END
    if expr.upper().startswith("CASE "):
        # Extract a simple CASE WHEN (cond) THEN val ELSE val END
        case_m = re.match(
            r"CASE\s+WHEN\s+\(\((.+?)\)\)\s+THEN\s+\((\w+)\)\s+ELSE\s+(\S+)\s+END",
            expr,
            re.IGNORECASE,
        )
        if case_m:
            cond, then_val, else_val = case_m.group(1), case_m.group(2), case_m.group(3)
            # Parse condition: col = val
            cond_m = re.match(r"(\w+)\s*=\s*(\d+)", cond)
            if cond_m:
                col, val = cond_m.group(1), cond_m.group(2)
                result = f'nw.when(nw.col("{col}") == {val}, {then_val}).otherwise({else_val})'
                if alias:
                    return f'{result}.alias("{alias}")'
                return result
        # Fallback for complex CASE
        safe = expr.replace('"', "'")[:80]
        if alias:
            return f'nw.lit(None).alias("{alias}")  # CASE: {safe}'
        return f"nw.lit(None)  # CASE: {safe}"

    # DISTINCT keyword — strip it and recurse
    if expr.upper().startswith("DISTINCT "):
        return col_expr_to_nw(expr[9:].strip() + (f" AS {alias}" if alias else ""))

    # NULL literal
    if expr.upper() == "NULL":
        if alias:
            return f'nw.lit(None).alias("{alias}")'
        return "nw.lit(None)"

    # String literal: 'value'
    str_m = re.match(r"^'(.+)'$", expr)
    if str_m:
        val = str_m.group(1)
        if alias:
            return f'nw.lit("{val}").alias("{alias}")'
        return f'nw.lit("{val}")'

    # left(col, n) — string slice
    left_m = re.match(r"'?left'?\((.+?),\s*(\d+)\)", expr, re.IGNORECASE)
    if left_m:
        col = left_m.group(1).strip()
        n = int(left_m.group(2))
        result = f"{_col_ref(col)}.str.slice(0, {n})"
        if alias:
            return f'{result}.alias("{alias}")'
        return result

    # Select all
    if expr.strip() == "*":
        return "nw.all()"

    # SQL quoted identifier: "name" -> name
    quoted_m = re.match(r'^"(\w+)"$', expr)
    if quoted_m:
        col = quoted_m.group(1)
        if alias:
            return f'nw.col("{col}").alias("{alias}")'
        return f'nw.col("{col}")'

    # col AS 'alias' (single-quoted alias in SQL)
    sq_alias_m = re.match(r"(.+?)\s+AS\s+'([^']+)'\s*$", expr, re.IGNORECASE)
    if sq_alias_m:
        inner = sq_alias_m.group(1).strip()
        sq_alias = sq_alias_m.group(2)
        ref = _col_ref(inner)
        return f'{ref}.alias("{sq_alias}")'

    # Simple column ref: table.col or col
    ref = _col_ref(expr)
    if alias and alias != expr.split(".")[-1]:
        return f'{ref}.alias("{alias}")'
    return ref


def _col_ref(expr: str) -> str:
    """Turn a SQL column reference into nw.col(...)."""
    expr = expr.strip()
    # Remove CAST / TRY_CAST wrapper
    cast_m = re.match(r"(?:TRY_)?CAST\((.+?)\s+AS\s+.+\)", expr, re.IGNORECASE)
    if cast_m:
        expr = cast_m.group(1).strip()
    # table.col -> just col
    if "." in expr and "(" not in expr:
        expr = expr.split(".")[-1]
    if expr.isidentifier():
        return f'nw.col("{expr}")'
    safe = expr.replace('"', "'")
    return f"nw.lit(None)  # complex: {safe}"


def _split_on_keyword(clause: str, keyword: str) -> list[str]:
    """Split a clause on a top-level keyword (AND/OR), respecting parentheses."""
    kw = f" {keyword} "
    kw_len = len(kw)
    parts = []
    depth = 0
    current: list[str] = []
    upper = clause.upper()
    i = 0
    while i < len(clause):
        if clause[i] == "(":
            depth += 1
            current.append(clause[i])
        elif clause[i] == ")":
            depth -= 1
            current.append(clause[i])
        elif depth == 0 and upper[i : i + kw_len] == kw:
            parts.append("".join(current).strip())
            current = []
            i += kw_len
            continue
        else:
            current.append(clause[i])
        i += 1
    if current:
        parts.append("".join(current).strip())
    return parts


def _split_and_clauses(clause: str) -> list[str]:
    return _split_on_keyword(clause, "AND")


def _split_or_clauses(clause: str) -> list[str]:
    return _split_on_keyword(clause, "OR")


def where_to_nw(where_clause: str) -> str:
    """Convert a WHERE clause to a narwhals filter expression."""
    clause = where_clause.strip()

    # Strip outer parens if balanced
    if clause.startswith("(") and clause.endswith(")"):
        depth = 0
        balanced_at_end = True
        for i, ch in enumerate(clause):
            if ch == "(":
                depth += 1
            elif ch == ")":
                depth -= 1
                if depth == 0 and i < len(clause) - 1:
                    balanced_at_end = False
                    break
        if balanced_at_end:
            clause = clause[1:-1].strip()

    # CAST('t' AS BOOLEAN) — always-true, skip it
    if re.match(r"CAST\('t'\s+AS\s+BOOLEAN\)", clause, re.IGNORECASE):
        return "nw.lit(True)"

    # NOT col (boolean negation)
    m = re.match(r"NOT\s+(\w+(?:\.\w+)?)\s*$", clause, re.IGNORECASE)
    if m:
        return f"~{_col_ref(m.group(1))}"

    # col IS NULL
    m = re.match(r"(\w+(?:\.\w+)?)\s+IS\s+NULL\s*$", clause, re.IGNORECASE)
    if m:
        return f"{_col_ref(m.group(1))}.is_null()"

    # col IS NOT NULL
    m = re.match(r"(\w+(?:\.\w+)?)\s+IS\s+NOT\s+NULL\s*$", clause, re.IGNORECASE)
    if m:
        return f"{_col_ref(m.group(1))}.is_not_null()"

    # col NOT IN ('a', 'b', ...)
    not_in_m = re.match(
        r"(\w+(?:\.\w+)?)\s+NOT\s+IN\s+\((.+?)\)\s*$", clause, re.IGNORECASE
    )
    if not_in_m:
        col = not_in_m.group(1)
        vals = not_in_m.group(2)
        return f"~{_col_ref(col)}.is_in([{vals}])"

    # col IN ('a', 'b', ...)
    in_m = re.match(r"(\w+(?:\.\w+)?)\s+IN\s+\((.+?)\)\s*$", clause, re.IGNORECASE)
    if in_m:
        col = in_m.group(1)
        vals = in_m.group(2)
        return f"{_col_ref(col)}.is_in([{vals}])"

    # col = 'string'
    m = re.match(r"(\w+(?:\.\w+)?)\s*=\s*'([^']+)'\s*$", clause)
    if m:
        return f'{_col_ref(m.group(1))} == "{m.group(2)}"'

    # col != 'string'
    m = re.match(r"(\w+(?:\.\w+)?)\s*!=\s*'([^']+)'\s*$", clause)
    if m:
        return f'{_col_ref(m.group(1))} != "{m.group(2)}"'

    # lower(col) != 'string'  /  lower(col) = 'string'
    m = re.match(
        r"lower\((\w+(?:\.\w+)?)\)\s*(!=|=)\s*'([^']+)'\s*$", clause, re.IGNORECASE
    )
    if m:
        col, op, val = m.group(1), m.group(2), m.group(3)
        nw_op = "!=" if op == "!=" else "=="
        return f'{_col_ref(col)}.str.to_lowercase() {nw_op} "{val}"'

    # date_diff('unit', a, b) >= N  (comparison with date_diff)
    dd_m = re.match(
        r"date_diff\('(\w+)',\s*(.+?),\s*(.+?)\)\s*(>=|<=|>|<|=|!=)\s*(\d+)\s*$",
        clause,
        re.IGNORECASE,
    )
    if dd_m:
        unit, a, b = dd_m.group(1), dd_m.group(2).strip(), dd_m.group(3).strip()
        op, val = dd_m.group(4), dd_m.group(5)
        nw_op = {">=": ">=", "<=": "<=", ">": ">", "<": "<", "=": "==", "!=": "!="}.get(
            op, op
        )
        # Strip CAST wrappers for cleaner output
        a = re.sub(r"CAST\((.+?)\s+AS\s+\w+\)", r"\1", a, flags=re.IGNORECASE)
        b = re.sub(r"CAST\((.+?)\s+AS\s+\w+\)", r"\1", b, flags=re.IGNORECASE)
        return f"({_col_ref(b)} - {_col_ref(a)}).dt.total_days() / 365 {nw_op} {val}  # date_diff {unit}"

    # substring(col, start, len) = 'val' or IN (...)
    sub_m = re.match(
        r'(?:main\.)?"?substring"?\((.+?),\s*(\d+),\s*(\d+)\)\s*(=|!=|IN)\s*(.+?)\s*$',
        clause,
        re.IGNORECASE,
    )
    if sub_m:
        col = sub_m.group(1).strip()
        start = int(sub_m.group(2)) - 1
        length = int(sub_m.group(3))
        op = sub_m.group(4).upper()
        right = sub_m.group(5).strip()
        ref = f"{_col_ref(col)}.str.slice({start}, {length})"
        if op == "IN":
            vals = right.strip("()")
            return f"{ref}.is_in([{vals}])"
        elif op == "=":
            return f"{ref} == {right}"
        else:
            return f"{ref} != {right}"

    # CAST(expr AS TYPE) <= val — unwrap CAST for comparison
    cast_cmp_m = re.match(
        r"CAST\((.+?)\s+AS\s+\w+\)\s*(>=|<=|>|<|=|!=)\s*(.+?)\s*$",
        clause,
        re.IGNORECASE,
    )
    if cast_cmp_m:
        inner = cast_cmp_m.group(1).strip()
        op = cast_cmp_m.group(2)
        right = cast_cmp_m.group(3).strip()
        nw_op = {">=": ">=", "<=": "<=", ">": ">", "<": "<", "=": "==", "!=": "!="}.get(
            op, op
        )
        return f"{_col_ref(inner)} {nw_op} {right}"

    # Generic comparison: expr OP val (simple col or function result)
    m = re.match(r"(\w+(?:\.\w+)?)\s*(>=|<=|!=|=|>|<)\s*(.+?)\s*$", clause)
    if m:
        left, op, right = m.group(1).strip(), m.group(2), m.group(3).strip()
        nw_op = {">=": ">=", "<=": "<=", ">": ">", "<": "<", "=": "==", "!=": "!="}.get(
            op, op
        )
        return f"{_col_ref(left)} {nw_op} {right}"

    # AND-combined: split on top-level AND (higher precedence than OR)
    if " AND " in clause.upper():
        parts = _split_and_clauses(clause)
        if len(parts) > 1:
            nw_parts = []
            comments = []
            for part in parts:
                sub = where_to_nw(part.strip())
                code, comment = (
                    (sub.split("  # ", 1) + [None])[:2]
                    if "  # " in sub
                    else (sub, None)
                )
                if code.lstrip().startswith("#"):
                    safe = clause.replace('"', "'")[:80]
                    return f"# complex filter: {safe}"
                nw_parts.append(f"({code})")
                if comment:
                    comments.append(comment)
            result = " & ".join(nw_parts)
            if comments:
                result += "  # " + "; ".join(comments)
            return result

    # OR-combined: split on top-level OR
    if " OR " in clause.upper():
        or_parts = _split_or_clauses(clause)
        if len(or_parts) > 1:
            nw_parts = []
            comments = []
            for part in or_parts:
                sub = where_to_nw(part.strip())
                code, comment = (
                    (sub.split("  # ", 1) + [None])[:2]
                    if "  # " in sub
                    else (sub, None)
                )
                if code.lstrip().startswith("#"):
                    safe = clause.replace('"', "'")[:80]
                    return f"# complex filter: {safe}"
                nw_parts.append(f"({code})")
                if comment:
                    comments.append(comment)
            result = " | ".join(nw_parts)
            if comments:
                result += "  # " + "; ".join(comments)
            return result

    safe = clause.replace('"', "'")[:80]
    return f"# complex filter: {safe}"


# ── Code generation ─────────────────────────────────────────────────────────


def generate_function(
    schema: str,
    view_name: str,
    sql: str,
    deps: list[tuple[str, str]],
    ops: list[str],
) -> tuple[str, set[str]]:
    """Generate a single narwhals expression function for a view."""
    func_name = to_func_name(view_name)
    imports: set[str] = set()

    # Build parameter list from dependencies
    params = []
    for dep_schema, dep_table in deps:
        param = ref_to_param(dep_schema, dep_table)
        params.append(f"{param}: FrameT")

    if not params:
        params = ["df: FrameT"]

    param_str = ", ".join(params)

    # Build docstring
    dep_names = [f"{s}.{t}" for s, t in deps]
    doc_deps = ", ".join(dep_names) if dep_names else "none"
    doc_ops = ", ".join(ops)

    lines: list[str] = []
    lines.append("@nw.narwhalify")
    lines.append(f"def {func_name}({param_str}) -> FrameT:")
    lines.append(f'    """Build {schema}.{view_name}')
    lines.append("")
    lines.append(f"    Operations: {doc_ops}")
    lines.append(f"    Dependencies: {doc_deps}")
    lines.append('    """')

    if "passthrough" in ops:
        p = params[0].split(":")[0]
        lines.append(f"    return {p}")
        return "\n".join(lines), imports

    oq = outer_query(sql)

    if "union" in ops:
        _gen_union_body(lines, deps, sql)
    elif "join" in ops:
        _gen_join_body(lines, deps, oq, ops)
    elif "group_by" in ops:
        _gen_group_by_body(lines, deps, oq)
    elif "filter" in ops:
        _gen_filter_body(lines, deps, oq)
    elif "window" in ops:
        _gen_window_body(lines, deps, oq)
    else:
        _gen_select_body(lines, deps, oq)

    return "\n".join(lines), imports


def _extract_group_by_keys(sql: str) -> list[str]:
    """Extract GROUP BY key expressions using paren-balanced parsing."""
    upper = sql.upper()
    gp = -1
    depth = 0
    for i in range(len(upper)):
        if sql[i] == "(":
            depth += 1
        elif sql[i] == ")":
            depth -= 1
        elif depth == 0 and upper[i : i + 9] == "GROUP BY ":
            gp = i + 9
            break
    if gp == -1:
        return []

    # Find the end: HAVING, ORDER, ;, or end of string (at top level)
    depth = 0
    end = len(sql)
    for i in range(gp, len(sql)):
        if sql[i] == "(":
            depth += 1
        elif sql[i] == ")":
            depth -= 1
        elif depth == 0:
            if upper[i : i + 7] == "HAVING " or upper[i : i + 6] == "ORDER ":
                end = i
                break
            if sql[i] == ";":
                end = i
                break

    raw = sql[gp:end]

    # Split on top-level commas
    parts = []
    depth = 0
    current: list[str] = []
    for ch in raw:
        if ch == "(":
            depth += 1
        elif ch == ")":
            depth -= 1
        elif ch == "," and depth == 0:
            parts.append("".join(current).strip())
            current = []
            continue
        current.append(ch)
    if current:
        parts.append("".join(current).strip())

    return parts


def _clean_key(k: str) -> str:
    """Strip table aliases and return just the column name."""
    k = k.strip()
    if "." in k and "(" not in k:
        k = k.split(".")[-1]
    return k


def _is_clean_key(k: str) -> bool:
    """Check if a parsed key looks like a valid column identifier."""
    return (
        bool(k)
        and k.replace("_", "").isalnum()
        and not any(c in k for c in "()\"' \t\n;=")
    )


def _is_safe_expr(expr: str) -> bool:
    """Check that a generated expression is safe to emit as Python code."""
    code = expr
    if "  # " in code:
        code = code.split("  # ", 1)[0]
    # Too long is suspicious
    if len(code) > 200:
        return False
    # Unbalanced parens
    if code.count("(") != code.count(")"):
        return False
    # Stray SQL keywords
    for kw in [
        "SELECT ",
        "FROM ",
        "WHERE ",
        "GROUP BY",
        "ORDER BY",
        "AS TIMESTAMP",
        "AS BIGINT",
        "AS INTEGER",
        "AS DOUBLE",
        "AS BOOLEAN",
        "AS VARCHAR",
        "INNER JOIN",
        "LEFT JOIN",
        "PARTITION BY",
        "UNION",
        "CASE ",
        "CAST(",
    ]:
        if kw in code.upper():
            return False
    return True


def _emit_exprs(lines: list[str], exprs: list[str]) -> None:
    """Emit a list of nw expressions into a multi-line call, safely."""
    for expr in exprs:
        if not _is_safe_expr(expr):
            # Replace with a safe placeholder
            safe_comment = expr.replace('"', "'").replace("\n", " ")[:80]
            lines.append(f"        # TODO: {safe_comment}")
            continue
        if "  # " in expr:
            code, comment = expr.split("  # ", 1)
            lines.append(f"        # {comment}")
            lines.append(f"        {code},")
        else:
            lines.append(f"        {expr},")


def _primary_dep(deps: list[tuple[str, str]]) -> str:
    if deps:
        return ref_to_param(deps[0][0], deps[0][1])
    return "df"


def _gen_select_body(lines: list[str], deps: list[tuple[str, str]], sql: str) -> None:
    """Generate body for a simple select (column projection)."""
    p = _primary_dep(deps)
    cols = extract_select_columns(sql)
    nw_cols = [col_expr_to_nw(c) for c in cols if "pipeline_last_run" not in c.lower()]
    if nw_cols:
        lines.append(f"    return {p}.select(")
        _emit_exprs(lines, nw_cols)
        lines.append("    )")
    else:
        lines.append(f"    return {p}")


def _gen_filter_body(lines: list[str], deps: list[tuple[str, str]], sql: str) -> None:
    p = _primary_dep(deps)
    clause = _extract_where_clause(sql)
    if clause:
        nw_filter = where_to_nw(clause)
        # Separate code from informational comments (# date_diff, etc.)
        if "  # " in nw_filter:
            code_part, comment_part = nw_filter.split("  # ", 1)
        else:
            code_part, comment_part = nw_filter, None
        # Check if the code part starts with "# " — that's a stub/TODO
        is_stub = code_part.lstrip().startswith("#")
        if not is_stub and _is_safe_expr(code_part):
            if comment_part:
                lines.append(f"    # {comment_part}")
            lines.append(f"    return {p}.filter(")
            lines.append(f"        {code_part}")
            lines.append("    )")
        else:
            safe = clause.replace('"', "'")[:80]
            lines.append(f"    # TODO: complex filter: {safe}")
            lines.append(f"    return {p}")
    else:
        lines.append("    # WHERE clause could not be parsed")
        lines.append(f"    return {p}")


def _extract_join_keys(on_clause: str) -> tuple[list[str], list[str]] | None:
    """Extract left/right join key column names from an ON clause.

    Handles simple col=col and also CAST(col AS type) = CAST(col AS type).
    """
    # Try simple: alias.col = alias.col
    parts = re.findall(r"(\w+)\.(\w+)\s*=\s*(\w+)\.(\w+)", on_clause)
    if parts:
        return [p[1] for p in parts], [p[3] for p in parts]
    # Try: col = col (no alias)
    parts = re.findall(r"\b(\w+)\s*=\s*(\w+)\b", on_clause)
    if parts:
        return [p[0] for p in parts], [p[1] for p in parts]
    return None


def _gen_join_body(
    lines: list[str], deps: list[tuple[str, str]], sql: str, ops: list[str]
) -> None:
    p = _primary_dep(deps)
    lines.append(f"    result = {p}")

    # Find JOIN clauses — try strict regex first, then loose
    matches = list(JOIN_RE.finditer(sql))
    if not matches:
        matches = list(JOIN_LOOSE_RE.finditer(sql))

    for jm in matches:
        join_type = (jm.group(1) or "inner").strip().lower()
        j_schema, j_table = jm.group(2), jm.group(3)
        on_clause = jm.group(5)

        j_param = ref_to_param(j_schema, j_table)
        how = "inner" if "inner" in join_type else join_type.split()[0]

        keys = _extract_join_keys(on_clause)
        if keys:
            left_keys, right_keys = keys
            if left_keys == right_keys:
                keys_str = ", ".join(f'"{k}"' for k in left_keys)
                lines.append(
                    f'    result = result.join({j_param}, on=[{keys_str}], how="{how}")'
                )
            else:
                lk = ", ".join(f'"{k}"' for k in left_keys)
                rk = ", ".join(f'"{k}"' for k in right_keys)
                lines.append(
                    f'    result = result.join({j_param}, left_on=[{lk}], right_on=[{rk}], how="{how}")'
                )
        else:
            safe_on = on_clause.replace('"', "'").replace("\n", " ")[:60]
            lines.append(f"    # join ON: {safe_on}")
            lines.append(
                f'    result = result.join({j_param}, how="{how}", suffix="_r")'
            )

    # Apply filter if present
    if "filter" in ops:
        clause = _extract_where_clause(sql)
        if clause:
            nw_filter = where_to_nw(clause)
            if "  # " in nw_filter:
                code_part, comment_part = nw_filter.split("  # ", 1)
            else:
                code_part, comment_part = nw_filter, None
            is_stub = code_part.lstrip().startswith("#")
            if not is_stub and _is_safe_expr(code_part):
                if comment_part:
                    lines.append(f"    # {comment_part}")
                lines.append(f"    result = result.filter({code_part})")
            else:
                safe = clause.replace('"', "'")[:80]
                lines.append(f"    # TODO: complex filter: {safe}")

    # Apply group_by if present
    if "group_by" in ops:
        raw_keys = _extract_group_by_keys(sql)
        keys = [_clean_key(k) for k in raw_keys if _is_clean_key(_clean_key(k))]
        if keys:
            keys_str = ", ".join(f'"{k}"' for k in keys)
            lines.append(f"    result = result.group_by({keys_str}).agg()")
        else:
            lines.append("    # TODO: complex GROUP BY")

    lines.append("    return result")


def _gen_group_by_body(lines: list[str], deps: list[tuple[str, str]], sql: str) -> None:
    p = _primary_dep(deps)
    raw_keys = _extract_group_by_keys(sql)
    if not raw_keys:
        lines.append(f"    return {p}")
        return

    # Clean keys: skip CASE expressions and other complex expressions
    keys = []
    for rk in raw_keys:
        ck = _clean_key(rk)
        if _is_clean_key(ck):
            keys.append(ck)
    if not keys:
        # All keys were complex (e.g. CASE) — use column names from SELECT as fallback
        cols = extract_select_columns(sql)
        for c in cols:
            c_clean = _clean_key(
                c.split(" AS ")[-1].strip() if " AS " in c.upper() else c
            )
            if _is_clean_key(c_clean):
                keys.append(c_clean)
                break
    if not keys:
        lines.append("    # TODO: complex GROUP BY could not be parsed")
        lines.append(f"    return {p}")
        return
    keys_str = ", ".join(f'"{k}"' for k in keys)

    # Extract aggregate columns
    cols = extract_select_columns(sql)
    agg_cols = []
    for c in cols:
        c_stripped = (
            c.strip().split(".")[-1] if "." in c and "(" not in c else c.strip()
        )
        if c_stripped in keys:
            continue
        if "pipeline_last_run" in c.lower():
            continue
        nw_expr = col_expr_to_nw(c)
        agg_cols.append(nw_expr)

    if agg_cols:
        lines.append(f"    return {p}.group_by({keys_str}).agg(")
        _emit_exprs(lines, agg_cols)
        lines.append("    )")
    else:
        lines.append(f"    return {p}.group_by({keys_str}).agg()")


def _gen_window_body(lines: list[str], deps: list[tuple[str, str]], sql: str) -> None:
    p = _primary_dep(deps)
    cols = extract_select_columns(sql)
    nw_cols = []
    for c in cols:
        if "pipeline_last_run" in c.lower():
            continue
        nw_cols.append(col_expr_to_nw(c))

    lines.append(f"    return {p}.with_columns(")
    _emit_exprs(lines, nw_cols)
    lines.append("    )")


def _gen_union_body(lines: list[str], deps: list[tuple[str, str]], sql: str) -> None:
    params = [ref_to_param(s, t) for s, t in deps]
    if len(params) >= 2:
        lines.append("    return nw.concat([")
        for p in params:
            lines.append(f"        {p},")
        lines.append("    ])")
    elif params:
        lines.append(f"    return {params[0]}")
    else:
        lines.append("    return df")


# ── Module generation ───────────────────────────────────────────────────────


def generate_schema_module(
    views: list[tuple[str, str, str]],
    schema: str,
) -> str:
    """Generate a full Python module for all views in a schema."""
    header_imports = {
        "from __future__ import annotations",
        "import narwhals as nw",
        "from narwhals.typing import FrameT",
    }

    functions: list[str] = []

    for _, view_name, sql in views:
        deps = extract_deps(sql)
        ops = classify_sql(sql)
        func_code, extra_imports = generate_function(schema, view_name, sql, deps, ops)
        header_imports |= extra_imports
        functions.append(func_code)

    sorted_imports = sorted(header_imports)
    header = "\n".join(sorted_imports)
    body = "\n\n\n".join(functions)

    return f"{header}\n\n\n{body}\n"


# ── Main ────────────────────────────────────────────────────────────────────


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Generate narwhals expression modules from DuckDB views"
    )
    from conf import path as _conf_path

    parser.add_argument(
        "--db",
        default=str(_conf_path("db.aco")),
        help="Path to DuckDB file",
    )
    parser.add_argument("--out", default="../src/aco/express", help="Output directory")
    args = parser.parse_args()

    db_path = Path(args.db)
    out_dir = Path(args.out)
    out_dir.mkdir(parents=True, exist_ok=True)

    con = duckdb.connect(str(db_path), read_only=True)

    all_views = con.execute("""
        SELECT schema_name, view_name, sql
        FROM duckdb_views()
        WHERE schema_name NOT IN ('information_schema', 'pg_catalog')
        ORDER BY schema_name, view_name
    """).fetchall()

    # Group by schema
    by_schema: dict[str, list[tuple[str, str, str]]] = defaultdict(list)
    for schema, name, sql in all_views:
        by_schema[schema].append((schema, name, sql))

    all_modules: list[str] = []

    for schema in sorted(by_schema):
        safe_name = schema.replace("-", "_")
        views = by_schema[schema]
        module_path = out_dir / f"{safe_name}.py"
        code = generate_schema_module(views, schema)
        module_path.write_text(code)
        print(f"Wrote {module_path} ({len(views)} expressions)")
        all_modules.append(safe_name)

    # Write __init__.py
    init_lines = ['"""Generated narwhals expressions for DuckDB view logic."""\n']
    for mod in sorted(all_modules):
        init_lines.append(f"from . import {mod}")
    init_lines.append("")
    (out_dir / "__init__.py").write_text("\n".join(init_lines))
    print(f"\nWrote {out_dir / '__init__.py'}")

    con.close()

    total = sum(len(v) for v in by_schema.values())
    print(f"\nDone: {len(by_schema)} schemas, {total} expressions -> {out_dir}/")


if __name__ == "__main__":
    main()