import sqlite3
 
def create_sample_db(db_path: str = "ecommerce.db"):
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
 
    cursor.executescript("""
    CREATE TABLE IF NOT EXISTS products (
        id INTEGER PRIMARY KEY,
        name TEXT NOT NULL,
        category TEXT NOT NULL,
        price REAL NOT NULL,
        stock INTEGER DEFAULT 0
    );
 
    CREATE TABLE IF NOT EXISTS orders (
        id INTEGER PRIMARY KEY,
        customer_name TEXT NOT NULL,
        order_date DATE NOT NULL,
        total_amount REAL NOT NULL
    );
 
    CREATE TABLE IF NOT EXISTS order_items (
        id INTEGER PRIMARY KEY,
        order_id INTEGER REFERENCES orders(id),
        product_id INTEGER REFERENCES products(id),
        quantity INTEGER NOT NULL,
        unit_price REAL NOT NULL
    );
 
    INSERT OR IGNORE INTO products VALUES
        (1, 'Wireless Earbuds', 'Audio', 89.99, 150),
        (2, 'USB Hub 7-Port', 'Accessories', 29.99, 300),
        (3, 'Mechanical Keyboard', 'Input Devices', 119.99, 80),
        (4, '4K Webcam', 'Cameras', 64.99, 120),
        (5, 'Portable SSD 1TB', 'Storage', 79.99, 200);
    """)
 
    conn.commit()
    conn.close()

def extract_schema(db_path: str) -> str:
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
 
    cursor.execute(
        "SELECT name FROM sqlite_master WHERE type='table'"
    )
    tables = cursor.fetchall()
 
    schema_parts = []
    for (table_name,) in tables:
        cursor.execute(f"PRAGMA table_info({table_name})")
        columns = cursor.fetchall()
        col_defs = []
        for col in columns:
            col_defs.append(
                f"  {col[1]} {col[2]}"
                f"{'  PRIMARY KEY' if col[5] else ''}"
                f"{'  NOT NULL' if col[3] else ''}"
            )
        schema_parts.append(
            f"TABLE {table_name} (\n"
            + ",\n".join(col_defs)
            + "\n)"
        )
 
    conn.close()
    return "\n\n".join(schema_parts)

from google import genai
 
client = genai.Client(api_key="YOUR_API_KEY")
 
SYSTEM_PROMPT = """You are an SQL expert.
Convert the user's natural language question into a
SQLite-compatible SQL query based on the provided schema.
 
Rules:
1. Only generate SELECT statements (no INSERT/UPDATE/DELETE)
2. Return only the SQL query without any explanation
3. Only use table and column names that exist in the schema
4. For ambiguous questions, use the most reasonable interpretation
"""
 
def generate_sql(question: str, schema: str) -> str:
    prompt = f"""## Database Schema
{schema}
 
## User Question
{question}
 
## SQL Query
"""
    response = client.models.generate_content(
        model="gemini-2.5-flash",
        contents=prompt,
        config={
            "system_instruction": SYSTEM_PROMPT,
            "temperature": 0.0,
        },
    )
    sql = response.text.strip()
    # Strip code block formatting if present
    if sql.startswith("```"):
        sql = sql.split("\n", 1)[1]
    if sql.endswith("```"):
        sql = sql.rsplit("```", 1)[0]
    return sql.strip()

import re
 
BLOCKED_KEYWORDS = [
    "DROP", "DELETE", "UPDATE", "INSERT",
    "ALTER", "CREATE", "TRUNCATE", "EXEC",
]
 
def validate_sql(sql: str) -> bool:
    upper_sql = sql.upper().strip()
    if not upper_sql.startswith("SELECT"):
        return False
    for keyword in BLOCKED_KEYWORDS:
        if re.search(rf"\b{keyword}\b", upper_sql):
            return False
    return True
 
def execute_query(db_path: str, sql: str) -> dict:
    if not validate_sql(sql):
        raise ValueError(
            "Unsafe query detected"
        )
    conn = sqlite3.connect(db_path)
    conn.execute("PRAGMA query_only = ON")
    cursor = conn.cursor()
    cursor.execute(sql)
    columns = [desc[0] for desc in cursor.description]
    results = cursor.fetchall()
    conn.close()
    return {"columns": columns, "rows": results}

def ask_database(question: str, db_path: str = "ecommerce.db"):
    schema = extract_schema(db_path)
    sql = generate_sql(question, schema)
    print(f"Generated SQL: {sql}")
 
    result = execute_query(db_path, sql)
    print(f"Columns: {result['columns']}")
    for row in result["rows"]:
        print(row)
    return result
 
# Usage examples
ask_database("Show me products with more than 100 in stock, sorted by price descending")
# → SELECT name, price, stock FROM products
#   WHERE stock >= 100 ORDER BY price DESC
 
