utils.py

import os
import binascii
import re
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor, as_completed, wait, FIRST_COMPLETED

# for now hardcoded
TEXT_EXTENSIONS = [".txt", ".md"]
EXCLUDE = ["node_modules"]


FUNCTION_PATTERNS = {
    # Python: def func_name(params):
    "py": r"^\s*def\s+([a-zA-Z_]\w*)\s*\(([^)]*)\)",
    # JavaScript: function funcName(params) { OR const funcName = (params) =>
    "js": r"^\s*(?:function\s+([a-zA-Z_]\w*)\s*\(([^)]*)\)|const\s+([a-zA-Z_]\w*)\s*=\s*\(([^)]*)\))",
    # Java: modifiers returnType funcName(params)
    "java": r"^\s*(?:public|private|protected)?\s*(?:static\s+)?[a-zA-Z_<>\[\]]+\s+([a-zA-Z_]\w*)\s*\(([^)]*)\)",
    # C: returnType funcName(params)
    "c": r"^\s*(?:[a-zA-Z_]\w*[\s\*]+)+([a-zA-Z_]\w*)\s*\(([^)]*)\)",
    # C++: template, modifiers, ClassName::funcName(params)
    "cpp": r"^\s*(?:template\s*<[^>]+>\s*)?(?:inline\s+|virtual\s+|static\s+)?[a-zA-Z_]\w*(?:\s*[*&:<>]\s*|\s+)+([a-zA-Z_]\w*)\s*\(([^)]*)\)",
    # C#: modifiers returnType funcName(params)
    "cs": r"^\s*(?:public|private|protected)?\s*(?:static\s+)?[a-zA-Z_<>\[\]]+\s+([a-zA-Z_]\w*)\s*\(([^)]*)\)",
    # Ruby: def func_name(params)
    "rb": r"^\s*def\s+([a-zA-Z_]\w*)\s*\(?([^)]*)\)?",
    # TypeScript: function funcName(params) OR const funcName = (params) =>
    "ts": r"^\s*(?:function\s+([a-zA-Z_]\w*)\s*\(([^)]*)\)|const\s+([a-zA-Z_]\w*)\s*=\s*\(([^)]*)\))",
}

COMPILED_PATTERNS = {ext: re.compile(pat) for ext, pat in FUNCTION_PATTERNS.items()}
b64d = lambda x: binascii.a2b_base64(x)
b64e = lambda x: binascii.b2a_base64(x, newline=False).decode()


def is_text_file(file_path: str):
    return file_path.split(".")[-1] in TEXT_EXTENSIONS


def extract_functions_from_file(file_path: Path, regex: re.Pattern, ext: str) -> str:
    """Extract functions using a precompiled regex."""
    functions = []
    try:
        with file_path.open("r", encoding="utf-8", errors="ignore") as f:
            for line in f:
                match = regex.match(line)
                if match:
                    groups = match.groups()
                    if ext in ("js", "ts"):
                        if groups[0]:
                            functions.append(f"{groups[0]}({groups[1]})")
                        elif groups[2]:
                            functions.append(f"{groups[2]}({groups[3]})")
                    else:
                        functions.append(f"{groups[0]}({groups[1]})")
    except (OSError, UnicodeDecodeError):
        return ""
    return "\n".join(functions)


def read_n_bytes_from_file(file_path: Path, n: int) -> bytes:
    try:
        with file_path.open("rb") as f:
            return f"{str(file_path)}: {f.read(n).decode('utf-8', errors='ignore')}"
    except OSError:
        return f"{str(file_path)}: ''"


def process_file(file_path: Path) -> tuple[str, str]:
    """Return (file_path, text)."""
    ext = file_path.suffix.lstrip(".").lower()
    regex = COMPILED_PATTERNS.get(ext)
    if regex:
        return str(file_path), extract_functions_from_file(file_path, regex, ext)
    else:
        return read_n_bytes_from_file(file_path, 200)


def scan_dir(path: Path):
    """List all files in a directory."""
    try:
        with os.scandir(path) as it:
            return [Path(entry.path) for entry in it if entry.is_file()]
    except OSError:
        return []


