src/sql_retriever.py

"""SQL retriever: execute validated queries against the SQLite knowledge base."""

import os
import re
import sqlite3
import time
from pathlib import Path


_DANGEROUS_KEYWORDS = re.compile(
    r"\b(INSERT|UPDATE|DELETE|DROP|ALTER|CREATE|ATTACH|DETACH|PRAGMA|"
    r"LOAD_EXTENSION|UNION|VACUUM|REINDEX|SAVEPOINT|RELEASE|"
    r"ROLLBACK|BEGIN|COMMIT|GRANT|REVOKE|EXPLAIN|WITH)\b",
    re.IGNORECASE,
)


def _strip_quoted(sql: str) -> str:
    """Remove quoted strings and identifiers for safe keyword checking."""
    # Remove single-quoted strings (handles '' escaping)
    result = re.sub(r"'(?:[^']|'')*'", "", sql)
    # Remove double-quoted identifiers
    result = re.sub(r'"[^"]*"', "", result)
    return result


def _strip_sql_comments(sql: str) -> str:
    """Remove SQL comments (single-line -- and multi-line /* */) from a query.

    Respects single-quoted string literals: ``--`` and ``/* */`` inside
    quotes are preserved.
    """
    result = []
    i = 0
    n = len(sql)
    while i < n:
        # Single-quoted string literal — copy verbatim (handles '' escaping)
        if sql[i] == "'":
            result.append("'")
            i += 1
            while i < n:
                if sql[i] == "'" and i + 1 < n and sql[i + 1] == "'":
                    result.append("''")
                    i += 2
                elif sql[i] == "'":
                    result.append("'")
                    i += 1
                    break
                else:
                    result.append(sql[i])
                    i += 1
        # Block comment
        elif sql[i] == '/' and i + 1 < n and sql[i + 1] == '*':
            i += 2
            while i < n:
                if sql[i] == '*' and i + 1 < n and sql[i + 1] == '/':
                    i += 2
                    break
                i += 1
        # Line comment
        elif sql[i] == '-' and i + 1 < n and sql[i + 1] == '-':
            i += 2
            while i < n and sql[i] != '\n':
                i += 1
        else:
            result.append(sql[i])
            i += 1
    return ''.join(result)


def _validate_sql(sql: str) -> bool:
    """Validate that a SQL string is a safe SELECT query.

    Layer 1 of SQL injection protection (Layer 2 is the read-only connection).
    Rejects non-SELECT statements, semicolons, dangerous keywords,
    and subqueries (multiple SELECT keywords).
    """
    stripped = sql.strip()
    if not stripped:
        return False
    if ";" in stripped:
        return False
    if not stripped.upper().startswith("SELECT "):
        return False
    unquoted = _strip_quoted(stripped)
    if _DANGEROUS_KEYWORDS.search(unquoted):
        return False
    # Block access to SQLite system tables
    if re.search(r'\bsqlite_(master|schema|temp_master|temp_schema)\b', stripped, re.IGNORECASE):
        return False
    # Block subqueries: reject if more than one SELECT keyword
    if len(re.findall(r'\bSELECT\b', unquoted, re.IGNORECASE)) > 1:
        return False
    return True


def _get_db_path(cfg: dict) -> str:
    """Resolve the SQLite database path from config."""
    sql_db_dir = cfg.get("paths", {}).get("sql_db", "sql_db")
    if not os.path.isabs(sql_db_dir):
        project_root = Path(__file__).resolve().parent.parent
        sql_db_dir = os.path.join(str(project_root), sql_db_dir)
    return os.path.join(sql_db_dir, "knowledge_base.db")


