backend/app/services/sql_generator.py

from langchain_community.utilities import SQLDatabase
from langchain.chains import create_sql_query_chain
from langchain_openai import ChatOpenAI  # or your LLaMA adapter
from config.settings import settings
from langchain.prompts import PromptTemplate, ChatPromptTemplate
from app.db.schema_reader import get_schema
from app.db.db_connector import get_db_engine
from app.db.db_connector import get_connection
import re

# llm = ChatOpenAI(model="gpt-4o-mini", temperature=0, api_key=settings.OPENAI_API_KEY)
  
print(f"OPENAI_API_KEY: {settings.OPENAI_API_KEY}")
    
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0, api_key=settings.OPENAI_API_KEY)

def detect_foreign_key_relationships_only(schema):
    """Only detect foreign key relationships without any assumptions about data types"""
    relationships = []
    
    for table_name, columns in schema.items():
        for column in columns:
            # Look for foreign key patterns: *_id columns
            if column.endswith('_id') and column != 'id':
                # Extract potential referenced table name
                ref_name = column[:-3]  # Remove '_id'
                
                # Check for exact match or plural/singular variations
                possible_refs = [ref_name, ref_name + 's', ref_name[:-1] if ref_name.endswith('s') else ref_name + 's']
                
                for possible_ref in possible_refs:
                    if possible_ref in schema:
                        relationships.append({
                            'from_table': table_name,
                            'from_column': column,
                            'to_table': possible_ref,
                            'to_column': 'id'
                        })
                        break
    
    return relationships

def generate_pure_schema_context(schema):
    """Generate schema context without any assumptions - let AI figure out everything"""
    
    # Just provide the raw schema structure
    context = "DATABASE SCHEMA:\n\n"
    
    # List all tables and their columns
    context += "TABLES AND COLUMNS:\n"
    for table_name, columns in schema.items():
        context += f"- {table_name}: {', '.join(columns)}\n"
    
    # Detect only foreign key relationships (no assumptions about data types)
    relationships = detect_foreign_key_relationships_only(schema)
    
    # Show detected relationships
    if relationships:
        context += "\nDETECTED FOREIGN KEY RELATIONSHIPS:\n"
        for rel in relationships:
            context += f"- {rel['from_table']}.{rel['from_column']} → {rel['to_table']}.{rel['to_column']}\n"
    
    return context


def generate_sql_query(user_question: str):
    schema = get_schema()
    
    # Generate pure schema context - no assumptions, let AI figure it out
    schema_context = generate_pure_schema_context(schema)

    prompt = ChatPromptTemplate.from_template("""
    You are an expert SQLite query generator with advanced analytical capabilities.
    
    {schema_context}

    INTELLIGENT ANALYSIS INSTRUCTIONS:
    1. EXAMINE the schema above - look at table names and column names to understand what data each table contains
    2. ANALYZE the column names to infer what type of information is stored (names, addresses, amounts, dates, etc.)
    3. USE the detected foreign key relationships to understand how tables connect
    4. For queries needing data from multiple tables, follow the relationship paths and use JOINs
    5. INFER from context - if user asks about location and you see city/address columns, use those tables
    6. THINK logically about which tables would contain the requested information

    QUERY GENERATION STRATEGY:
    - Read the user's question carefully
    - Identify what data they want (customers, transactions, locations, etc.)
    - Look at the schema to find which tables likely contain that data
    - Check if multiple tables are needed and use the foreign key relationships
    - Generate appropriate SQL with JOINs if data spans multiple tables

    RESPONSE FORMAT:
    - Return ONLY valid SQLite syntax
    - No explanations, comments, or extra text
    - Use proper JOIN syntax when needed
    - If you cannot determine how to query the available schema, respond: "I cannot generate this query."

    Example thinking process:
    - User asks "customers in New York" → Look for tables with customer info AND location info → Use relationships to JOIN them
    - User asks "transactions over $1000" → Look for tables with transaction/amount data
    - User asks "branch managers" → Look for tables with employee/staff info and management roles

    User question: {question}
    """)

    messages = prompt.format_messages(schema_context=schema_context, question=user_question)
    sql_query = llm.invoke(messages).content.strip()

    # If the LLM says it cannot generate a query, treat as normal chat
    if sql_query.lower().startswith("i cannot generate this query"):
        # Use the LLM as a chatbot for normal conversation
        chat_prompt = (
            "You are an AI assistant whose primary role is to generate SQL queries from natural language questions. "
            "If the user's question is not related to SQL or databases, respond conversationally as a helpful assistant."
            f"\n\nUser question: {user_question}"
        )
        chat_response = llm.invoke([{"role": "user", "content": chat_prompt}]).content.strip()
        return {"chat_message": chat_response}

    return sql_query