def read_files_from_dir(dir_path: str, max_workers=8):
    """Yield (file_path, text) from all files under dir_path as soon as they are found."""
    base_path = Path(dir_path)
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = {executor.submit(scan_dirs, base_path): base_path}
        file_futures = set()

        while futures or file_futures:
            # Wait for any future to complete (either dir scan or file processing)
            done, _ = wait(futures.keys() | file_futures, return_when=FIRST_COMPLETED)

            for future in done:
                if future in futures:
                    # Directory scanning completed
                    current_dir = futures.pop(future)
                    try:
                        subdirs = future.result()
                    except OSError:
                        continue

                    # Submit new directory scan futures for subdirectories
                    for subdir in subdirs:
                        futures[executor.submit(scan_dirs, subdir)] = subdir

                    # Also scan and submit file processing tasks for files in current_dir
                    try:
                        files = scan_dir(current_dir)
                        for file_path in files:
                            file_futures.add(executor.submit(process_file, file_path))
                    except OSError:
                        pass

                elif future in file_futures:
                    # File processing completed
                    file_futures.remove(future)
                    yield future.result()


def scan_dirs(path: Path):
    """List all directories inside a given directory."""
    try:
        with os.scandir(path) as it:
            return [
                Path(entry.path) for entry in it if entry.is_dir(follow_symlinks=False)
            ]
    except OSError:
        return []


def read_dirs_from_dir(dir_path: str, max_workers=8):
    """Yield directories under dir_path as soon as they are found."""
    base_path = Path(dir_path)
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = {executor.submit(scan_dirs, base_path): base_path}

        while futures:
            for future in as_completed(futures.keys()):
                futures.pop(future)
                try:
                    subdirs = future.result()
                except OSError:
                    # future does not exist
                    continue

                for subdir in subdirs:
                    yield str(subdir)
                    futures[executor.submit(scan_dirs, subdir)] = subdir


def read_dirs_from_dir_no_threads(dir_path: str):
    """Yield directory paths sequentially."""
    base_path = Path(dir_path)

    for p in base_path.rglob("*"):
        if p.is_dir():
            yield p


def get_doc_content(file: str, transformer, cu=None) -> tuple:
    """
    Return (text, shasum, vectors). If file is binary, return None.

    returns file content if text file create summary AI summary
    """
    from shorter import generate_summary
    import numpy as np
    import hashlib

    with open(file, "r") as f:
        print(f"Reading '{file}'")
        try:
            if os.path.getsize(file) > 1024 * 1024 * 10:
                text = f.read(5000)
            else:
                text = f.read()
            shasum = hashlib.sha3_256(text.encode("utf-8")).hexdigest()
        except:
            text = ""
            shasum = ""
            # save just name
            print(f"Binary file: '{file}'")

        if cu is not None:
            if (
                cu.execute(
                    "SELECT id FROM documents WHERE hash = ?", (shasum,)
                ).fetchone()
                is not None
            ):
                return None

        # file too big = take first 5000 bytes
        # file is text = take summary
        # file is programming language = extract functions
        if len(text) > 0:
            if extract_functions_from_file(file):
                funcs = extract_functions_from_file(file)
                if len(funcs) > 0:
                    text = b64e(
                        "\n".join(f"{name}({params})" for name, params in funcs).encode(
                            "utf-8"
                        )
                    )
                else:
                    text = b64e(text.encode("utf-8"))
            else:
                if os.path.getsize(file) > 1024 * 1024 * 10:
                    text = b64e(text[:5000].encode("utf-8"))
                else:
                    if is_text_file(file):
                        text = b64e(generate_summary(text).encode("utf-8"))
                    else:
                        text = b64e(text.encode("utf-8"))

        vectors = transformer.encode(f"{file}: {text}", normalize_embeddings=True)
        vectors = np.array(vectors, dtype=np.float32)

    return text, shasum, vectors


# from sentence_transformers import SentenceTransformer
# encoder = SentenceTransformer("all-MiniLM-L6-v2")
# prepare_database(dir_path, encoder)