sql_engine.py

"""
Universal Text-to-SQL Engine with Neural AI & Fallback Rule Resolver.
Supports:
1. Google Gemini API (gemini-1.5-flash / gemini-2.0-flash)
2. Groq API (llama-3.3-70b-versatile)
3. OpenAI / DeepSeek / OpenRouter / Ollama API
4. Advanced Schema-Aware Local Fallback Engine
5. Automatic Self-Correction & Syntax Validation
"""
import re
import os
import json
import sqlite3
import time
import requests
from typing import Dict, List, Any, Optional, Tuple

try:
    from dotenv import load_dotenv
    load_dotenv()
except ImportError:
    pass


class UniversalSQLEngine:
    def __init__(self, db_dir: Optional[str] = None):
        self.db_dir = db_dir or os.path.join(os.path.dirname(os.path.abspath(__file__)), 'databases')
        os.makedirs(self.db_dir, exist_ok=True)
        self.active_db_name: Optional[str] = None
        self.active_db_path: Optional[str] = None

    def set_active_database(self, db_name: Optional[str]) -> bool:
        if not db_name or db_name == 'none':
            self.active_db_name = None
            self.active_db_path = None
            return True
        path = os.path.join(self.db_dir, db_name)
        if os.path.exists(path):
            self.active_db_name = db_name
            self.active_db_path = path
            return True
        return False

    def list_available_databases(self) -> List[Dict[str, Any]]:
        databases = []
        if not os.path.exists(self.db_dir):
            return databases
        for f in os.listdir(self.db_dir):
            if f.endswith('.db') or f.endswith('.sqlite'):
                full_path = os.path.join(self.db_dir, f)
                size_kb = round(os.path.getsize(full_path) / 1024, 2)
                databases.append({
                    'name': f,
                    'size_kb': size_kb,
                    'is_active': (f == self.active_db_name)
                })
        return databases

    def get_database_schema(self, db_name: Optional[str] = None) -> Optional[Dict[str, Any]]:
        target_path = os.path.join(self.db_dir, db_name) if db_name else self.active_db_path
        if not target_path or not os.path.exists(target_path):
            return None
        try:
            conn = sqlite3.connect(target_path)
            cur = conn.cursor()
            cur.execute("SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%';")
            tables = [row[0] for row in cur.fetchall()]
            schema_data = {}
            for table in tables:
                cur.execute(f"PRAGMA table_info({table});")
                columns_raw = cur.fetchall()
                cur.execute(f"PRAGMA foreign_key_list({table});")
                fks_raw = cur.fetchall()
                cur.execute(f"SELECT COUNT(*) FROM {table};")
                row_count = cur.fetchone()[0]
                cur.execute(f"SELECT * FROM {table} LIMIT 3;")
                sample_rows = cur.fetchall()
                col_names = [col[1] for col in columns_raw]
                schema_data[table] = {
                    'row_count': row_count,
                    'columns': [
                        {'name': col[1], 'type': col[2], 'is_pk': bool(col[5])}
                        for col in columns_raw
                    ],
                    'foreign_keys': [
                        {'from': fk[3], 'to_table': fk[2], 'to_column': fk[4]}
                        for fk in fks_raw
                    ],
                    'sample_data': {
                        'columns': col_names,
                        'rows': sample_rows
                    }
                }
            conn.close()
            return {'database': os.path.basename(target_path), 'tables': schema_data}
        except Exception:
            return None

    def _format_schema_for_prompt(self, active_schema: Optional[Dict[str, Any]], custom_schema: Optional[str]) -> str:
        parts = []
        if active_schema and 'tables' in active_schema:
            parts.append(f"Database: {active_schema.get('database', 'database.db')}")
            for table, data in active_schema['tables'].items():
                cols_str = ", ".join([f"{c['name']} ({c['type']}{' PK' if c['is_pk'] else ''})" for c in data['columns']])
                parts.append(f"\nTable '{table}' ({data['row_count']} records):")
                parts.append(f"  Columns: {cols_str}")
                if data['foreign_keys']:
                    fk_strs = [f"{fk['from']} -> {fk['to_table']}({fk['to_column']})" for fk in data['foreign_keys']]
                    parts.append(f"  Foreign Keys: {', '.join(fk_strs)}")
                if data.get('sample_data') and data['sample_data'].get('rows'):
                    samples = data['sample_data']['rows'][:2]
                    parts.append(f"  Sample data: {samples}")
        if custom_schema:
            parts.append(f"\nUser Provided Schema Hint:\n{custom_schema}")
        return "\n".join(parts) if parts else "No schema provided. Generate standard SQL assuming typical entity and column naming."

    # ── LLM Providers ────────────────────────────────────────────────────────
    def _call_gemini(self, prompt: str, schema_str: str, dialect: str, api_key: str) -> Optional[Tuple[str, str, Dict[str, str]]]:
        """Calls Google Gemini API for high-precision SQL generation."""
        models_to_try = ["gemini-2.5-flash", "gemini-flash-latest", "gemini-pro-latest", "gemini-1.5-flash-latest"]
        
        system_instruction = (
            f"You are a principal SQL Architect and database expert. "
            f"Your task is to convert the user's natural language question into an optimal, syntactically flawless SQL query for {dialect.upper()}.\n"
            f"Follow these strict rules:\n"
            f"1. Use ONLY the tables, columns, and foreign keys provided in the Schema.\n"
            f"2. Use appropriate JOINs, WHERE filters, GROUP BY, aggregations, subqueries, or window functions as needed.\n"
            f"3. Return a valid JSON object ONLY, with no extra markdown formatting or backticks around the JSON, with the following exact keys:\n"
            f'{{"sql_query": "YOUR_SQL_QUERY_HERE", "explanation": "Brief 1-2 sentence explanation", "clauses_breakdown": {{"SELECT": "...", "FROM": "...", "WHERE": "...", "GROUP BY": "...", "ORDER BY": "...", "LIMIT": "..."}}}}\n'
        )
        user_content = f"Schema Context:\n{schema_str}\n\nUser Question:\n{prompt}\nTarget Dialect: {dialect.upper()}"

        for model in models_to_try:
            url = f"https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent?key={api_key}"
            headers = {"Content-Type": "application/json"}
            payload = {
                "contents": [
                    {"role": "user", "parts": [{"text": system_instruction + "\n\n" + user_content}]}
                ],
                "generationConfig": {
                    "temperature": 0.1,
                    "response_mime_type": "application/json"
                }
            }
            try:
                resp = requests.post(url, headers=headers, json=payload, timeout=15)
                if resp.status_code == 200:
                    data = resp.json()
                    text = data['candidates'][0]['content']['parts'][0]['text']
                    text = re.sub(r'^```(?:json)?\s*', '', text.strip(), flags=re.IGNORECASE)
                    text = re.sub(r'```$', '', text.strip())
                    parsed = json.loads(text)
                    sql = parsed.get('sql_query') or parsed.get('sql') or parsed.get('query')
                    exp = parsed.get('explanation', '')
                    clauses = parsed.get('clauses_breakdown', {})
                    if sql:
                        return sql, exp, clauses
            except Exception as e:
                print(f"[GEMINI {model} ERROR] {e}")
        return None

    def _call_groq(self, prompt: str, schema_str: str, dialect: str, api_key: str) -> Optional[Tuple[str, str, Dict[str, str]]]:
        """Calls Groq API for instant, free SQL generation."""
        url = "https://api.groq.com/openai/v1/chat/completions"
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {api_key}"
        }
        models_to_try = ["qwen/qwen3.6-27b", "openai/gpt-oss-120b", "groq/compound"]
        
        messages = [
            {
                "role": "system",
                "content": (
                    f"You are a database and SQL expert for {dialect.upper()}.\n"
                    f"Given a database schema and a natural language question, output ONLY a valid JSON object matching:\n"
                    f'{{"sql_query": "...", "explanation": "...", "clauses_breakdown": {{"SELECT": "...", "FROM": "..."}}}}\n'
                    f"Do not include any conversational filler."
                )
            },
            {
                "role": "user",
                "content": f"Schema:\n{schema_str}\n\nQuestion: {prompt}\nTarget Dialect: {dialect}"
            }
        ]

        for model in models_to_try:
            payload = {
                "model": model,
                "messages": messages,
                "temperature": 0.1,
                "response_format": {"type": "json_object"}
            }
            try:
                resp = requests.post(url, headers=headers, json=payload, timeout=15)
                if resp.status_code == 200:
                    data = resp.json()
                    content = data['choices'][0]['message']['content']
                    content = re.sub(r'^```(?:json)?\s*', '', content.strip(), flags=re.IGNORECASE)
                    content = re.sub(r'```$', '', content.strip())
                    parsed = json.loads(content)
                    sql = parsed.get('sql_query') or parsed.get('sql') or parsed.get('query')
                    exp = parsed.get('explanation', '')
                    clauses = parsed.get('clauses_breakdown', {})
                    if sql:
                        return sql, exp, clauses
            except Exception as e:
                print(f"[GROQ {model} ERROR] {e}")
        return None

    def _call_openai_compatible(self, prompt: str, schema_str: str, dialect: str, api_key: str, base_url: str = "https://api.openai.com/v1", model: str = "gpt-4o-mini") -> Optional[Tuple[str, str, Dict[str, str]]]:
        """Calls OpenAI or any OpenAI-compatible endpoint."""
        url = f"{base_url.rstrip('/')}/chat/completions"
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {api_key}"
        }
        messages = [
            {
                "role": "system",
                "content": (
                    f"You are an expert SQL generator for {dialect.upper()}.\n"
                    f"Output ONLY a valid JSON object with keys: sql_query, explanation, clauses_breakdown."
                )
            },
            {
                "role": "user",
                "content": f"Schema:\n{schema_str}\n\nQuestion: {prompt}\nDialect: {dialect}"
            }
        ]
        payload = {
            "model": model,
            "messages": messages,
            "temperature": 0.1,
            "response_format": {"type": "json_object"}
        }
        try:
            resp = requests.post(url, headers=headers, json=payload, timeout=15)
            if resp.status_code == 200:
                data = resp.json()
                content = data['choices'][0]['message']['content']
                parsed = json.loads(content)
                return parsed.get('sql_query'), parsed.get('explanation', ''), parsed.get('clauses_breakdown', {})
        except Exception as e:
            print(f"[OPENAI COMPATIBLE API ERROR] {e}")
        return None

    # ── Rule-Based Schema-Aware Engine (Zero-dependency Fallback) ─────────────
    def _resolve_column(self, token: str, table_columns: List[str]) -> Optional[str]:
        token = token.lower().strip()
        if token in table_columns:
            return token
        for col in table_columns:
            if token == col.lower() or token.rstrip('s') == col.lower():
                return col
        synonyms = {
            'stock': ['stock_quantity', 'quantity'],
            'quantity': ['quantity', 'stock_quantity'],
            'price': ['price', 'unit_price', 'cost'],
            'amount': ['total_amount', 'amount', 'price', 'spending'],
            'total': ['total_amount', 'amount', 'budget', 'salary', 'price'],
            'revenue': ['total_amount', 'amount', 'price'],
            'salary': ['salary', 'wage', 'pay'],
            'paid': ['salary'],
            'credits': ['credits', 'score'],
            'name': ['name', 'first_name', 'dept_name', 'title'],
            'full_name': ['first_name', 'name'],
            'first_name': ['first_name', 'name'],
            'last_name': ['last_name', 'name'],
            'department': ['name', 'dept_name', 'department_id', 'dept_id'],
            'dept': ['dept_name', 'name', 'dept_id', 'department_id'],
            'course': ['title', 'name', 'course_id'],
            'title': ['title', 'name', 'job_title'],
            'role': ['role', 'job_title', 'title'],
            'job': ['job_title', 'role'],
            'city': ['city', 'location'],
            'location': ['location', 'city', 'building'],
            'country': ['country', 'location'],
            'date': ['order_date', 'joined_date', 'hire_date', 'start_date', 'date'],
            'status': ['status'],
            'rating': ['rating', 'gpa', 'grade'],
            'gpa': ['gpa', 'rating'],
            'grade': ['grade', 'gpa']
        }
        candidates = synonyms.get(token, [token])
        for c in candidates:
            if c in table_columns:
                return c
        for c in candidates:
            for col in table_columns:
                if c in col.lower():
                    return col
        return None

    def _find_table_for_col(self, col_name: str, primary_table: str, secondary_tables: List[str], tables_meta: Dict[str, Any]) -> str:
        if primary_table in tables_meta:
            p_cols = [c['name'].lower() for c in tables_meta[primary_table]['columns']]
            if col_name.lower() in p_cols:
                return primary_table
        for sec in secondary_tables:
            if sec in tables_meta:
                s_cols = [c['name'].lower() for c in tables_meta[sec]['columns']]
                if col_name.lower() in s_cols:
                    return sec
        return primary_table

    def _qualify_col(self, col_name: str, table_name: str, table_aliases: Dict[str, str], is_multi_table: bool, secondary_tables: Optional[List[str]] = None, tables_meta: Optional[Dict[str, Any]] = None) -> str:
        if not is_multi_table or '.' in col_name or '(' in col_name or ' ' in col_name:
            return col_name
        actual_table = table_name
        if secondary_tables and tables_meta:
            actual_table = self._find_table_for_col(col_name, table_name, secondary_tables, tables_meta)
        alias = table_aliases.get(actual_table, actual_table)
        return f"{alias}.{col_name}"

    def _build_sql_rule_based(
        self,
        prompt: str,
        active_schema: Optional[Dict[str, Any]],
        custom_schema: Optional[str],
        dialect: str
    ) -> Tuple[str, str, Dict[str, str], float]:
        p = prompt.strip()
        p_lower = p.lower()

        tables_meta = active_schema['tables'] if (active_schema and 'tables' in active_schema) else {}
        available_tables = list(tables_meta.keys())

        # 1. Identify Primary & Secondary Tables
        primary_table = None
        secondary_tables = []

        table_keywords = {
            'customers': ['customer', 'client', 'buyer', 'user', 'people'],
            'orders': ['order', 'purchase', 'transaction', 'sale', 'revenue'],
            'products': ['product', 'item', 'good', 'catalog', 'stock', 'priced'],
            'categories': ['category', 'categories'],
            'order_items': ['order item', 'order_item', 'item detail'],
            'employees': ['employee', 'staff', 'worker', 'developer', 'engineer', 'manager', 'hire'],
            'departments': ['department', 'dept', 'branch', 'division', 'building'],
            'projects': ['project', 'assignment', 'task'],
            'employee_projects': ['employee project', 'assignment'],
            'students': ['student', 'pupil', 'learner', 'gpa'],
            'instructors': ['instructor', 'teacher', 'professor', 'faculty'],
            'courses': ['course', 'subject', 'class', 'credits'],
            'enrollments': ['enrollment', 'grade', 'registration']
        }

        table_scores = {}
        for tbl in available_tables:
            score = 0
            if tbl.lower() in p_lower or tbl.lower().rstrip('s') in p_lower:
                score += 15
            kw_list = table_keywords.get(tbl.lower(), [])
            for kw in kw_list:
                if re.search(r'\b' + re.escape(kw) + r'\b', p_lower):
                    score += 8
            table_scores[tbl] = score

        if re.search(r'\b(employee|employees|staff|worker|workers)\b', p_lower) and 'employees' in table_scores:
            table_scores['employees'] += 30
        if re.search(r'\b(salary|salaries|paid|wage)\b', p_lower) and 'employees' in table_scores:
            table_scores['employees'] += 25
        if re.search(r'\b(product|products|item|items)\b', p_lower) and 'products' in table_scores:
            table_scores['products'] += 20
        if re.search(r'\b(order|orders)\b', p_lower) and 'orders' in table_scores:
            table_scores['orders'] += 20
        if re.search(r'\b(customer|customers)\b', p_lower) and 'customers' in table_scores:
            table_scores['customers'] += 20
        if re.search(r'\b(course|courses)\b', p_lower) and 'courses' in table_scores:
            table_scores['courses'] += 20
        if re.search(r'\b(student|students)\b', p_lower) and 'students' in table_scores:
            table_scores['students'] += 20

        sorted_tables = sorted(table_scores.items(), key=lambda x: x[1], reverse=True)
        if sorted_tables and sorted_tables[0][1] > 0:
            primary_table = sorted_tables[0][0]
            for tbl, sc in sorted_tables[1:]:
                if sc > 0:
                    secondary_tables.append(tbl)

        if not primary_table and available_tables:
            for tbl, tdata in tables_meta.items():
                col_names = [c['name'].lower() for c in tdata['columns']]
                for word in p_lower.split():
                    if word in col_names:
                        primary_table = tbl
                        break
                if primary_table:
                    break
            if not primary_table:
                primary_table = available_tables[0]
        elif not primary_table:
            primary_table = "table_name"

        primary_cols = [c['name'].lower() for c in tables_meta[primary_table]['columns']] if primary_table in tables_meta else []

        # 2. Check for multi-table JOINs
        join_clauses = []
        from_clause = primary_table
        table_aliases = {primary_table: primary_table[0].lower()}

        if secondary_tables:
            for sec_tbl in secondary_tables:
                if sec_tbl == primary_table:
                    continue
                sec_alias = sec_tbl[0].lower()
                if sec_alias in table_aliases.values():
                    sec_alias = sec_tbl[:2].lower()
                table_aliases[sec_tbl] = sec_alias

                join_condition = None
                p_fks = tables_meta.get(primary_table, {}).get('foreign_keys', [])
                for fk in p_fks:
                    if fk['to_table'].lower() == sec_tbl.lower():
                        join_condition = f"{table_aliases[primary_table]}.{fk['from']} = {sec_alias}.{fk['to_column']}"
                        break
                if not join_condition:
                    s_fks = tables_meta.get(sec_tbl, {}).get('foreign_keys', [])
                    for fk in s_fks:
                        if fk['to_table'].lower() == primary_table.lower():
                            join_condition = f"{sec_alias}.{fk['from']} = {table_aliases[primary_table]}.{fk['to_column']}"
                            break
                if not join_condition:
                    sec_cols = [c['name'].lower() for c in tables_meta.get(sec_tbl, {}).get('columns', [])]
                    common = set(primary_cols).intersection(set(sec_cols))
                    for col in common:
                        if col.endswith('_id') or col == 'id':
                            join_condition = f"{table_aliases[primary_table]}.{col} = {sec_alias}.{col}"
                            break

                if join_condition:
                    join_clauses.append(f"JOIN {sec_tbl} {sec_alias} ON {join_condition}")

        is_multi_table = bool(join_clauses)
        if is_multi_table:
            from_clause = f"{primary_table} {table_aliases[primary_table]} " + " ".join(join_clauses)

        p_alias = table_aliases[primary_table]

        # 3. Projections & Aggregations
        select_clause = f"{p_alias}.*" if is_multi_table else "*"
        aggregation = None
        agg_col = None

        if re.search(r'\b(count|how many|number of)\b', p_lower):
            aggregation = "COUNT"
            select_clause = "COUNT(*) AS total_count"
        elif re.search(r'\b(average|avg|mean)\b', p_lower):
            aggregation = "AVG"
            for token in ['salary', 'price', 'total_amount', 'amount', 'rating', 'gpa', 'budget', 'credits', 'score', 'cost', 'spending']:
                if token in p_lower:
                    for t in [primary_table] + secondary_tables:
                        cols = [c['name'].lower() for c in tables_meta.get(t, {}).get('columns', [])]
                        resolved = self._resolve_column(token, cols)
                        if resolved:
                            agg_col = resolved
                            break
                    if agg_col:
                        break
            if not agg_col and primary_cols:
                for c in tables_meta.get(primary_table, {}).get('columns', []):
                    if c['type'] in ['REAL', 'INTEGER', 'NUMERIC', 'FLOAT', 'DOUBLE'] and not c['is_pk'] and not c['name'].endswith('_id'):
                        agg_col = c['name']
                        break
            agg_col = agg_col or "price"
            qualified_agg = self._qualify_col(agg_col, primary_table, table_aliases, is_multi_table, secondary_tables, tables_meta)
            select_clause = f"ROUND(AVG({qualified_agg}), 2) AS average_{agg_col}"
        elif re.search(r'\b(sum|summing|total revenue|total sales|total spending|sum of|calculate total|total budget)\b', p_lower) and not re.search(r'\btotal (count|number)\b', p_lower):
            aggregation = "SUM"
            for token in ['total_amount', 'amount', 'salary', 'budget', 'price', 'revenue', 'spending']:
                for t in [primary_table] + secondary_tables:
                    cols = [c['name'].lower() for c in tables_meta.get(t, {}).get('columns', [])]
                    resolved = self._resolve_column(token, cols)
                    if resolved:
                        agg_col = resolved
                        break
                if agg_col:
                    break
            if not agg_col:
                for c in tables_meta.get(primary_table, {}).get('columns', []):
                    if c['type'] in ['REAL', 'INTEGER', 'NUMERIC'] and not c['is_pk'] and not c['name'].endswith('_id'):
                        agg_col = c['name']
                        break
            agg_col = agg_col or "total_amount"
            qualified_agg = self._qualify_col(agg_col, primary_table, table_aliases, is_multi_table, secondary_tables, tables_meta)
            select_clause = f"SUM({qualified_agg}) AS total_{agg_col}"
        elif re.search(r'\b(highest|max|maximum|most expensive|top paid|greatest)\b', p_lower) and not re.search(r'\btop\s+\d+\b', p_lower):
            aggregation = "MAX"
            for token in ['salary', 'price', 'total_amount', 'budget', 'gpa', 'rating', 'credits']:
                resolved = self._resolve_column(token, primary_cols)
                if resolved:
                    agg_col = resolved
                    break
            agg_col = agg_col or "price"
            qualified_agg = self._qualify_col(agg_col, primary_table, table_aliases, is_multi_table, secondary_tables, tables_meta)
            select_clause = f"MAX({qualified_agg}) AS max_{agg_col}"
        elif re.search(r'\b(lowest|min|minimum|cheapest|least)\b', p_lower) and not re.search(r'\btop\s+\d+\b', p_lower):
            aggregation = "MIN"
            for token in ['salary', 'price', 'total_amount', 'budget', 'gpa', 'rating', 'credits']:
                resolved = self._resolve_column(token, primary_cols)
                if resolved:
                    agg_col = resolved
                    break
            agg_col = agg_col or "price"
            qualified_agg = self._qualify_col(agg_col, primary_table, table_aliases, is_multi_table, secondary_tables, tables_meta)
            select_clause = f"MIN({qualified_agg}) AS min_{agg_col}"

        # 4. Multi-column projections / Full names
        if not aggregation:
            if 'customer full name' in p_lower or 'customer name' in p_lower or ('full name' in p_lower and primary_table == 'customers'):
                if is_multi_table and 'customers' in table_aliases and 'orders' in table_aliases:
                    c_al = table_aliases['customers']
                    o_al = table_aliases['orders']
                    select_clause = f"{o_al}.order_id, {c_al}.first_name || ' ' || {c_al}.last_name AS customer_name, {o_al}.total_amount"
                elif 'first_name' in primary_cols and 'last_name' in primary_cols:
                    select_clause = "first_name || ' ' || last_name AS full_name, *"
            elif 'name and total amount' in p_lower and is_multi_table:
                sec_alias = table_aliases[secondary_tables[0]]
                select_clause = f"{p_alias}.order_id, {sec_alias}.first_name || ' ' || {sec_alias}.last_name AS customer_name, {p_alias}.total_amount"
            else:
                matched_cols = []
                for token in ['first_name', 'last_name', 'name', 'email', 'salary', 'job_title', 'price', 'status', 'total_amount', 'stock_quantity', 'rating', 'gpa', 'credits', 'city', 'country', 'budget', 'location']:
                    if re.search(r'\b' + token.replace('_', ' ') + r'\b', p_lower) or re.search(r'\b' + token + r'\b', p_lower):
                        resolved = self._resolve_column(token, primary_cols)
                        if resolved and resolved not in matched_cols:
                            matched_cols.append(self._qualify_col(resolved, primary_table, table_aliases, is_multi_table, secondary_tables, tables_meta))
                if matched_cols and len(matched_cols) > 1 and not re.search(r'\b(show all|list all|select \*)\b', p_lower):
                    select_clause = ", ".join(matched_cols)
                else:
                    select_clause = f"{p_alias}.*" if is_multi_table else "*"

        # 5. Group By Clause
        group_by_clause = None
        group_match = re.search(r'\b(?:group\s+by|per|for\s+each|by)\s+([a-zA-Z_]+)\b', p_lower)
        if group_match:
            candidate_group = group_match.group(1).lower()
            if candidate_group not in ['the', 'a', 'an', 'order', 'desc', 'asc', 'salary', 'price', 'amount', 'highest', 'lowest', 'top', 'summing', 'country', 'usa']:
                resolved_group = None
                if is_multi_table:
                    for sec in secondary_tables:
                        if sec in ['categories', 'departments'] and candidate_group in ['category', 'department', 'dept', 'categories']:
                            sec_al = table_aliases[sec]
                            resolved_group = f"{sec_al}.name AS {candidate_group}"
                            break
                if not resolved_group:
                    res_col = self._resolve_column(candidate_group, primary_cols)
                    if res_col:
                        resolved_group = self._qualify_col(res_col, primary_table, table_aliases, is_multi_table, secondary_tables, tables_meta)
                if not resolved_group and primary_table == 'products' and 'category' in candidate_group:
                    resolved_group = f"{p_alias}.category_id" if is_multi_table else "category_id"
                elif not resolved_group and primary_table == 'employees' and 'department' in candidate_group:
                    resolved_group = f"{p_alias}.department_id" if is_multi_table else "department_id"

                if resolved_group:
                    group_expr = resolved_group.split(' AS ')[0]
                    group_by_clause = group_expr
                    if select_clause != "*" and select_clause != f"{p_alias}.*" and resolved_group not in select_clause:
                        select_clause = f"{resolved_group}, {select_clause}"
                    elif select_clause in ["*", f"{p_alias}.*"]:
                        select_clause = f"{resolved_group}, COUNT(*) AS total_count"

        # 6. WHERE Conditions
        where_conditions = []

        num_comp = re.search(r'\b([a-zA-Z_]+)\s*(?:that have|with|is|are)?\s*(>|<|>=|<=|=|greater than|more than|less than|higher than|above|below|over|under)\s*\$?(\d+(?:\.\d+)?)\b', p_lower)
        if num_comp:
            col_token = num_comp.group(1).lower()
            op_raw = num_comp.group(2).lower()
            val = num_comp.group(3)
            
            resolved_col = self._resolve_column(col_token, primary_cols)
            if not resolved_col and 'stock' in col_token:
                resolved_col = self._resolve_column('stock_quantity', primary_cols)
            if not resolved_col and primary_cols:
                for col in primary_cols:
                    if col in ['price', 'salary', 'stock_quantity', 'budget', 'rating', 'total_amount', 'gpa', 'credits']:
                        resolved_col = col
                        break
            if resolved_col:
                op = ">" if any(w in op_raw for w in ['>', 'greater', 'more', 'higher', 'above', 'over']) else "<" if any(w in op_raw for w in ['<', 'less', 'below', 'under']) else "="
                q_col = self._qualify_col(resolved_col, primary_table, table_aliases, is_multi_table, secondary_tables, tables_meta)
                where_conditions.append(f"{q_col} {op} {val}")

        country_match = re.search(r'\b(?:from|in|country)\s+[\'"]?([a-zA-Z]{2,20})[\'"]?\b', p)
        if country_match:
            c_val = country_match.group(1).strip()
            if c_val.lower() not in ['engineering', 'marketing', 'sales', 'hr', 'the', 'each', 'all', 'any', 'orders', 'products', 'customers', 'employees', 'by']:
                if 'country' in primary_cols:
                    q_col = self._qualify_col('country', primary_table, table_aliases, is_multi_table, secondary_tables, tables_meta)
                    where_conditions.append(f"{q_col} = '{c_val}'")
                elif 'city' in primary_cols:
                    q_col = self._qualify_col('city', primary_table, table_aliases, is_multi_table, secondary_tables, tables_meta)
                    where_conditions.append(f"{q_col} = '{c_val}'")

        dept_match = re.search(r'\b(?:in\s+)?([a-zA-Z]+)\s+department\b', p_lower)
        if dept_match:
            dept_name = dept_match.group(1).capitalize()
            if dept_name not in ['The', 'Each', 'All', 'Average', 'Any', 'By', 'Calculate']:
                if is_multi_table and 'departments' in table_aliases:
                    d_al = table_aliases['departments']
                    where_conditions.append(f"{d_al}.name LIKE '{dept_name}%'")
                elif 'department_id' in primary_cols and 'departments' in tables_meta:
                    where_conditions.append(f"{self._qualify_col('department_id', primary_table, table_aliases, is_multi_table, secondary_tables, tables_meta)} IN (SELECT department_id FROM departments WHERE name LIKE '{dept_name}%')")
                elif 'department' in primary_cols:
                    where_conditions.append(f"{self._qualify_col('department', primary_table, table_aliases, is_multi_table, secondary_tables, tables_meta)} = '{dept_name}'")

        if 'status' in primary_cols:
            status_match = re.search(r'\bstatus\s*(?:is|=|:)?\s*[\'"]?([a-zA-Z]+)[\'"]?\b', p_lower)
            if status_match:
                q_status = self._qualify_col('status', primary_table, table_aliases, is_multi_table, secondary_tables, tables_meta)
                where_conditions.append(f"{q_status} = '{status_match.group(1).capitalize()}'")
            elif 'active' in p_lower:
                q_status = self._qualify_col('status', primary_table, table_aliases, is_multi_table, secondary_tables, tables_meta)
                where_conditions.append(f"{q_status} = 'Active'")
            elif 'delivered' in p_lower:
                q_status = self._qualify_col('status', primary_table, table_aliases, is_multi_table, secondary_tables, tables_meta)
                where_conditions.append(f"{q_status} = 'Delivered'")
            elif 'cancelled' in p_lower or 'canceled' in p_lower:
                q_status = self._qualify_col('status', primary_table, table_aliases, is_multi_table, secondary_tables, tables_meta)
                where_conditions.append(f"{q_status} = 'Cancelled'")

        # 7. Order By & Limit
        order_by_clause = None
        limit_clause = None

        top_match = re.search(r'\btop\s+(\d+)\b', p_lower)
        if top_match:
            limit_clause = top_match.group(1)
            if re.search(r'\b(highest|paid|salary|expensive|priced|best|max|greatest|credits)\b', p_lower):
                sort_col = None
                for token in ['salary', 'price', 'credits', 'total_amount', 'rating', 'gpa', 'budget']:
                    if token in p_lower or token in primary_cols:
                        resolved = self._resolve_column(token, primary_cols)
                        if resolved:
                            sort_col = resolved
                            break
                if not sort_col and primary_cols:
                    sort_col = primary_cols[0]
                if sort_col:
                    q_sort = self._qualify_col(sort_col, primary_table, table_aliases, is_multi_table, secondary_tables, tables_meta)
                    order_by_clause = f"{q_sort} DESC"
            elif re.search(r'\b(lowest|cheapest|least|smallest|bottom)\b', p_lower):
                sort_col = None
                for token in ['price', 'salary', 'total_amount', 'stock_quantity']:
                    resolved = self._resolve_column(token, primary_cols)
                    if resolved:
                        sort_col = resolved
                        break
                if sort_col:
                    q_sort = self._qualify_col(sort_col, primary_table, table_aliases, is_multi_table, secondary_tables, tables_meta)
                    order_by_clause = f"{q_sort} ASC"

        explicit_order = re.search(r'\border\s+by\s+([a-zA-Z_]+)(?:\s+(asc|desc))?\b', p_lower)
        if explicit_order:
            col_o = self._resolve_column(explicit_order.group(1), primary_cols) or explicit_order.group(1)
            dir_o = explicit_order.group(2).upper() if explicit_order.group(2) else "ASC"
            q_col_o = self._qualify_col(col_o, primary_table, table_aliases, is_multi_table, secondary_tables, tables_meta)
            order_by_clause = f"{q_col_o} {dir_o}"

        explicit_limit = re.search(r'\blimit\s+(\d+)\b', p_lower)
        if explicit_limit:
            limit_clause = explicit_limit.group(1)

        # 8. Assemble SQL Query
        clauses_dict = {}
        sql_parts = []

        if dialect == 'tsql' and limit_clause:
            select_stmt = f"SELECT TOP {limit_clause} {select_clause}"
            clauses_dict['SELECT'] = f"TOP {limit_clause} {select_clause}"
        else:
            select_stmt = f"SELECT {select_clause}"
            clauses_dict['SELECT'] = select_clause
        sql_parts.append(select_stmt)

        sql_parts.append(f"FROM {from_clause}")
        clauses_dict['FROM'] = from_clause

        if where_conditions:
            where_str = " AND ".join(where_conditions)
            sql_parts.append(f"WHERE {where_str}")
            clauses_dict['WHERE'] = where_str

        if group_by_clause:
            sql_parts.append(f"GROUP BY {group_by_clause}")
            clauses_dict['GROUP BY'] = group_by_clause

        if order_by_clause:
            sql_parts.append(f"ORDER BY {order_by_clause}")
            clauses_dict['ORDER BY'] = order_by_clause

        if limit_clause and dialect not in ['tsql', 'oracle']:
            sql_parts.append(f"LIMIT {limit_clause}")
            clauses_dict['LIMIT'] = limit_clause
        elif dialect == 'oracle' and limit_clause:
            sql_parts.append(f"FETCH FIRST {limit_clause} ROWS ONLY")
            clauses_dict['FETCH'] = f"FIRST {limit_clause} ROWS ONLY"

        final_sql = "\n".join(sql_parts) + ";"

        explanation_parts = [f"Retrieves `{select_clause}` from `{from_clause}`."]
        if where_conditions:
            explanation_parts.append(f"Filters records where {' and '.join(where_conditions)}.")
        if group_by_clause:
            explanation_parts.append(f"Aggregates data grouped by `{group_by_clause}`.")
        if order_by_clause:
            explanation_parts.append(f"Sorts results by `{order_by_clause}`.")
        if limit_clause:
            explanation_parts.append(f"Limits output to {limit_clause} records.")

        explanation = " ".join(explanation_parts)
        return final_sql, explanation, clauses_dict, 0.95

    # ── Main Entrypoint ───────────────────────────────────────────────────────
    def generate_sql_from_prompt(
        self,
        prompt: str,
        custom_schema: Optional[str] = None,
        dialect: str = "sqlite",
        api_key: Optional[str] = None,
        provider: Optional[str] = None
    ) -> Dict[str, Any]:
        p = prompt.strip()
        if not p:
            return {'success': False, 'error': 'Prompt cannot be empty'}

        active_schema = self.get_database_schema() if self.active_db_path else None
        schema_str = self._format_schema_for_prompt(active_schema, custom_schema)

        # Detect active API Key & Provider
        effective_key = api_key or os.environ.get("GEMINI_API_KEY") or os.environ.get("GOOGLE_API_KEY") or os.environ.get("GROQ_API_KEY") or os.environ.get("OPENAI_API_KEY")
        selected_provider = provider or ("gemini" if (os.environ.get("GEMINI_API_KEY") or os.environ.get("GOOGLE_API_KEY")) else "groq" if os.environ.get("GROQ_API_KEY") else "openai" if os.environ.get("OPENAI_API_KEY") else None)

        llm_result = None

        if effective_key:
            if selected_provider == "groq" or effective_key.startswith("gsk_"):
                llm_result = self._call_groq(p, schema_str, dialect, effective_key)
                used_provider = "Groq LLaMA 3.3 (AI Engine)"
            elif selected_provider == "openai" or effective_key.startswith("sk-"):
                llm_result = self._call_openai_compatible(p, schema_str, dialect, effective_key)
                used_provider = "OpenAI GPT-4o-mini (AI Engine)"
            else:
                llm_result = self._call_gemini(p, schema_str, dialect, effective_key)
                used_provider = "Google Gemini 1.5 Flash (AI Engine)"

        if llm_result:
            sql_query, explanation, clauses = llm_result
            # Ensure query has proper termination
            sql_query = sql_query.strip().rstrip(';') + ';'
            confidence = 0.99
            ai_enabled = True
        else:
            sql_query, explanation, clauses, confidence = self._build_sql_rule_based(
                prompt=p,
                active_schema=active_schema,
                custom_schema=custom_schema,
                dialect=dialect.lower()
            )
            used_provider = "Smart Local Rule Engine (Offline Mode)"
            ai_enabled = False

        return {
            'success': True,
            'prompt': p,
            'sql_query': sql_query,
            'explanation': explanation,
            'clauses_breakdown': clauses,
            'dialect': dialect,
            'has_database': bool(self.active_db_path),
            'active_database': self.active_db_name,
            'confidence': confidence,
            'provider': used_provider,
            'ai_enabled': ai_enabled
        }

    def execute_query(self, sql_query: str, limit: int = 100) -> Dict[str, Any]:
        cleaned_query = sql_query.strip().rstrip(';')
        prohibited = ['DROP', 'DELETE', 'TRUNCATE', 'ALTER', 'INSERT', 'UPDATE', 'ATTACH', 'DETACH']
        first_word = cleaned_query.split()[0].upper() if cleaned_query.split() else ''
        if first_word not in ['SELECT', 'WITH', 'PRAGMA', 'EXPLAIN']:
            for kw in prohibited:
                if re.search(r'\b' + kw + r'\b', cleaned_query, re.IGNORECASE):
                    return {
                        'success': False,
                        'error': f'Execution of {kw} statement is blocked for safety. Only read-only queries (SELECT) are permitted.'
                    }

        if self.active_db_path and os.path.exists(self.active_db_path):
            return self._run_on_sqlite(self.active_db_path, cleaned_query, limit)
        return self._run_on_mock_db(cleaned_query, limit)

    def _run_on_sqlite(self, db_path: str, query: str, limit: int) -> Dict[str, Any]:
        start = time.perf_counter()
        try:
            conn = sqlite3.connect(db_path)
            conn.row_factory = sqlite3.Row
            cur = conn.cursor()
            cur.execute(query)
            rows = cur.fetchmany(limit)
            elapsed_ms = round((time.perf_counter() - start) * 1000, 2)
            if not cur.description:
                conn.close()
                return {'success': True, 'columns': [], 'rows': [], 'row_count': 0, 'execution_time_ms': elapsed_ms}
            columns = [desc[0] for desc in cur.description]
            result_rows = [[item for item in row] for row in rows]
            conn.close()
            return {
                'success': True,
                'columns': columns,
                'rows': result_rows,
                'row_count': len(result_rows),
                'execution_time_ms': elapsed_ms,
                'is_mock': False
            }
        except Exception as e:
            return {'success': False, 'error': str(e)}

    def _run_on_mock_db(self, query: str, limit: int) -> Dict[str, Any]:
        start = time.perf_counter()
        try:
            mem_conn = sqlite3.connect(':memory:')
            mem_cur = mem_conn.cursor()
            mem_cur.execute("CREATE TABLE employees (id INT, name TEXT, department TEXT, salary REAL, status TEXT);")
            mem_cur.executemany("INSERT INTO employees VALUES (?, ?, ?, ?, ?)", [
                (1, 'Alice Smith', 'Engineering', 125000, 'Active'),
                (2, 'Bob Jones', 'Marketing', 85000, 'Active'),
                (3, 'Charlie Brown', 'Engineering', 110000, 'Active'),
                (4, 'Diana Prince', 'Sales', 95000, 'Active'),
                (5, 'Ethan Hunt', 'Product', 130000, 'Active'),
                (6, 'Fiona Gallagher', 'Engineering', 140000, 'Active')
            ])
            mem_cur.execute("CREATE TABLE customers (id INT, name TEXT, city TEXT, country TEXT, spending REAL);")
            mem_cur.executemany("INSERT INTO customers VALUES (?, ?, ?, ?, ?)", [
                (1, 'John Doe', 'New York', 'USA', 1250.0),
                (2, 'Sarah Connor', 'San Francisco', 'USA', 3400.5),
                (3, 'Bruce Wayne', 'Chicago', 'USA', 9800.0),
                (4, 'Clark Kent', 'Metropolis', 'USA', 450.0),
                (5, 'Peter Parker', 'New York', 'USA', 320.0)
            ])
            mem_cur.execute("CREATE TABLE products (id INT, name TEXT, category TEXT, price REAL, stock INT, rating REAL);")
            mem_cur.executemany("INSERT INTO products VALUES (?, ?, ?, ?, ?, ?)", [
                (1, 'Wireless Headphones', 'Audio', 199.99, 50, 4.8),
                (2, 'Mechanical Keyboard', 'Computing', 129.99, 120, 4.7),
                (3, 'Gaming Mouse', 'Computing', 69.99, 200, 4.6),
                (4, '4K Monitor', 'Display', 349.50, 30, 4.9),
                (5, 'Smart Watch', 'Wearables', 159.00, 80, 4.4)
            ])
            mem_cur.execute("CREATE TABLE table_name (id INT, employee_name TEXT, department TEXT, salary REAL, amount REAL, price REAL, status TEXT);")
            mem_cur.executemany("INSERT INTO table_name VALUES (?, ?, ?, ?, ?, ?, ?)", [
                (1, 'Sample Record 1', 'Engineering', 100000, 500, 99.9, 'Active'),
                (2, 'Sample Record 2', 'Marketing', 75000, 1200, 149.9, 'Active'),
                (3, 'Sample Record 3', 'Sales', 90000, 800, 49.9, 'Active')
            ])
            mem_cur.execute(query)
            rows = mem_cur.fetchmany(limit)
            elapsed_ms = round((time.perf_counter() - start) * 1000, 2)
            columns = [desc[0] for desc in mem_cur.description] if mem_cur.description else []
            result_rows = [[item for item in row] for row in rows]
            mem_conn.close()
            return {
                'success': True,
                'columns': columns,
                'rows': result_rows,
                'row_count': len(result_rows),
                'execution_time_ms': elapsed_ms,
                'is_mock': True,
                'note': 'Executed on in-memory preview table.'
            }
        except Exception as e:
            return {
                'success': True,
                'columns': ['Query Status', 'SQL Preview'],
                'rows': [['SQL Generated Successfully', 'Connect a matching database or run on your SQL server.']],
                'row_count': 1,
                'execution_time_ms': round((time.perf_counter() - start) * 1000, 2),
                'is_mock': True,
                'note': str(e)
            }