#!/usr/bin/env python3
"""Text-to-SQL pipeline demo: schema-aware prompting with validation.

Demonstrates the full pipeline: schema grounding -> SQL generation -> validation
-> self-correction. Uses sqlite3 for validation (parse + schema check).
"""
from __future__ import annotations

import re
import sqlite3


# Simulated LLM (replace with real API call in production)
def llm_generate_sql(prompt: str) -> str:
    """
    In production, this would call an LLM API.
    Here we simulate with rule-based templates for demo purposes.
    """
    prompt_lower = prompt.lower()

    if "all orders" in prompt_lower and "us" in prompt_lower and "customer" in prompt_lower:
        return """SELECT c.name, o.total, o.status
FROM orders o
JOIN customers c ON c.id = o.customer_id
WHERE c.country = 'US'
ORDER BY o.total DESC
LIMIT 10"""

    if "count" in prompt_lower and "orders" in prompt_lower and "status" in prompt_lower:
        return """SELECT status, COUNT(*) as order_count
FROM orders
GROUP BY status
ORDER BY order_count DESC"""

    # Simulate a hallucinated column
    if "revenue" in prompt_lower:
        return """SELECT SUM(revenue) FROM orders"""  # BUG: no 'revenue' column

    return """SELECT * FROM nonexistent_table"""  # deliberate error for demo


# --- Schema grounding ---

SCHEMA = """
Tables:
- customers (id INT PK, name TEXT, country TEXT, signup_date TEXT)
- orders (id INT PK, customer_id INT FK->customers.id, total REAL, status TEXT, created_at TEXT)

Columns:
- customers.country: two-letter country code (US, UK, DE, FR)
- orders.status: one of 'pending', 'paid', 'shipped', 'delivered', 'cancelled'
- orders.total: order total in USD

Join: orders.customer_id = customers.id
"""


def schema_prompt(question: str) -> str:
    return f"""Given this database schema:
{SCHEMA}

Generate a SQL query to answer: {question}

Return ONLY the SQL, no explanation. Use LIMIT for safety."""


# --- Validation ---

FORBIDDEN_KEYWORDS = ["DROP", "TRUNCATE", "DELETE", "INSERT", "UPDATE", "ALTER", "CREATE"]


def validate_sql(sql: str, conn: sqlite3.Connection) -> tuple[bool, str]:
    """Validate SQL: syntax, schema, and safety."""
    sql_upper = sql.upper().strip()

    # Safety: reject forbidden keywords
    for kw in FORBIDDEN_KEYWORDS:
        if kw in sql_upper:
            return False, f"Rejected: contains forbidden keyword '{kw}'"

    # Schema: check referenced tables/columns exist
    try:
        cursor = conn.execute("SELECT name FROM sqlite_master WHERE type='table'")
        existing_tables = {row[0] for row in cursor.fetchall()}
    except Exception:
        existing_tables = {"customers", "orders"}

    # Simple table reference check
    table_refs = set(re.findall(r'FROM\s+(\w+)', sql, re.IGNORECASE))
    table_refs |= set(re.findall(r'JOIN\s+(\w+)', sql, re.IGNORECASE))
    for t in table_refs:
        if t.lower() not in existing_tables:
            return False, f"Table '{t}' does not exist"

    # Syntax: try EXPLAIN (sqlite parses but doesn't execute)
    try:
        conn.execute(f"EXPLAIN {sql}")
    except sqlite3.OperationalError as e:
        return False, f"SQL syntax error: {e}"

    return True, "OK"


def repair_prompt(sql: str, error: str, question: str) -> str:
    return f"""The previous SQL was invalid:
SQL: {sql}
Error: {error}

Generate a corrected SQL for: {question}
Return ONLY the SQL, no explanation."""


# --- Pipeline ---

def text_to_sql(question: str, conn: sqlite3.Connection, max_retries: int = 2) -> tuple[str | None, list[str]]:
    """Full pipeline: generate -> validate -> retry -> return safe SQL."""
    log: list[str] = []

    prompt = schema_prompt(question)
    sql = llm_generate_sql(prompt)
    log.append(f"Generated: {sql[:80]}")

    for attempt in range(1 + max_retries):
        ok, error = validate_sql(sql, conn)
        if ok:
            log.append(f"Validated OK (attempt {attempt})")
            # Add safety LIMIT if missing
            if "LIMIT" not in sql.upper():
                sql += " LIMIT 100"
                log.append("Added LIMIT 100")
            return sql, log

        log.append(f"Validation failed (attempt {attempt}): {error}")
        if attempt < max_retries:
            prompt = repair_prompt(sql, error, question)
            sql = llm_generate_sql(prompt)
            log.append(f"Retry {attempt + 1}: {sql[:80]}")

    log.append("Failed after all retries")
    return None, log


def main() -> None:
    conn = sqlite3.connect(":memory:")

    questions = [
        "Show all orders from US customers with their names and totals",
        "Count orders by status",
        "What is the total revenue?",  # triggers hallucinated column
        "Delete all orders",  # triggers safety reject
    ]

    for q in questions:
        print(f"\n{'='*60}")
        print(f"Q: {q}")
        sql, log = text_to_sql(q, conn)
        for entry in log:
            print(f"  {entry}")
        if sql:
            print(f"  Final SQL: {sql}")
        else:
            print(f"  RESULT: No valid SQL generated")

    conn.close()


if __name__ == "__main__":
    main()