import os
import openai
from io import StringIO 
import sys

import gradio as gr
from langchain.agents import create_csv_agent
from langchain.chat_models import AzureChatOpenAI
from langchain.chat_models import ChatOpenAI
from langchain.llms import AzureOpenAI

# os.environ["OPENAI_API_TYPE"] = openai.api_type = "azure"
# os.environ["OPENAI_API_BASE"] = openai.api_base = "https://aoai-southcentral.openai.azure.com/"
os.environ["OPENAI_API_KEY"] = openai.api_key = "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"
# openai.api_version = "2023-03-15-preview"



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(input_file, question):
    llm = ChatOpenAI(temperature=0)
    agent = create_csv_agent(llm, input_file.name, verbose=True)
    question = question
    with Capturing() as printed_text:
        try:
            answer = agent.run(question)
        except Exception as e:
            answer = "LLM: " + str(e)
       
    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

with gr.Blocks(css="footer {visibility: hidden}", title="CSV Copilot") as demo:
    csv_file = gr.State([])
    csv_file = gr.File(label="CSV File", accept=".csv")
    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=[csv_file, question], outputs=text_box)
    


demo.launch()