import streamlit as st
from langchain import HuggingFaceHub
from langchain.prompts import ChatPromptTemplate
from langchain_google_genai import ChatGoogleGenerativeAI, HarmBlockThreshold, HarmCategory
from langchain.chains.question_answering import load_qa_chain
from dotenv import load_dotenv
import os
import json
import pandas as pd
import numpy as np
import sqlite3


load_dotenv()

conn = sqlite3.connect('your_database.db')
c = conn.cursor()


def get_database_schema():
    c.execute("SELECT name FROM sqlite_master WHERE type='table';")
    tables = c.fetchall()
    schema_dict = {}
    for table in tables:
        c.execute(f"PRAGMA table_info({table[0]})")
        schema = c.fetchall()
        schema_dict[table[0]] = schema
    return schema_dict


def get_chain():
    template = """
    You are a SQLite query generator.
    return output executable query only or empyt text .
    output should be json with key "query" and tabel name with key "table_name"
    Based on user question, generate SQLite Query  for given database Schema.
    Schema : \n{context}\n 
    Question : {question}       
    """
    hf_model = HuggingFaceHub(repo_id="Mistralai/Mistralaichat", model_kwargs={"temperature": 0.4})
    prompt=ChatPromptTemplate.from_template(template)
    chain=LLMChain(llm=hf_model.load_chat(),prompt=prompt)
    return chain

def get_query(txt1):
   start_index = txt1.find('{')
   end_index = txt1.find('}') + 1
   json_string = txt1[start_index:end_index]
   data = json.loads(json_string)
 
   return data


def query_with_database(query, tablename=None):
    if tablename:
        c.execute(query)
    else:
        c.execute(query)
    query_output = c.fetchall()
    df = pd.DataFrame(query_output, columns=[desc[1] for desc in c.description])
    st.table(df)


def user_input(user_question):
    schema = get_database_schema()
    chain = get_chain()
    response = chain.run({"context": schema, "question": user_question})
    generated_query = get_query(response)
    st.write("Generated Query : ", generated_query["query"])
    query_with_database(generated_query["query"])


def main():
   st.header("Chat With SQLite DATABASE - HUGGING FACE MISTRALAI/MISTRALAICHAT")
   with st.sidebar:
       st.write("Develop by : Rehan Shekh")
       st.write("Model : Hugging Face Mistralai/Mistralaichat")
   buttoclick=st.text_input(label="Ask from Database")
   st.write("Demo Database only have one table :- User_details")
   if buttoclick:
   
      user_input(str(buttoclick))


if __name__=="__main__":
    main()