import os
from dotenv import load_dotenv
from utils.document_management import load_document, transcript_chunker, merge_chunks
from utils.vector_db import VectorDB
from langchain.schema import Document
from uuid import uuid4

# Load environment variables
load_dotenv()

# Initialize services
# Load all the transcriptions from the "transcription_data" directory
working_dir = os.path.join(os.getcwd(), "_backend","data" ,"honduras_data")
if os.getenv("ENVIRONMENT") == "production":
    working_dir = os.path.join(os.getcwd(), "data","honduras_data")
vector_db_name = "transcription"
vector_db_embedding_length = 1536

# First, collect all documents
all_documents = []
all_files = []

# Walk through the directory to find all *_transcribed.srt files
for root, dirs, files in os.walk(working_dir):
    for file in files:
        if file.endswith("transcribed.srt"):
            file_path = os.path.join(root, file)
            # clipping everything before honduras_data 
            relative_path = os.path.relpath(file_path, working_dir)
            all_files.append(relative_path)

for file_path in all_files:
    path = os.path.join(working_dir, file_path)
    
    docuemnt = load_document(path)

    # Chunk the document into smaller parts
    chunks = transcript_chunker(docuemnt, filename=file_path)

    merged_chunks = merge_chunks(chunks)

    # Create LangChain Document objects
    documents = [Document(page_content=chunk["text"], metadata={**chunk, "chunk_id": i, "id": str(uuid4()), "filename": file_path}) for i, chunk in enumerate(merged_chunks)]

    print(f"Processing {file_path} with {len(documents)} chunks")
    all_documents.extend(documents)


# Now create the vector database with all documents at once (this will delete and recreate)
if all_documents:
    if os.path.exists(os.path.join(working_dir, "chroma_db")):
        import shutil
        # Remove the existing vector database directory
        shutil.rmtree(os.path.join(working_dir, "chroma_db"))
        print("Removed existing vector database directory.")

        
    vector_db = VectorDB(
        collection_name=vector_db_name,
        embedding_length=vector_db_embedding_length,
        open_ai_api_key=os.getenv("OPENAI_API_KEY"),
        working_dir=working_dir,
        documents=all_documents
    )
    print(f"Created new vector database with {len(all_documents)} documents")
else:
    print("No documents found to process")

print("Vector database reset and data upload complete with ChromaDB.")