"""
init_knowledge_base.py

Script for loading documents from a directory into the vector store.

Usage:
    python init_knowledge_base.py <documents_directory>

Supported file types: .txt, .md
"""

import sys
import os
from pathlib import Path
from vector_store import add_documents

SUPPORTED_EXTENSIONS = {".txt", ".md"}


def load_documents_from_directory(directory: str) -> None:
    dir_path = Path(directory)
    if not dir_path.exists() or not dir_path.is_dir():
        print(f"Error: '{directory}' is not a valid directory.")
        sys.exit(1)

    files = [
        f for f in dir_path.rglob("*")
        if f.is_file() and f.suffix.lower() in SUPPORTED_EXTENSIONS
    ]

    if not files:
        print(f"No supported files found in '{directory}'.")
        return

    print(f"Found {len(files)} file(s) to load.\n")
    total_chunks = 0

    for file_path in files:
        try:
            content = file_path.read_text(encoding="utf-8")
            title = file_path.stem
            num_chunks = add_documents(content=content, title=title)
            total_chunks += num_chunks
            print(f"  [OK] '{file_path.name}' -> {num_chunks} chunk(s) added.")
        except Exception as e:
            print(f"  [ERROR] '{file_path.name}': {e}")

    print(f"\nDone. Total chunks added: {total_chunks}")


def main():
    if len(sys.argv) < 2:
        print("Usage: python init_knowledge_base.py <documents_directory>")
        sys.exit(1)

    directory = sys.argv[1]
    print(f"Loading documents from: {directory}\n")
    load_documents_from_directory(directory)


if __name__ == "__main__":
    main()