from langchain_community.utilities import SQLDatabase
from langchain_openai import ChatOpenAI
from langchain_community.utilities import SQLDatabase
from langchain_openai import ChatOpenAI
# from langchain_classic.chains import create_sql_query_chain
from langchain.chains import create_sql_query_chain
from langchain_core.prompts import PromptTemplate
import os

def ask_question(question: str):

    # Connect to your DB (same as FastAPI)
    db = SQLDatabase.from_uri("sqlite:///./manabCommercial1.db")
    # db = SQLDatabase.from_uri("sqlite:////data/manabCommercial1.db")
    # LLM (use your API key)
    api_key = os.getenv("OPENAI_API_KEY")
    if api_key:
        os.environ["OPENAI_API_KEY"] = api_key
    llm = ChatOpenAI(
        model="gpt-4o-mini", 
        temperature=0
        )
    table_info = db.get_table_info()
    print("Table Info:", table_info)  # Debugging line to check schema retrieval
    top_k = 5  # Limit to top 5 rows
    input=question
    # Custom prompt for your schema
    PROMPT = """
    You are a SQLite expert. Given an input question, create a syntactically correct SQLite query to run.

    Schema: {table_info}
    

    Only use: SELECT from manabCommercial1. No INSERT/UPDATE/DELETE.
    Limit to top {top_k} rows.

    Question: {input}
    SQLQuery:
    """
    prompt = PromptTemplate.from_template(PROMPT)

    # Create chain
    chain = create_sql_query_chain(llm, db, prompt, top_k=top_k)

    # Query example
    # question = "Show status for EFileNo containing 'ABC' where completed=Yes"

    sql_query = chain.invoke({"input": question})
    print(sql_query)  # "SELECT status FROM manabCommercial1 WHERE EFileNo LIKE '%ABC%' AND completed='Yes'"

    # Execute + Answer
    result = db.run(sql_query)
    full_response = llm.invoke(f"Query: {sql_query}\nResult: {result}\nAnswer the original question: {question}")
    # print(full_response.content)
    return full_response.content