import pandas as pd
from langchain_community.document_loaders import CSVLoader
from langchain.schema import Document
from langchain_openai import OpenAIEmbeddings
from langchain_community.vectorstores import Chroma
import openai 
from dotenv import load_dotenv
import os
import shutil

# Load environment variables. Assumes that project contains .env file with API keys
load_dotenv()
#---- Set OpenAI API key
openai.api_key = os.environ['OPENAI_API_KEY']

CHROMA_PATH = "chroma"
DATA_PATH = "data/wikis/patient_reviews_with_symptoms_automated.csv"

import nltk
nltk.download('punkt_tab')
nltk.download('averaged_perceptron_tagger_eng')

def main():
    generate_data_store()


def generate_data_store():
    documents = load_documents_with_pandas()
    save_to_chroma(documents)

def load_documents():
    loader = CSVLoader(DATA_PATH, encoding="windows-1252")
    documents = loader.load()
    return documents

def load_documents_with_pandas():
    # Read CSV file using Pandas
    df = pd.read_csv(DATA_PATH, encoding="utf-8")

    # Convert each row to a Document object
    documents = [
        Document(
            page_content=row['review'],
            metadata={"rating": row['rating text'],"DrugName": row['drugName'],"condition": row['condition']}
        )
        for _, row in df.iterrows()
    ]
    return documents

def save_to_chroma(chunks: list[Document]):
    # Clear out the database first.
    if os.path.exists(CHROMA_PATH):
        shutil.rmtree(CHROMA_PATH)

    # Create a new DB from the documents.
    db = Chroma.from_documents(
        chunks, OpenAIEmbeddings(), persist_directory=CHROMA_PATH
    )
    db.persist()
    print(f"Saved {len(chunks)} chunks to {CHROMA_PATH}.")


if __name__ == "__main__":
    main()