from py2neo import Graph, Node, Relationship
import pandas as pd

# Connect to Neo4j
graph = Graph("bolt://localhost:7687", auth=("username", "password"))

def load_data(row):
    # Check for NaN in critical fields and handle it
    if pd.isna(row['director_etd_str']):
        director_name = "Unknown Director"
    else:
        director_name = row['director_etd_str']
    if pd.notna(row['all_roles']) :
        actors = row['all_roles']

    # Cypher query to create/find director and connect with movie
    query = """
    MERGE (director:Director {name: $director_name})
    CREATE (movie:Movie {id: $movie_id, title: $title})
    MERGE (actors:Actor {name: $actor_name})
    MERGE (movie)-[:DIRECTED_BY]->(director)
    MERGE (movie)-[:ACTED_BY]->(actors)
    """
    # Execute the query
    graph.run(query, director_name=director_name, movie_id=row['movieId'], title=row['titleKo'], actor_name=actors)

# 영화 노드가 이미 존재하는 경우, 청킹된 시놉시스를 연결하는 함수
def upload_to_existing_movies(df):
    for index, row in df.iterrows():
        if row['textembedding_tolist'] is not None :
            # Movie 노드를 movieId로 조회
            movie_node = graph.nodes.match("Movie", id=row['movieId']).first()

            if movie_node:
                # Synopsis 청킹된 텍스트와 임베딩 값 노드 생성 및 연결
                for chunk in row['textembedding_tolist']:
                    synopsis_node = Node("Synopsis",
                                        seqId=chunk['seqId'],
                                        text=chunk['text'],
                                        textEmbedding=chunk['textEmbedding'])
                    graph.create(synopsis_node)

                    # 관계 생성
                    rel = Relationship(movie_node, "HAS_SYNOPSIS", synopsis_node)
                    graph.create(rel)

# Applying the function to each row of the DataFrame
df.apply(load_data, axis=1)

# 데이터프레임을 Neo4j에 적재
upload_to_existing_movies(df)

CREATE VECTOR INDEX queryVector IF NOT EXISTS
FOR (s:Synopsis)
ON s.textEmbedding
OPTIONS {indexConfig: {
 `vector.dimensions`: 1536,
 `vector.similarity_function`: 'cosine'
}}

CALL db.index.vector.queryNodes('queryVector', 5, $queryVector)
YIELD node AS doc, score

from langchain.chat_models import ChatOpenAI
from langchain.schema import HumanMessage, SystemMessage, AIMessage
from langchain.prompts.prompt import PromptTemplate
from langchain_core.messages import SystemMessage
from langchain_core.messages import HumanMessage
from langchain.embeddings.openai import OpenAIEmbeddings

from retry import retry
from timeit import default_timer as timer
# import streamlit as st

from neo4j import GraphDatabase
from json import loads
import json

import os
from dotenv import load_dotenv

load_dotenv()

os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY")
openai_key = os.environ["OPENAI_API_KEY"]
os.environ["NEO4J_URI"] = os.getenv("NEO4J_URI")
os.environ["NEO4J_USERNAME"] =os.getenv("NEO4J_USERNAME")
os.environ["NEO4J_PASSWORD"] =os.getenv("NEO4J_PASSWORD")

neo4j_uri = os.environ["NEO4J_URI"]
neo4j_user = os.environ["NEO4J_USERNAME"]
neo4j_password = os.environ["NEO4J_PASSWORD"]

SYSTEM_PROMPT = """You are a movie expert who recommends movies.
* Create answers in Korean
* Don't answer the same sentence repeatedly.

"""

PROMPT_TEMPLATE = """

{questions}

Here is the context in JSON format. This dataset contains information about movies that will be recommended to the user.

<context>
{context}
</context>

When recommending movies to a user related to a question, make sure to recommend at least five movies included in the context!
Create answers in Korean

The following is an example of a response when recommending a movie to a user :
Hello! My name is 😎 Agent SH , a movie recommendation chatbot that specializes in movie recommendations. I recommend movies based on GraphRAG.
Based on your questions, I'll recommend movies you might like.

🎬 Movie Title: Monster
🎥 Film director: Bong Joon-ho
🕴️ Actors: Song Kang-ho, Bae Doo-na, etc ...
📄 Synopsis Summary: Brief synopsis summary
✓ Reasons similar to the movie : It's similar to the movie you asked about in that it features monsters.

"""
PROMPT = PromptTemplate(
    input_variables=["questions","context"], template= PROMPT_TEMPLATE
)

embeddings = OpenAIEmbeddings(model="text-embedding-3-small")

def run_query(uri, user, password, query, params):
    driver = GraphDatabase.driver(uri, auth=(user, password))
    with driver.session() as session:
        # print(params)
        result = session.run(query, params)
        # print(result)
        return [record for record in result]

def vector_graph_qa(query):
    query_vector = embeddings.embed_query(query)
    # print(query_vector)
    uri = neo4j_uri
    user = neo4j_user
    password = neo4j_password
    params = {'queryVector':query_vector}
    cypher_query = """
    CALL db.index.vector.queryNodes('queryVector', 5, $queryVector)
    YIELD node AS doc, score
    match (doc)<-[s:HAS_SYNOPSIS]-(movie:Movie)
    match (movie)-[d:DIRECTED_BY]->(director:Director)
    match (movie)-[a:ACTED_BY]->(actor:Actor)
    RETURN movie.title AS MovieTitle,
        collect('The director of the movie ' + movie.title + ' is ' + director.name +
        ', the actor is ' + actor.name + ' and the synopsis of the movie is ' + doc.text) AS MovieInfo,
        score
    ORDER BY score DESC LIMIT 5
    """
    result = run_query(uri, user, password, cypher_query, params)
    # print(result)
    return result

@retry(tries=5, delay=5)
def get_results(question):
    start = timer()
    try:
        chat_llm = ChatOpenAI(
            model = "gpt-3.5-turbo-0125",
            temperature=0,
            openai_api_key=openai_key,
        )

        df = vector_graph_qa(question)
        # print('df : ', df)
        # ctx = df_to_context(df)
        ans = PROMPT.format(questions=question, context=df)
        print('ans : ', ans)
        messages = chat_llm(
        [
            SystemMessage(content="SYSTEM_PROMPT"),
            HumanMessage(content=ans)
        ]
        )

        result = messages.content
        r = {'context': df, 'result': result}

        return result
    finally:
        print('Response Generation Time : {}'.format(timer() - start))