import gradio as gr
from langchain_groq import ChatGroq
from langchain_community.graphs.networkx_graph import NetworkxEntityGraph
from langchain.chains import GraphQAChain
from langchain_community.document_loaders import TextLoader
from langchain.text_splitter import CharacterTextSplitter
from langchain_community.vectorstores import Pinecone
from langchain_community.embeddings import HuggingFaceEmbeddings  # Correct import
from langchain_community.llms import HuggingFaceHub  # Correct import
from langchain.schema.runnable import RunnablePassthrough
from langchain.schema.output_parser import StrOutputParser
from langchain.prompts import PromptTemplate
from langchain_core.documents import Document
from neo4j import GraphDatabase
import networkx as nx
import pinecone
import os
from datetime import datetime
import gspread
from oauth2client.service_account import ServiceAccountCredentials

os.system("pip install sentence-transformers")
os.system("pip install gspread oauth2client")
os.system("pip install -U langchain-huggingface langchain-community")

# Google Sheets Setup
scope = ["https://spreadsheets.google.com/feeds", "https://www.googleapis.com/auth/drive"]
creds = ServiceAccountCredentials.from_json_keyfile_name(r"./gen-lang-client-0300122402-f2f67b3e8c27.json", scope)
client = gspread.authorize(creds)

spreadsheet_id = "10953i4ZOhvpNAyyrFdo-4TvVwfVWppt0H2xXM7yYQF4"
sheet = client.open_by_key(spreadsheet_id).sheet1

def store_feedback_in_sheet(feedback, question, rag_response, graphrag_response):
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    row = [timestamp, question, rag_response, graphrag_response, feedback]
    sheet.append_row(row)

def load_data():
    data = sheet.get_all_records()
    return data[-10:], len(data)

def add_review(question, rag_response, graphrag_response, feedback):
    store_feedback_in_sheet(feedback, question, rag_response, graphrag_response)
    return load_data()

# RAG Setup
text_path = r"./text_chunks.txt"
loader = TextLoader(text_path, encoding='utf-8')
documents = loader.load()
text_splitter = CharacterTextSplitter(chunk_size=3000, chunk_overlap=4)
docs = text_splitter.split_documents(documents)

embeddings = HuggingFaceEmbeddings()

# Define the repo ID and connect to Mixtral model on Huggingface
repo_id = "meta-llama/Meta-Llama-3-8B"
llm = HuggingFaceHub(
  repo_id=repo_id, 
  model_kwargs={"temperature": 0.8, "top_k": 50}, 
  huggingfacehub_api_token=os.getenv('HUGGINGFACEHUB_API_TOKEN')
)

pinecone.init(
    api_key=os.getenv('PINECONE_API_KEY', '6396a319-9bc0-49b2-97ba-400e96eff377'),
    environment='gcp-starter'
)

index_name = "langchain-demo"
if index_name not in pinecone.list_indexes():
    pinecone.create_index(name=index_name, metric="cosine", dimension=768)
    docsearch = Pinecone.from_documents(docs, embeddings, index_name=index_name)
else:
    docsearch = Pinecone.from_existing_index(index_name, embeddings)

rag_llm = ChatGroq(
    model="Llama3-8b-8192",
    temperature=0,
    max_tokens=None,
    timeout=None,
    max_retries=5,
    groq_api_key='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'
)

template = """
You are a Thai rice assistant. These humans will ask you questions about Thai rice.
Answer the question only in Thai language.
Use the following piece of context to answer the question. 
If you don't know the answer, just say you don't know. 
Keep the answer within 2 sentences and concise.
Context: {context}
Question: {question}
Answer: 
"""

prompt = PromptTemplate(
  template=template, 
  input_variables=["context", "question"]
)

rag_chain = (
    {"context": docsearch.as_retriever(), "question": RunnablePassthrough()}
    | prompt
    | rag_llm
    | StrOutputParser()
)

graphrag_llm = ChatGroq(
    model="Llama3-8b-8192",
    temperature=0,
    max_tokens=None,
    timeout=None,
    max_retries=5,
    groq_api_key='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'
)

uri = "neo4j+s://46084f1a.databases.neo4j.io"
user = "neo4j"
password = "FwnX0ige_QYJk8eEYSXSF0l081mWWGIS7TFg6t8rLZc"
driver = GraphDatabase.driver(uri, auth=(user, password))

def fetch_nodes(tx):
    query = "MATCH (n) RETURN id(n) AS id, labels(n) AS labels"
    result = tx.run(query)
    return result.data()

def fetch_relationships(tx):
    query = "MATCH (n)-[r]->(m) RETURN id(n) AS source, id(m) AS target, type(r) AS relation"
    result = tx.run(query)
    return result.data()

def populate_networkx_graph():
    G = nx.Graph()
    with driver.session() as session:
        nodes = session.execute_read(fetch_nodes)  # Updated to use execute_read
        relationships = session.execute_read(fetch_relationships)  # Updated to use execute_read
        for node in nodes:
            G.add_node(node['id'], labels=node['labels'])
        for relationship in relationships:
            G.add_edge(
                relationship['source'],
                relationship['target'],
                relation=relationship['relation']
            )
    return G

networkx_graph = populate_networkx_graph()
graph = NetworkxEntityGraph()
graph._graph = networkx_graph

graphrag_chain = GraphQAChain.from_llm(
    llm=graphrag_llm,
    graph=graph,
    verbose=True
)

def get_rag_response(question):
    response = rag_chain.invoke(question)
    return response

def get_graphrag_response(question):
    system_prompt = "You are a Thai rice assistant that gives concise and direct answers. Do not explain the process, just provide the answer, provide the answer only in Thai."
    formatted_question = f"System Prompt: {system_prompt}\n\nQuestion: {question}"
    response = graphrag_chain.run(formatted_question)
    return response

def compare_models(question):
    rag_response = get_rag_response(question)
    graphrag_response = get_graphrag_response(question)
    return rag_response, graphrag_response

# Gradio Interface
with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            question_input = gr.Textbox(label="ถามคำถามเกี่ยวกับข้าว:", placeholder="Enter your question about Thai rice")
            submit_btn = gr.Button(value="ถาม")
            rag_output = gr.Textbox(label="Model A", interactive=False)
            graphrag_output = gr.Textbox(label="Model B", interactive=False)
            feedback = gr.Radio(label="Which response is better?", choices=["A ดีกว่า", "B ดีกว่า", "เท่ากัน", "แย่ทั้งคู่"])
            submit_feedback = gr.Button(value="Submit Feedback")

    # Function to handle question submission and display responses
    submit_btn.click(fn=compare_models, inputs=[question_input], outputs=[rag_output, graphrag_output])

    # Function to handle feedback submission and update data display
    submit_feedback.click(fn=add_review, inputs=[question_input, rag_output, graphrag_output, feedback])

    # Load initial data display
    demo.load(fn=load_data, inputs=None)

demo.launch(share=True)