#!/usr/bin/env python3
"""Initialization client - loads documents from a directory into the vector store."""

import os
from typing import List

from langchain_core.documents import Document
from langchain_community.documenloaders import TextLoader

 from vector_store import get_vector_store, add_documents_to_store


def load_documents_from_dir(dir_path: str) -> List[Document]:
    """Load all .txt files from a directory as LangChain Documents.

    Args:
        dir_path: Path to directory containing .txt files.

    Returns:
        List of Document objects.
    """
    documents = []
    for filename in os.listdir(dir_path):
        if filename.endswith(".txt"):
            file_path = os.path.join(dir_path, filename)
            loader = TextLoader(file_path, encoding="utf-8")
            docs = loader.load()
            # Add title metadata from filename
            title = os.path.splitext(filename)[0]
            for doc in docs:
                doc.metadata["title"] = title
            documents.extend(docs)
    return documents


def init_knowledge_base(dir_path: str) -> int:
    """Initialize the knowledge base by loading documents from a directory.

    Args:
        dir_path: Path to directory containing .txt files.

    Returns:
        Number of chunks added to the store.
    """
    store = get_vector_store()
    documents = load_documents_from_dir(dir_path)

    if not documents:
        print(f"No .txt files found in {dir_path}")
        return 0

    ids = add_documents_to_store(store, documents)
    print(f"Loaded {len(documents)} document(s) into knowledge base ({len(ids)} chunks).")
    return len(ids)


if __name__ == "__main__":
    import sys

    if len(sys.argv) < 2:
        print("Usage: python init_client.py <directory_path>")
        sys.exit(1)

    dir_path = sys.argv[1]
    if not os.path.isdir(dir_path):
        print(f"Error: {dir_path} is not a directory")
        sys.exit(1)

    init_knowledge_base(dir_path)