import os
from dotenv import load_dotenv
from langchain_community.document_loaders import WebBaseLoader, TextLoader
from langchain_chroma import Chroma
from langchain_ollama import OllamaEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter

load_dotenv()

TEXT_FILE_PATH = "./data/aboutPromtior.txt"
WEBSITE_URL = os.getenv("WEBSITE_URL")
VECTOR_STORE_PATH = "./vectorstore"
USER_AGENT = os.getenv("USER_AGENT")



def main():
    generate_data_store()

def generate_data_store():
    documents = load_documents()
    chunks = split_text(documents)
    save_to_chroma(chunks)

def load_documents():
    doc_text = load_text_content()
    doc_web = load_web_content(WEBSITE_URL)
    return doc_text + doc_web


def load_text_content():
    loader = TextLoader(TEXT_FILE_PATH)
    documents = loader.load()
    return documents

def load_web_content(url):
    loader = WebBaseLoader(url)
    documents = loader.load()
    return documents

def split_text(documents):
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200, add_start_index=True)
    chunks = text_splitter.split_documents(documents)
    return chunks

def save_to_chroma(chunks):
    embeddings = OllamaEmbeddings(model="llama3")
    Chroma.from_documents(
        chunks,
        embeddings,
        persist_directory=VECTOR_STORE_PATH,
    )
    print(f"Saved {len(chunks)} chunks to {VECTOR_STORE_PATH}")

if __name__ == "__main__":
    main()