"""Загрузка документов из директории в векторное хранилище."""
from __future__ import annotations

import argparse
from pathlib import Path

from knowledge_base import KnowledgeBase

SUPPORTED_SUFFIXES = {".txt", ".md", ".markdown"}


def load_directory(knowledge_dir: Path, kb: KnowledgeBase | None = None) -> int:
    kb = kb or KnowledgeBase()
    if not knowledge_dir.exists():
        knowledge_dir.mkdir(parents=True, exist_ok=True)
        print(f"Создана пустая директория: {knowledge_dir}")
        return 0

    total_chunks = 0
    files = [p for p in knowledge_dir.rglob("*") if p.suffix.lower() in SUPPORTED_SUFFIXES]
    if not files:
        print(f"Нет файлов {SUPPORTED_SUFFIXES} в {knowledge_dir}")
        return 0

    for file_path in sorted(files):
        chunks = kb.add_file(file_path)
        total_chunks += chunks
        print(f"  + {file_path.name}: {chunks} чанк(ов)")

    print(f"Готово: {len(files)} файл(ов), {total_chunks} чанк(ов).")
    return total_chunks


def main() -> None:
    parser = argparse.ArgumentParser(description="Инициализация базы знаний из папки")
    parser.add_argument(
        "--dir",
        default="./knowledge",
        help="Путь к директории с документами (.txt, .md)",
    )
    args = parser.parse_args()
    load_directory(Path(args.dir))


if __name__ == "__main__":
    main()