from langchain.agents import create_sql_agent
from langchain.agents.agent_toolkits import SQLDatabaseToolkit
from langchain.sql_database import SQLDatabase
from langchain.llms.openai import OpenAI
from langchain.agents import AgentExecutor
import urllib, os

from io import StringIO 
import sys

class Capturing(list):
    def __enter__(self):
        self._stdout = sys.stdout
        sys.stdout = self._stringio = StringIO()
        return self
    def __exit__(self, *args):
        self.extend(self._stringio.getvalue().splitlines())
        del self._stringio    # free up some memory
        sys.stdout = self._stdout



def answer_question(question):
    with Capturing() as printed_text:
        answer = agent_executor.run("what are the top 3 most expensive items and how many customers bought them?")
    import re
    text = '\n'.join(printed_text) + '\n' + str(answer)
    # Remove all escape characters
    text = re.sub(r"\x1b\[\d+(;\d+)?m", "", text)

    # Remove all characters inside angle brackets
    text = re.sub(r"<.*?>", "", text)

    # Remove all leading/trailing whitespaces
    text = text.strip()
    return text

db = SQLDatabase.from_uri("mssql+pyodbc:///?odbc_connect=Driver={ODBC Driver 18 for SQL Server};Server=tcp:tesserversean.database.windows.net,1433;Database=testdb-sean;Uid=sean;Pwd=abc123456!;Encrypt=yes;TrustServerCertificate=no;Connection Timeout=30;")
toolkit = SQLDatabaseToolkit(db=db)


agent_executor = create_sql_agent(
    llm = OpenAI(model_name="gpt-4", temperature=0.0),
    toolkit=toolkit,
    verbose=True
)



import gradio as gr

with gr.Blocks(css="footer {visibility: hidden}", title="SQL Chat") as demo:
    csv_file = gr.State([])
    question = gr.Textbox(label="Question")
    ask_question = gr.Button(label="Ask Question")
    text_box = gr.TextArea(label="Output", lines=10)

    ask_question.click(answer_question, inputs=[question], outputs=text_box)
    


demo.launch()