from langchain.embeddings import HuggingFaceEmbeddings
#from langchain.vectorstores import FAISS
from langchain.schema import Document

#from langchain.vectorstores import Chroma 
from langchain.llms import AzureMLOnlineEndpoint
from langchain.chat_models.azureml_endpoint import ContentFormatterBase
import json

from langchain.chains import create_sql_query_chain

import chainlit as cl


from typing import Dict


# Now we can create the agent, adjusting the standard SQL Agent suffix to consider our use case. 
# Although the most straightforward way to handle this would be to include it just in the tool description, 
# this is often not enough and we need to specify it in the agent prompt using the suffix argument in the constructor.

from langchain.agents import create_sql_agent, AgentType
from langchain.agents.agent_toolkits import SQLDatabaseToolkit
from langchain.utilities import SQLDatabase
from langchain.chat_models import ChatOpenAI

import os
OPENAI_API_KEY = os.environ['OPENAI_API_KEY']




def create_agent():
   # conn_str = "mssql+pyodbc://" + SQL_USR_NM + ":" + PWD + "@" + SQL_HOST + "/" + SQL_TBL + "?driver=ODBC+Driver+18+for+SQL+Server"

    # Create the SQLDatabase object
    db = SQLDatabase.from_uri('sqlite:///spm.db')
    llm = ChatOpenAI(temperature=0.05, model="gpt-3.5-turbo-16k-0613")
    db_chain = SQLDatabaseChain.from_llm(llm, db, verbose=True)
    return db_chain

    #toolkit = SQLDatabaseToolkit(db=db, llm=llm)

custom_suffix = """
Compose a query in the All_data table in the db database.
Here is a description of each column:
destn_area_name: The name of the destination area.
destn_district_name: The name of the destination district.
score: The score of the destination area.
avg_days_todelr: The average number of days to deliver to the destination area.
time_per: The time period of the data.
orgn_area: The code of the origin area.
orgn_dist: The code of the origin district.
orgn_area_name: The name of the origin area.
orgn_dist_name: The name of the origin district.
destn_area: The code of the destination area.
destn_dist: The code of the destination district.
destn_area_name: The name of the destination area.
destn_dist_name: The name of the destination district.
prodt: The product type.
rptg_start_date: The start date of the reporting period.
rptg_end_date: The end date of the reporting period.
mo: The month of the reporting period.
pstl_qtr: The quarter of the Postal reporting period.
pstl_yr: The year of the Postal reporting period.
score: The score of the destination area.
score_plus_1: The score of the destination area plus 1.
"""

#agent = create_sql_agent(llm=llm,
#                         toolkit=toolkit,
#                         verbose=False,
#                         agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
#                         extra_tools=custom_tool_list,
#                         suffix=custom_suffix, 
#                         handle_parsing_errors=True
#                        )

from langchain.prompts import PromptTemplate


def build_sql_chain(llm, db):
    
    dialect = "Azure SQL"
    table_info = "All_data"
    few_shots = {"What are the top 10 performing areas?": "SELECT TOP 10 destn_area_name, AVG(score) AS AvgScore FROM All_data GROUP BY destn_area_name ORDER BY AvgScore DESC",
                        "What are the worst 10 performing areas?": "SELECT TOP 10 destn_area_name, AVG(score) AS AvgScore FROM All_data GROUP BY destn_area_name ORDER BY AvgScore ASC",
                        "What districts have the highest volume of mail?": "SELECT TOP 10 destn_district_name, COUNT(*) AS Volume FROM All_data GROUP BY destn_district_name ORDER BY Volume DESC",
                        "What districts have the lowest volume of mail?": "SELECT TOP 10 destn_district_name, COUNT(*) AS Volume FROM All_data GROUP BY destn_district_name ORDER BY Volume ASC",
                        "What are the top 10 performing districts?": "SELECT TOP 10 destn_district_name, AVG(score) AS AvgScore FROM All_data GROUP BY destn_district_name ORDER BY AvgScore DESC",
                        "What are the worst 10 performing districts?": "SELECT TOP 10 destn_district_name, AVG(score) AS AvgScore FROM All_data GROUP BY destn_district_name ORDER BY AvgScore ASC",
                        "What districts gave the fastest delivery time?": "SELECT TOP 10 destn_district_name, AVG(avg_days_todelr) AS AvgDeliveryTime FROM All_data GROUP BY destn_district_name ORDER BY AvgDeliveryTime ASC"}
    fs = str(few_shots)
                        

    TEMPLATE = """Given an input question, first create a syntactically correct {dialect} query to run, then look at the results of the query and return the answer.
    Use the following format:

    Question: "Question here"
    SQLQuery: "SQL Query to run"
    SQLResult: "Result of the SQLQuery"
    Answer: "Final answer here"

    Only use the following tables:

    {table_info}.

    Some examples of SQL queries that correspond to questions are:

    \{"What are the top 10 performing areas?": "SELECT TOP 10 destn_area_name, AVG(score) AS AvgScore FROM All_data GROUP BY destn_area_name ORDER BY AvgScore DESC",
     "What are the worst 10 performing areas?": "SELECT TOP 10 destn_area_name, AVG(score) AS AvgScore FROM All_data GROUP BY destn_area_name ORDER BY AvgScore ASC",
     "What districts have the highest volume of mail?": "SELECT TOP 10 destn_district_name, COUNT(*) AS Volume FROM All_data GROUP BY destn_district_name ORDER BY Volume DESC",
     "What districts have the lowest volume of mail?": "SELECT TOP 10 destn_district_name, COUNT(*) AS Volume FROM All_data GROUP BY destn_district_name ORDER BY Volume ASC",
     "What are the top 10 performing districts?": "SELECT TOP 10 destn_district_name, AVG(score) AS AvgScore FROM All_data GROUP BY destn_district_name ORDER BY AvgScore DESC",
     "What are the worst 10 performing districts?": "SELECT TOP 10 destn_district_name, AVG(score) AS AvgScore FROM All_data GROUP BY destn_district_name ORDER BY AvgScore ASC",
     "What districts gave the fastest delivery time?": "SELECT TOP 10 destn_district_name, AVG(avg_days_todelr) AS AvgDeliveryTime FROM All_data GROUP BY destn_district_name ORDER BY AvgDeliveryTime ASC"\}

    Question: {input}"""

    CUSTOM_PROMPT = PromptTemplate(
        input_variables=["input", "table_info", "dialect"], template=TEMPLATE
    )

    # Set verbose=True to see the full prompt:
    return create_sql_query_chain(llm=llm, db=db)

#from langchain.llms import OpenAI
from langchain_experimental.sql import SQLDatabaseChain
#sql_chain = build_sql_chain(llm, db)


@cl.on_chat_start
async def main():
    # Parse the command line arguments
   # args = parse_arguments()
    await cl.Message(content="Welcome to GeoData!").send()
    
    # activate/deactivate the streaming StdOut callback for LLMs
    #callbacks = [StreamingStdOutCallbackHandler()]

    #sql_chain = build_sql_chain(llm, db)


@cl.on_message
async def msg(message: str):
    # Retrieve the chain from the user session
    #sql_chain = cl.user_session.get("sql_chain")  # type: RetrievalQA
    agent = create_agent()
    m = message.content
    #res = sql_chain.invoke({"question": m})
    res = agent.run({"query": m})
    # Call the chain asynchronously

    print(res)
    await cl.Message(content=res).send()