def execute_sql_query(sql_query: str, cfg: dict) -> list[dict]:
    """Execute a validated SELECT query against the knowledge base SQLite DB.

    Args:
        sql_query: A SQL SELECT query (generated by the LLM).
        cfg: App config dict.

    Returns:
        List of row dicts, or empty list on any error.
    """
    # Strip markdown code fences if present (LLMs often wrap SQL in ```sql...```)
    sql_query = sql_query.strip()
    if sql_query.startswith("```"):
        sql_query = sql_query.split("\n", 1)[-1].rsplit("```", 1)[0].strip()

    # Strip comments before validation AND execution
    sql_query = _strip_sql_comments(sql_query).strip()

    # Strip trailing semicolons — LLMs add them ~90% of the time.
    # Must happen AFTER comment removal but BEFORE validation, so that
    # _validate_sql() still rejects internal semicolons (statement chaining).
    sql_query = sql_query.rstrip(";").rstrip()

    if not _validate_sql(sql_query):
        return []

    db_path = _get_db_path(cfg)
    if not os.path.exists(db_path):
        return []

    max_rows = cfg.get("sql", {}).get("max_rows", 200)

    try:
        conn = sqlite3.connect(f"file:{db_path}?mode=ro", uri=True, timeout=5)
        try:
            conn.row_factory = sqlite3.Row
            # Abort long-running queries after 10 seconds
            _start = time.monotonic()
            def _progress_check():
                if time.monotonic() - _start > 10:
                    return 1  # non-zero = abort
                return 0
            conn.set_progress_handler(_progress_check, 10000)
            cursor = conn.execute(sql_query)
            rows = cursor.fetchmany(max_rows)
            return [dict(row) for row in rows]
        finally:
            conn.close()
    except Exception:
        return []


def format_sql_results_as_context(
    rows: list[dict], sql_query: str, source_file: str,
    table_info: dict = None, max_rows: int = None,
) -> str:
    """Format SQL result rows as context for the verification pipeline.

    Uses [CHUNK-SQL-NNN] IDs for internal anchoring, consistent with
    CHUNK-LOCAL and CHUNK-WEB patterns.

    Args:
        rows: Query result rows as dicts.
        sql_query: The SQL query that produced these rows.
        source_file: Original source file path.
        table_info: Schema entry for the table (includes column descriptions).
        max_rows: The max_rows limit used in the query. When len(rows) >= max_rows,
            a truncation notice is appended so the LLM knows data may be incomplete.
    """
    if not rows:
        return ""

    row_label = f"Rows returned: {len(rows)}"
    if max_rows is not None and len(rows) >= max_rows:
        row_label += f" (truncated — more rows may exist, limit was {max_rows})"

    parts = [
        "=== SQL Query Results (PRIMARY — from local dataset) ===\n",
        f"Query: {sql_query}",
        f"Source: {source_file}",
        f"{row_label}\n",
    ]

    # Include table and column descriptions so the LLM understands the data
    if table_info:
        table_desc = table_info.get("table_description", "")
        if table_desc:
            parts.append(f"Dataset description: {table_desc}\n")
        col_descs = {
            c.get("name", "unknown"): c.get("description", "")
            for c in table_info.get("columns", [])
        }
        desc_lines = [f"  {name}: {desc}" for name, desc in col_descs.items() if desc]
        if desc_lines:
            parts.append("Column descriptions:")
            parts.extend(desc_lines)
            parts.append("")

    for i, row in enumerate(rows, 1):
        fields = ", ".join(
            f"{k} = {v if v is not None else 'N/A'}" for k, v in row.items()
        )
        parts.append(f"[CHUNK-SQL-{i:03d}] {fields}")
    return "\n".join(parts)


_FUZZY_STOPWORDS = frozenset({
    "the", "and", "for", "are", "but", "not", "you", "all",
    "can", "had", "her", "was", "one", "our", "out", "has",
    "what", "how", "who", "which", "when", "where", "with",
    "from", "that", "this", "than", "then", "they", "been",
    "south", "north", "east", "west", "new", "old", "united",
    "democratic", "republic", "people", "state", "states",
    "islamic", "federal", "kingdom",
})


