import gradio as gr
from huggingface_hub import InferenceClient
import os
from groq import Groq
from sqlalchemy import text

"""
For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
"""

import pandas as pd
import pandasql
from collections import defaultdict

# Create a sample DataFrame
data = [
    {"Name": "John", "Age": 25, "Gender": "male", "Votes": 100},
    {"Name": "Mary", "Age": 30, "Gender": "female", "Votes": 200},
    {"Name": "Bob", "Age": 28, "Gender": "male", "Votes": 150},
    {"Name": "Alice", "Age": 24, "Gender": "female", "Votes": 120},
]

# types = defaultdict(str,'Votes'=int,'Votes Percentage'=float)
# {'Votes':int,'Votes Percentage':float}
#
# Create a pandas dataframe from the list of dictionaries
df = pd.read_csv('election_results.csv',dtype={'Votes':'int64'})


def respond(
    message,
    history: list[tuple[str, str]],
):

    client = Groq()
    messages = [
        {
            "role": "system",
            "content": "You are an election result analysis bot. Your task is to convert the user's natural language query into a SQL SELECT statement suitable for a pandas DataFrame named df.\n\nAvailable Columns:\n\nCandidate (string)\nParty (string)\nVotes (integer)\nVotes Percentage (decimal)\nConstituency (string)\nState (string)\n\nCase-Insensitive Text Comparisons:\n\nUse UPPER function to convert text columns (Candidate, Party, Constituency, State) to uppercase for case-insensitive comparisons.\nEmploy LIKE operator with wildcards (%) for pattern matching when appropriate.\n\nUser Input:\n\nThe system will provide the user's query in natural language.\n\nOutput Format:\n\nGenerate the SQL SELECT statement only, starting with SELECT. Do not add any other extra instruction text, your output will be directly executed, so send only the SQL Statement\n\nExample:\n\nUser Input: Find all candidates from Coimbatore who received more than 50% of the votes.\n\nYour Output:\nSELECT * FROM df\nWHERE UPPER(State) = 'COIMBATORE' AND Votes Percentage > 50"
        }
    ]
    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})

    messages.append({"role": "user", "content": message})
    completion = client.chat.completions.create(
        model="llama3-70b-8192",
        messages=messages,
        temperature=1,
        max_tokens=2048,
        top_p=1,
        stream=False,
        stop=None,
    )

    sql_command = completion.choices[0].message.content
    print(sql_command)

    if  sql_command.startswith('SELECT'):
        result = pandasql.sqldf(sql_command, globals())
        yield result.to_string()
    else:
        yield str(sql_command)

"""
For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
"""
demo = gr.ChatInterface(
    respond
)


if __name__ == "__main__":
    demo.launch()