import time
import openai
from config import app_config
from langchain_openai import ChatOpenAI
from langchain_community.agent_toolkits import create_sql_agent
from langchain_community.utilities import SQLDatabase
from sqlalchemy import create_engine
# Load API Key
openai.api_key = app_config.OPENAI_API_KEY
db_conn = app_config.SQLALCHEMY_DATABASE_URI
#  Optimized DB Connection Pooling (No Session Required)
engine = create_engine(db_conn, pool_size=150, max_overflow=50, echo=False)
db = SQLDatabase(engine)
#  Optimized GPT Model (GPT-4-Turbo for Speed)
llm = ChatOpenAI(temperature=0, model_name='gpt-4o', max_tokens=300)
#  Initialize SQL Agent Once (Prevents Repeated Instantiations)
agent_executor = create_sql_agent(llm, db=db, agent_type="openai-tools", verbose=False)
#  Optimized SQL Query Generator with Auto-Retry
def generate_sql_query(natural_language_query, max_retries=3, base_delay=2):
    """
    Generates an SQL query using the LLM with automatic retry mechanism.
    
    - Retries up to `max_retries` times if an error occurs.
    - Uses exponential backoff (`base_delay` seconds) between retries.
    """
    attempt = 0
    while attempt < max_retries:
        try:
            #  Generate Optimized SQL Query
            result = agent_executor.invoke({"input": natural_language_query})
            # print(time.strftime("%Y%m%d-%H%M%S"))
            if result:
                return result  # Return Successful Result
            
            raise ValueError("No response from LLM")  # Trigger retry if response is empty
        except Exception as e:
            error_message = f"Unexpected error: {str(e)}"

        attempt += 1
        if attempt < max_retries:
            wait_time = base_delay * (2 ** (attempt - 1))  # Exponential Backoff
            # print(f"Attempt {attempt}/{max_retries} failed: {error_message}. Retrying in {wait_time}s...")
            time.sleep(wait_time)  # ⏳ Wait before retrying
     
    return {"error": error_message}  # Return Final Failure Message