-- Пример NL2SQL: вопрос на русском → SQL

-- Вопрос: "Покажи топ-5 клиентов по сумме заказов за последний квартал"
SELECT
    c.name AS "Клиент",
    SUM(o.total_amount) AS "Сумма заказов",
    COUNT(o.id) AS "Количество заказов"
FROM customers c
INNER JOIN orders o ON c.id = o.customer_id
WHERE o.created_at >= DATE_TRUNC('quarter', CURRENT_DATE - INTERVAL '3 months')
GROUP BY c.id, c.name
ORDER BY "Сумма заказов" DESC
LIMIT 5;
    

# nl2sql_prompt.py — построение системного промпта для NL2SQL
def build_system_prompt(schema_ddl: str, descriptions: dict) -> str:
    return f"""Ты — SQL-ассистент для аналитической базы данных.
Отвечай ТОЛЬКО SQL-запросом (без Markdown-форматирования).
СУБД: PostgreSQL 15.

== СХЕМА БАЗЫ ДАННЫХ ==
{schema_ddl}

== ОПИСАНИЯ ТАБЛИЦ И КОЛОНОК ==
{descriptions}

== ПРАВИЛА ==
1. Только SELECT-запросы. Запрещены INSERT/UPDATE/DELETE/DROP/ALTER.
2. Всегда добавляй LIMIT 100, если пользователь не указал иное.
3. Используй кириллические алиасы для колонок (AS "Название").
4. Для фильтрации по датам используй ISO-формат: '2026-06-15'.
5. Агрегации: SUM, AVG, COUNT, MAX, MIN — явно указывай GROUP BY.
6. Если вопрос неоднозначен — уточни, какие колонки использовать.

== ПРИМЕРЫ (FEW-SHOT) ==
Вопрос: Сколько заказов в этом месяце?
SQL: SELECT COUNT(*) AS "Количество заказов"
FROM orders WHERE created_at >= DATE_TRUNC('month', CURRENT_DATE);

Вопрос: Покажи клиентов из Москвы с заказами больше 100 000 руб
SQL: SELECT c.name AS "Клиент", c.city AS "Город",
       SUM(o.total_amount) AS "Сумма"
FROM customers c JOIN orders o ON c.id = o.customer_id
WHERE c.city = 'Москва'
GROUP BY c.id, c.name, c.city
HAVING SUM(o.total_amount) > 100000
ORDER BY "Сумма" DESC
LIMIT 100;

Вопрос: Средний чек по месяцам за 2025 год
SQL: SELECT TO_CHAR(created_at, 'YYYY-MM') AS "Месяц",
       ROUND(AVG(total_amount), 2) AS "Средний чек"
FROM orders WHERE created_at BETWEEN '2025-01-01' AND '2025-12-31'
GROUP BY "Месяц" ORDER BY "Месяц";
"""
    

# schema_extractor.py — извлечение схемы из PostgreSQL
import asyncpg
from typing import Dict

async def extract_schema(conn: asyncpg.Connection, schema: str = "public") -> Dict:
    # Получаем список таблиц
    tables = await conn.fetch("""
        SELECT table_name
        FROM information_schema.tables
        WHERE table_schema = $1 AND table_type = 'BASE TABLE'
        ORDER BY table_name
    """, schema)

    schema_info = {"tables": {}, "ddl": ""}

    for (table_name,) in tables:
        # Получаем колонки с типами и комментариями
        columns = await conn.fetch("""
            SELECT
                c.column_name,
                c.data_type,
                c.is_nullable,
                c.column_default,
                pgd.description AS comment
            FROM information_schema.columns c
            LEFT JOIN pg_catalog.pg_description pgd
                ON pgd.objoid = (quote_ident($1) || '.' || quote_ident(c.table_name))::regclass
                AND pgd.objsubid = c.ordinal_position
            WHERE c.table_schema = $1 AND c.table_name = $2
            ORDER BY c.ordinal_position
        """, schema, table_name)

        schema_info["tables"][table_name] = [
            {"name": col["column_name"], "type": col["data_type"],
             "nullable": col["is_nullable"] == "YES",
             "default": col["column_default"],
             "comment": col["comment"] or ""}
            for col in columns
        ]

        # Генерируем DDL
        ddl = f"CREATE TABLE {table_name} (\n"
        col_defs = []
        for col in columns:
            nullable = "" if col["is_nullable"] == "YES" else " NOT NULL"
            comment = f" -- {col['comment']}" if col["comment"] else ""
            col_defs.append(f"  {col['column_name']} {col['data_type']}{nullable}{comment}")
        ddl += ",\n".join(col_defs) + "\n);\n"
        schema_info["ddl"] += ddl + "\n"

    return schema_info
    

# sql_validator.py — валидация и безопасное выполнение SQL
import re
import asyncpg

FORBIDDEN_KEYWORDS = [
    r'\bDROP\b', r'\bINSERT\b', r'\bUPDATE\b', r'\bDELETE\b',
    r'\bTRUNCATE\b', r'\bALTER\b', r'\bCREATE\b', r'\bGRANT\b',
    r'\bREVOKE\b', r'\bCOPY\b', r'\bEXECUTE\b', r'\bVACUUM\b'
]

