# -*- coding: utf-8 -*-
"""RAG по своей кодовой базе (№1): индексация + поиск + ответ по проекту.

Запуск:
  python rag.py index <папка>              — проиндексировать проект
  python rag.py ask "где у меня логика авторизации?" [backend]

Работает без тяжёлых зависимостей (пересечение слов). Для лучшего качества можно
поставить sentence-transformers и заменить search() на эмбеддинги с косинусной близостью.
"""
import json
import os
import re
import sys

from brain import pick, generate

INDEX = "rag_index.json"
SKIP = {"venv", ".venv", ".git", "__pycache__", "node_modules", ".idea"}


def chunk_file(path):
    code = open(path, encoding="utf-8").read()
    lines = code.splitlines()
    chunks = []
    for i in range(0, len(lines), 50):
        chunks.append((path, i + 1, "\n".join(lines[i:i + 50])))
    return chunks


def index(folder):
    chunks = []
    for root, dirs, files in os.walk(folder):
        dirs[:] = [d for d in dirs if d not in SKIP]
        for f in files:
            if f.endswith(".py"):
                chunks.extend(chunk_file(os.path.join(root, f)))
    json.dump(chunks, open(INDEX, "w", encoding="utf-8"), ensure_ascii=False)
    print(f"проиндексировано чанков: {len(chunks)} -> {INDEX}")


def search(query, top=5):
    if not os.path.exists(INDEX):
        print(f"сначала проиндексируй: rag.py index <папка>")
        return []
    chunks = json.load(open(INDEX, encoding="utf-8"))
    qwords = set(re.findall(r"\w+", query.lower()))
    scored = []
    for path, line, text in chunks:
        twords = set(re.findall(r"\w+", text.lower()))
        s = len(qwords & twords)
        if s:
            scored.append((s, path, line, text))
    scored.sort(reverse=True)
    return scored[:top]


def ask(query, backend):
    hits = search(query)
    if not hits:
        print("ничего релевантного не найдено")
        return
    context = "\n\n".join(f"--- {p}:{l} ---\n{t[:800]}" for _, p, l, t in hits)
    msgs = [{"role": "system",
             "content": "Ты анализируешь кодовую базу. Отвечай по предоставленному "
                        "контексту, указывай файлы и строки."},
            {"role": "user",
             "content": f"Контекст из кодовой базы:\n{context}\n\nВопрос: {query}"}]
    print(generate(backend, msgs, 1024))


def main():
    args = sys.argv[1:]
    if len(args) < 2:
        print("usage: rag.py index <папка> | rag.py ask <вопрос> [backend]")
        sys.exit(2)
    if args[0] == "index":
        index(args[1])
    elif args[0] == "ask":
        ask(" ".join(args[1:]), pick())


if __name__ == "__main__":
    main()