ask_database("How many products are there in each category?")
# → SELECT category, COUNT(*) as count
#   FROM products GROUP BY category

FEW_SHOT_EXAMPLES = """
Question: What is the most expensive product?
SQL: SELECT name, price FROM products ORDER BY price DESC LIMIT 1
 
Question: How many orders were placed in March 2026?
SQL: SELECT COUNT(*) FROM orders WHERE order_date BETWEEN '2026-03-01' AND '2026-03-31'
 
Question: Which customer has the highest total spending?
SQL: SELECT customer_name, SUM(total_amount) as total FROM orders GROUP BY customer_name ORDER BY total DESC LIMIT 1
"""

ANNOTATED_SCHEMA = """
TABLE products (
  id INTEGER PRIMARY KEY -- Product ID
  name TEXT NOT NULL -- Product name
  category TEXT NOT NULL -- Product category
  price REAL NOT NULL -- Price including tax (USD)
  stock INTEGER -- Current inventory count
)
"""

from google.genai import types
 
response = client.models.generate_content(
    model="gemini-2.5-flash",
    contents=prompt,
    config=types.GenerateContentConfig(
        system_instruction=SYSTEM_PROMPT,
        temperature=0.0,
        response_mime_type="application/json",
        response_schema={
            "type": "object",
            "properties": {
                "sql": {"type": "string"},
                "explanation": {"type": "string"},
            },
            "required": ["sql"],
        },
    ),
)

from functools import lru_cache
 
@lru_cache(maxsize=256)
def cached_generate_sql(question: str, schema_hash: str):
    schema = get_schema_by_hash(schema_hash)
    return generate_sql(question, schema)

def ask_with_retry(question, db_path, max_retries=2):
    schema = extract_schema(db_path)
    sql = generate_sql(question, schema)
 
    for attempt in range(max_retries + 1):
        try:
            return execute_query(db_path, sql)
        except Exception as e:
            if attempt == max_retries:
                raise
            fix_prompt = (
                f"The following SQL produced an error.\n"
                f"SQL: {sql}\n"
                f"Error: {e}\n"
                f"Generate a corrected SQL query."
            )
            sql = generate_sql(fix_prompt, schema)

GOLDEN_SET = [
    {
        "question": "Which products have at least 100 in stock?",
        "expected_sql": "SELECT name FROM products WHERE stock >= 100",
    },
    {
        "question": "What is the average price per category?",
        "expected_sql": (
            "SELECT category, AVG(price) FROM products GROUP BY category"
        ),
    },
    # In practice, prepare 40-50 questions
]
 
 
def normalize_rows(result: dict) -> set:
    # frozenset so comparison ignores column and row order
    return {frozenset(map(str, row)) for row in result["rows"]}
 
 
def evaluate(db_path: str) -> float:
    hits = 0
    for case in GOLDEN_SET:
        gen_sql = generate_sql(case["question"], extract_schema(db_path))
        try:
            got = normalize_rows(execute_query(db_path, gen_sql))
            want = normalize_rows(execute_query(db_path, case["expected_sql"]))
            if got == want:
                hits += 1
        except Exception:
            pass  # count execution failures as wrong
    return hits / len(GOLDEN_SET)

def extract_identifiers(sql: str) -> set:
    # rough identifier extraction, excluding string literals
    cleaned = re.sub(r"'[^']*'", "", sql)
    return set(re.findall(r"[A-Za-z_][A-Za-z0-9_]*", cleaned))
 
 
SQL_RESERVED = {
    "SELECT", "FROM", "WHERE", "GROUP", "BY", "ORDER",
    "LIMIT", "AS", "AND", "OR", "DESC", "ASC", "COUNT",
    "SUM", "AVG", "MAX", "MIN", "JOIN", "ON", "DISTINCT",
}
 
 
def validate_columns(sql: str, allowed: set) -> list:
    used = {t for t in extract_identifiers(sql) if t.upper() not in SQL_RESERVED}
    # allowed holds table and column names in lowercase
    unknown = {t for t in used if t.lower() not in allowed}
    return sorted(unknown)

RELATIONSHIP_HINTS = """
## Relationships
- order_items.order_id references orders.id
- order_items.product_id references products.id
- Per-product sales = SUM(order_items.quantity * order_items.unit_price)
- "sales" is computed from order_items; orders.total_amount is the per-order total
"""

import json, time, logging
 
logger = logging.getLogger("text2sql")
 
 
def ask_with_logging(question: str, db_path: str, allowed: set):
    start = time.perf_counter()
    schema = extract_schema(db_path)
    sql = generate_sql(question, schema)
    elapsed_ms = round((time.perf_counter() - start) * 1000)
 
    record = {
        "question": question,
        "generated_sql": sql,
        "unknown_columns": validate_columns(sql, allowed),
        "latency_ms": elapsed_ms,
    }
    logger.info(json.dumps(record, ensure_ascii=False))
    return record