# So the graphRAG has been pre-built this just adds the knowledge graph + embeddings into a retriever object
# To build a new GraphRAG, first define the new text to be used to build in the input folder
# After that go to settings.yaml (define the chunking strategy, API to be used and all other strategies to be used)
# After that run command in your terminal python -m graphrag.index --root PATH (can be just . instead of PATH if already in the appropriate directory) 
from graphrag.query.indexer_adapters import (
    read_indexer_entities,
    read_indexer_relationships,
    read_indexer_reports,
    read_indexer_text_units,
)
import tiktoken

from graphrag.query.context_builder.entity_extraction import EntityVectorStoreKey

from graphrag.query.llm.oai.chat_openai import ChatOpenAI
from graphrag.query.llm.oai.typing import OpenaiApiType
from graphrag.query.structured_search.local_search.mixed_context import LocalSearchMixedContext
from graphrag.query.llm.oai.embedding import OpenAIEmbedding

import pandas as pd

from graphrag.query.indexer_adapters import (
    read_indexer_entities,
    read_indexer_relationships,
    read_indexer_reports,
    read_indexer_text_units,
)
from graphrag.vector_stores.lancedb import LanceDBVectorStore
from graphrag.query.input.loaders.dfs import store_entity_semantic_embeddings
import os
# All api defined to run the embeddings
from dotenv import load_dotenv

load_dotenv()

api_key = os.getenv('OPENAI_API_KEY')
llm_model = "gpt-4o-mini"
llm = ChatOpenAI(
    api_key=api_key,
    model=llm_model,
    api_type=OpenaiApiType.OpenAI,
    max_retries=20,
)
embedding_model = "text-embedding-3-small"

token_encoder = tiktoken.get_encoding("cl100k_base")
text_embedder = OpenAIEmbedding(api_key=api_key, api_type=OpenaiApiType.OpenAI, model=embedding_model, max_retries=20)


INPUT_DIR = os.path.join('output', '20240727-210235', 'artifacts')

# Define the table names
ENTITY_TABLE = "create_final_nodes"
ENTITY_EMBEDDING_TABLE = "create_final_entities"
RELATIONSHIP_TABLE = "create_final_relationships"
COMMUNITY_REPORT_TABLE = "create_final_community_reports"
TEXT_UNIT_TABLE = "create_final_text_units"

# Construct file paths
entity_file = os.path.join(INPUT_DIR, f"{ENTITY_TABLE}.parquet")
report_file = os.path.join(INPUT_DIR, f"{COMMUNITY_REPORT_TABLE}.parquet")
entity_embedding_file = os.path.join(INPUT_DIR, f"{ENTITY_EMBEDDING_TABLE}.parquet")
relationship_file = os.path.join(INPUT_DIR, f"{RELATIONSHIP_TABLE}.parquet")
text_unit_file = os.path.join(INPUT_DIR, f"{TEXT_UNIT_TABLE}.parquet")

LANCEDB_URI = os.path.join(INPUT_DIR, 'lancedb')


COMMUNITY_LEVEL = 1
# Read the files
entity_df = pd.read_parquet(entity_file)
report_df = pd.read_parquet(report_file)
entity_embedding_df = pd.read_parquet(entity_embedding_file)
relationship_df = pd.read_parquet(relationship_file)
text_unit_df = pd.read_parquet(text_unit_file)



reports = read_indexer_reports(report_df, entity_df, COMMUNITY_LEVEL)
entities = read_indexer_entities(entity_df, entity_embedding_df, COMMUNITY_LEVEL)
relationships = read_indexer_relationships(relationship_df)
text_units = read_indexer_text_units(text_unit_df)


description_embedding_store = LanceDBVectorStore(collection_name="entity_description_embeddings")
description_embedding_store.connect(db_uri=LANCEDB_URI)
store_entity_semantic_embeddings(entities=entities, vectorstore=description_embedding_store)

# Add everything into a LocalSearch-based retriever!
graph_rag_retriever = LocalSearchMixedContext(
    community_reports=reports,
    text_units=text_units,
    entities=entities,
    relationships=relationships,
    entity_text_embeddings=description_embedding_store,
    embedding_vectorstore_key=EntityVectorStoreKey.ID,
    text_embedder=text_embedder,
    token_encoder=token_encoder,
)