def make_fuzzy_query(sql: str, word_level: bool = False) -> str | None:
    """Convert exact equality on text values to LIKE fuzzy matching.

    When ``word_level`` is False (default), transforms
    ``column = 'value'`` into ``column LIKE '%value%'``
    (phrase-level matching).

    When ``word_level`` is True, transforms
    ``column = 'South Korea'`` into
    ``(column LIKE '%Korea%')``
    by extracting significant content words (dropping stopwords and
    directional/geopolitical qualifiers). This handles cases where
    "South Korea" is stored as "Korea, Republic of".

    Returns the modified query, or None if no substitutions were made.
    """
    pattern = re.compile(
        r"""("[^"]+"\s*|[A-Za-z_]\w*\s*)=\s*'([^']+)'""",
    )

    if not word_level:
        def _phrase_replace(m):
            col_part = m.group(1)
            value = m.group(2).replace("'", "''")
            # Escape LIKE wildcards in the value
            value = value.replace("%", "\\%").replace("_", "\\_")
            return f"{col_part}LIKE '%{value}%' ESCAPE '\\'"
        new_sql, count = pattern.subn(_phrase_replace, sql)
        if count == 0:
            return None
        return new_sql

    # Word-level: extract significant words from each value
    def _word_replace(m):
        col_part = m.group(1)
        value = m.group(2).replace("'", "''")
        # Escape LIKE wildcards in the value
        value = value.replace("%", "\\%").replace("_", "\\_")
        words = [
            w.replace("'", "''").replace("%", "\\%").replace("_", "\\_")
            for w in re.split(r'\W+', m.group(2))
            if len(w) >= 3 and w.lower() not in _FUZZY_STOPWORDS
        ]
        esc = " ESCAPE '\\'"
        if not words:
            # Fallback: use the original value as phrase LIKE
            return f"{col_part}LIKE '%{value}%'{esc}"
        if len(words) == 1:
            return f"{col_part}LIKE '%{words[0]}%'{esc}"
        # Multiple words: join with AND on the same column
        conditions = [f"{col_part}LIKE '%{w}%'{esc}" for w in words]
        return "(" + " AND ".join(conditions) + ")"

    new_sql, count = pattern.subn(_word_replace, sql)
    if count == 0:
        return None
    return new_sql


def _lookup_source_file(table_name: str, schema: dict) -> str:
    """Look up the original source file for a SQL table name."""
    entry = schema.get(table_name)
    if entry:
        return entry.get("source_file", table_name)
    return table_name


def build_schema_summary(schema: dict) -> str:
    """Build a compact schema summary for injection into the QU prompt.

    Includes column descriptions (from codebook or LLM inference),
    sample values, and statistics so the LLM knows what data looks like
    for correct WHERE clauses and aggregation queries.
    """
    if not schema:
        return ""

    lines = ["Available SQL tables:"]
    for table_name, info in schema.items():
        row_count = info.get("row_count", 0)
        source = info.get("source_file", "")
        table_desc = info.get("table_description", "")

        header = f"\n- {table_name} ({row_count} rows, from {source}):"
        if table_desc:
            header += f"\n  Description: {table_desc}"
        lines.append(header)

        for c in info.get("columns", []):
            col_type = c.get("type", "TEXT")
            stats = c.get("stats", {})
            samples = c.get("sample", [])
            desc = c.get("description", "")

            # Build metadata string
            meta_parts = [col_type]
            unique_count = stats.get("unique_count")
            col_min = stats.get("min")
            col_max = stats.get("max")

            if col_min is not None and col_max is not None:
                meta_parts.append(f"range {col_min}\u2013{col_max}")
            if unique_count is not None:
                meta_parts.append(f"{unique_count} unique")

            meta_str = ", ".join(meta_parts)

            # Build sample string
            sample_str = ""
            if samples:
                quoted = ", ".join(f'"{s}"' for s in samples[:5])
                sample_str = f" e.g. {quoted}"

            col_name = c.get("name", "unknown")
            line = f"    {col_name} ({meta_str}){sample_str}"
            if desc:
                line += f" \u2014 {desc}"
            lines.append(line)
    return "\n".join(lines)