from __future__ import annotations

from pathlib import Path
from typing import List, Tuple

from genesis.rag.embeddings import Embedder
from genesis.rag.store import RagStore

_CHUNK = 800
_TEXT_EXTS = {
    ".java", ".py", ".js", ".ts", ".go", ".kt", ".scala", ".xml", ".yml",
    ".yaml", ".md", ".txt", ".json", ".csv", ".sql", ".html", ".css", ".sh",
}


def _split(text: str, size: int = _CHUNK) -> List[str]:
    paras = [p.strip() for p in text.split("\n") if p.strip()]
    out, buf = [], ""
    for p in paras:
        if len(buf) + len(p) > size and buf:
            out.append(buf)
            buf = p
        else:
            buf = (buf + "\n" + p).strip()
    if buf:
        out.append(buf)
    return out or [""]


class ImpactRAG:
    def __init__(self, store: RagStore, embedder: Embedder):
        self.store = store
        self.embedder = embedder

    def index(self, scope: str, sources: List[Tuple[str, str]]) -> None:
        chunks = []
        for name, text in sources:
            for piece in _split(text):
                # 跳过空片段，避免产生空向量噪声 chunk
                if piece.strip():
                    chunks.append(f"[{name}]\n{piece}")
        embs = self.embedder.embed(chunks)
        self.store.add(scope, chunks, embs)

    def index_dir(self, scope: str, root: str) -> int:
        """D1：索引目录内文本源文件（既有系统源码）。返回索引的文件数。"""
        root_path = Path(root)
        sources: List[Tuple[str, str]] = []
        if root_path.is_dir():
            for p in root_path.rglob("*"):
                if p.is_file() and p.suffix.lower() in _TEXT_EXTS:
                    try:
                        text = p.read_text(encoding="utf-8", errors="ignore")
                    except Exception:
                        # 读取异常的文件跳过，不中断整体索引
                        continue
                    # 空内容文件跳过
                    if not text.strip():
                        continue
                    # 含 NUL 字节的疑似二进制文件（误带白名单扩展名）跳过
                    if "\x00" in text:
                        continue
                    sources.append((str(p.relative_to(root_path)), text))
        if sources:
            self.index(scope, sources)
        return len(sources)

    def retrieve(self, scope: str, query: str, k: int = 5) -> List[str]:
        if not query.strip():
            return []
        qv = self.embedder.embed([query])[0]
        return self.store.search(scope, qv, k)