from langchain import SQLDatabaseChain
from langchain.sql_database import SQLDatabase
from langchain.llms.openai import OpenAI
from langchain.chat_models import ChatOpenAI
from langchain.prompts.prompt import PromptTemplate


def check_query(query):
    if query.startswith("### Query"):
        split = query.split('\n\n')
        q_text = split[0]
        t_text = split[1]

        if t_text.startswith("### Tables"):
            query_params = dict()
            tables = t_text.split('\n')
            query_params['tables'] = tables[1:]
            query_params['q'] = q_text.split('\n')[1]
            print(query_params)
            return query_params
        else:
            return 'error'
    return 'small'

def answer_question(query):
    query_check = check_query(query)
    if isinstance(query_check, dict):
        return("BIG TABLE")
    if query_check == 'small':
        return('SMALL TABLE')
    if query_check == 'error':
        return('ERROR: Wrong format for getting the big db schema')
  


    return("DONE: " + query)

if __name__ == "__main__":
    import gradio as gr
    # print(answer_question("Who is Harry's Father"))
    
    gr.Interface(
        answer_question,
        [
          gr.inputs.Textbox(lines=10, label="Query"),
        ],
        gr.outputs.Textbox(label="Response"),
        title="Ask NBA Stats",
        description=""" Ask NBA Stats is a tool that let's you ask a question with
                        the NBA SQL tables as a reference
                        
                        Ask a simple question to use the small database

                        If you would like to access the large DB use format
                        
                        ### Query
                        single line query
                        
                        ### Tables
                        tables to access line by line
                        table1
                        table2"""
    ).launch()