src/post_process.py

import re
from typing import List

_FROM_TABLE_RE = re.compile(r'\bFROM\s+table\b', flags=re.IGNORECASE)
_AGG_RE = re.compile(r'\b(COUNT|MAX|MIN|AVG|SUM)\s+([A-Za-z0-9_\s`"]+?)(?=\sFROM|\sWHERE|,|$)', flags=re.IGNORECASE)
_WHERE_RE = re.compile(r'\bWHERE\b\s+(.*)', flags=re.IGNORECASE)
_AND_OR_RE = re.compile(r'\s+\b(AND|OR)\b\s+', flags=re.IGNORECASE)
_COND_RE = re.compile(r'(.*?)\s*(=|>|<|>=|<=|!=)\s*(.*)')

def replace_table(sql: str, table_name: str = "data") -> str:
    """Replaces 'FROM table' (case-insensitive) with the actual table name."""
    return _FROM_TABLE_RE.sub(f'FROM "{table_name}"', sql)

def add_agg_parentheses(sql: str) -> str:
    """
    Wraps column names after SQL aggregations in parentheses.
    E.g., 'SELECT COUNT Player' -> 'SELECT COUNT(Player)'
    """
    def repl(match: re.Match) -> str:
        func = match.group(1)
        col = match.group(2).strip()
        # If it already has parens, leave it alone
        if col.startswith('(') and col.endswith(')'):
            return f"{func}{col}"
        return f"{func}({col})"
    
    return _AGG_RE.sub(repl, sql)

def quote_multiword_columns(sql: str, columns: List[str]) -> str:
    """
    Wraps multi-word column names in double quotes to prevent SQLite tokenization errors.
    Only quotes columns that need it: those containing spaces or special characters.
    Single-word, simple columns (e.g., 'Age', 'Team') are left unquoted.
    """
    # Filter to columns that actually need quoting (contain spaces or non-alphanumeric/underscore chars)
    needs_quoting = [col for col in columns if ' ' in col or not re.match(r'^[A-Za-z_][A-Za-z0-9_]*$', col)]
    
    # Sort by length descending so we replace longer names first (e.g. 'Player Name' before 'Player')
    sorted_cols = sorted(needs_quoting, key=len, reverse=True)
    
    for col in sorted_cols:
        # Find the column name NOT already surrounded by quotes/backticks.
        # We can't use \b word boundaries because columns may contain special chars
        # like parentheses or slashes (e.g., "Revenue ($)", "Win/Loss %").
        # Instead, match at string boundaries or after/before SQL-meaningful characters.
        pattern = rf'(?<![`"])(?<![A-Za-z0-9_]){re.escape(col)}(?![A-Za-z0-9_`"])'
        sql = re.sub(pattern, f'"{col}"', sql, flags=re.IGNORECASE)
        
    return sql

def quote_string_values(sql: str, columns: List[str], types: List[str]) -> str:
    """
    Checks column schema; if text, converts `= Value` to `= 'Value'`.
    Robust against AND/OR statements.
    """
    where_match = _WHERE_RE.search(sql)
    if not where_match:
        return sql
        
    where_clause = where_match.group(1)
    parts = _AND_OR_RE.split(where_clause)
    
    type_map = {c.lower(): t for c, t in zip(columns, types)}
    processed_parts = []
    
    for part in parts:
        if part.upper() in ['AND', 'OR']:
            processed_parts.append(part.upper())
            continue
            
        cond_match = _COND_RE.match(part)
        if cond_match:
            col_ref = cond_match.group(1).strip()
            op = cond_match.group(2)
            val = cond_match.group(3).strip()
            
            clean_col = col_ref.replace('"', '').replace('`', '').lower()
            
            if clean_col in type_map and type_map[clean_col] == 'text':
                if not (val.startswith("'") and val.endswith("'")):
                    val = f"'{val}'"
            
            processed_parts.append(f"{col_ref} {op} {val}")
        else:
            processed_parts.append(part)
            
    new_where = " ".join(processed_parts)
    return sql[:where_match.start()] + "WHERE " + new_where

def post_process_sql(sql: str, columns: List[str], types: List[str], table_name: str = "data") -> str:
    """Runs the full post-processing pipeline on the raw T5 output."""
    sql = replace_table(sql, table_name)
    sql = quote_multiword_columns(sql, columns)
    sql = add_agg_parentheses(sql)
    sql = quote_string_values(sql, columns, types)
    return sql