def validate_sql(sql: str) -> tuple[bool, str]:
    """Возвращает (валиден, сообщение_об_ошибке)."""
    sql_upper = sql.upper()

    # 1. Проверка запрещённых ключевых слов
    for pattern in FORBIDDEN_KEYWORDS:
        if re.search(pattern, sql_upper):
            return False, f"Запрос содержит запрещённую операцию: {pattern}"

    # 2. Проверка на SELECT
    if not sql_upper.strip().startswith("SELECT") and \
       not sql_upper.strip().startswith("WITH") and \
       not sql_upper.strip().startswith("EXPLAIN"):
        return False, "Разрешены только SELECT/WITH/EXPLAIN запросы"

    # 3. Проверка на наличие LIMIT
    if "LIMIT" not in sql_upper:
        return False, "Добавьте LIMIT для ограничения выборки"

    return True, "OK"

async def safe_execute_sql(conn: asyncpg.Connection, sql: str,
                          timeout_ms: int = 10000) -> dict:
    """Выполняет SQL с таймаутом и возвращает результат или ошибку."""
    is_valid, error = validate_sql(sql)
    if not is_valid:
        return {"status": "validation_error", "message": error}

    try:
        # Сначала EXPLAIN для проверки синтаксиса
        await conn.execute(f"EXPLAIN {sql}", timeout=timeout_ms / 1000)

        # Затем выполняем сам запрос
        rows = await conn.fetch(sql, timeout=timeout_ms / 1000)
        columns = [k for k in rows[0].keys()] if rows else []
        data = [[str(v) for v in row.values()] for row in rows]
        return {"status": "ok", "columns": columns, "rows": data, "count": len(rows)}
    except asyncpg.exceptions.PostgresSyntaxError as e:
        return {"status": "syntax_error", "message": str(e)}
    except Exception as e:
        return {"status": "execution_error", "message": str(e)}
    

# nl2sql_agent.py — основной цикл Text-to-SQL агента
import json
from openai import OpenAI

client = OpenAI()

async def nl2sql_agent(user_question: str, conn, schema_info: dict) -> dict:
    system_prompt = build_system_prompt(
        schema_info["ddl"],
        "\n".join(
            f"Таблица {t}: {', '.join(c['name'] + ' — ' + c['comment'] for c in cols if c['comment'])}"
            for t, cols in schema_info["tables"].items()
        )
    )

    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_question}
    ]

    for attempt in range(3):  # до 3 попыток исправления
        response = client.chat.completions.create(
            model="gpt-4", messages=messages, temperature=0.1
        )
        sql = response.choices[0].message.content.strip()

        # Убираем Markdown-форматирование, если модель его добавила
        if sql.startswith("```"):
            sql = sql.split("\n")[1:-1]
            sql = "\n".join(sql)

        result = await safe_execute_sql(conn, sql)

        if result["status"] == "ok":
            # Генерируем пояснение на естественном языке
            explanation = client.chat.completions.create(
                model="gpt-4", temperature=0.3,
                messages=[
                    {"role": "system", "content": "Кратко поясни результат SQL-запроса на русском, 1-2 предложения."},
                    {"role": "user", "content": f"Запрос: {sql}\nРезультат: {json.dumps(result['rows'], ensure_ascii=False)}"}
                ]
            )
            return {
                "sql": sql,
                "columns": result["columns"],
                "rows": result["rows"],
                "explanation": explanation.choices[0].message.content
            }

        # Self-correction: добавляем ошибку в контекст и просим исправить
        messages.append({"role": "assistant", "content": sql})
        messages.append({"role": "user", "content":
            f"Ошибка выполнения: {result['message']}. Исправь запрос."})

    return {"error": "Не удалось сгенерировать корректный SQL за 3 попытки"}
    

# advanced_nl2sql.py — граф связей и семантическая валидация
import asyncpg

async def build_join_graph(conn: asyncpg.Connection) -> str:
    """Извлекает foreign keys и строит граф связей для промпта."""
    fks = await conn.fetch("""
        SELECT
            tc.table_name AS from_table,
            kcu.column_name AS from_column,
            ccu.table_name AS to_table,
            ccu.column_name AS to_column
        FROM information_schema.table_constraints tc
        JOIN information_schema.key_column_usage kcu
            ON tc.constraint_name = kcu.constraint_name
        JOIN information_schema.constraint_column_usage ccu
            ON ccu.constraint_name = tc.constraint_name
        WHERE tc.constraint_type = 'FOREIGN KEY'
          AND tc.table_schema = 'public'
    """)

    graph_lines = ["== СВЯЗИ МЕЖДУ ТАБЛИЦАМИ (FOREIGN KEYS) =="]
    for fk in fks:
        graph_lines.append(
            f"{fk['from_table']}.{fk['from_column']} → {fk['to_table']}.{fk['to_column']}"
        )
    return "\n".join(graph_lines)

# Кэш запросов: LRU-словарь вопрос → SQL
from functools import lru_cache
import hashlib

query_cache = {}

def cache_key(question: str) -> str:
    return hashlib.md5(question.lower().strip().encode()).hexdigest()

def get_cached(question: str) -> str | None:
    return query_cache.get(cache_key(question))

def set_cache(question: str, sql: str):
    key = cache_key(question)
    if len(query_cache) > 1000:
        query_cache.pop(next(iter(query_cache)))  # простой LRU
    query_cache[key] = sql