#!/usr/bin/env python3
"""Build a small local vector database from a repository of files.

Architecture:
  - Ollama creates embeddings.
  - SQLite stores source files, chunks, metadata, and recoverable vectors.
  - FAISS IndexIDMap2(IndexFlatIP) accelerates cosine-similarity search.

Tested with Python 3.10+. Run `python repo_vector_db.py --help` for usage.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import os
import sqlite3
import sys
import time
import urllib.error
import urllib.request
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable, Sequence

import faiss
import numpy as np

DEFAULT_EXTENSIONS = {
    ".txt", ".md", ".markdown", ".rst", ".py", ".js", ".jsx", ".ts", ".tsx",
    ".java", ".c", ".h", ".cpp", ".hpp", ".cs", ".go", ".rs", ".rb", ".php",
    ".swift", ".kt", ".kts", ".sh", ".bash", ".zsh", ".ps1", ".sql", ".html",
    ".htm", ".css", ".scss", ".json", ".jsonl", ".yaml", ".yml", ".toml",
    ".ini", ".cfg", ".conf", ".csv", ".tsv", ".xml", ".pdf", ".docx", ".pptx",
}
IGNORE_DIRS = {
    ".git", ".hg", ".svn", ".idea", ".vscode", "node_modules", ".venv", "venv",
    "__pycache__", ".pytest_cache", ".mypy_cache", ".ruff_cache", "dist", "build",
    "target", ".next", ".cache",
}
SCHEMA_VERSION = "1"


@dataclass(frozen=True)
class FileRecord:
    relative_path: str
    absolute_path: Path
    sha256: str
    size: int
    mtime_ns: int


class OllamaEmbedder:
    def __init__(self, model: str = "nomic-embed-text", endpoint: str = "http://localhost:11434"):
        self.model = model
        self.endpoint = endpoint.rstrip("/")

    def embed(self, texts: Sequence[str]) -> np.ndarray:
        if not texts:
            return np.empty((0, 0), dtype=np.float32)
        payload = json.dumps({"model": self.model, "input": list(texts)}).encode("utf-8")
        request = urllib.request.Request(
            f"{self.endpoint}/api/embed",
            data=payload,
            headers={"Content-Type": "application/json"},
            method="POST",
        )
        try:
            with urllib.request.urlopen(request, timeout=300) as response:
                body = json.loads(response.read().decode("utf-8"))
        except (urllib.error.URLError, TimeoutError) as exc:
            raise RuntimeError(
                f"Could not reach Ollama at {self.endpoint}. Start it and pull {self.model!r}."
            ) from exc
        vectors = np.asarray(body.get("embeddings", []), dtype=np.float32)
        if vectors.ndim != 2 or vectors.shape[0] != len(texts):
            raise RuntimeError(f"Unexpected Ollama response shape: {vectors.shape}; response keys={list(body)}")
        return vectors


def connect(db_path: Path) -> sqlite3.Connection:
    db_path.parent.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(db_path)
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("PRAGMA foreign_keys=ON")
    conn.executescript(
        """
        CREATE TABLE IF NOT EXISTS settings (
            key TEXT PRIMARY KEY,
            value TEXT NOT NULL
        );
        CREATE TABLE IF NOT EXISTS files (
            path TEXT PRIMARY KEY,
            sha256 TEXT NOT NULL,
            size INTEGER NOT NULL,
            mtime_ns INTEGER NOT NULL,
            indexed_at REAL NOT NULL
        );
        CREATE TABLE IF NOT EXISTS chunks (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            path TEXT NOT NULL REFERENCES files(path) ON DELETE CASCADE,
            chunk_index INTEGER NOT NULL,
            start_line INTEGER NOT NULL,
            end_line INTEGER NOT NULL,
            text TEXT NOT NULL,
            vector BLOB NOT NULL,
            vector_dim INTEGER NOT NULL,
            UNIQUE(path, chunk_index)
        );
        CREATE INDEX IF NOT EXISTS chunks_path_idx ON chunks(path);
        """
    )
    conn.execute(
        "INSERT OR IGNORE INTO settings(key, value) VALUES('schema_version', ?)",
        (SCHEMA_VERSION,),
    )
    conn.commit()
    return conn


def sha256_file(path: Path, block_size: int = 1024 * 1024) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        while block := handle.read(block_size):
            digest.update(block)
    return digest.hexdigest()


def discover_files(root: Path, max_bytes: int, extensions: set[str]) -> list[FileRecord]:
    root = root.resolve()
    found: list[FileRecord] = []
    for path in sorted(root.rglob("*")):
        if not path.is_file() or path.is_symlink():
            continue
        relative_parts = path.relative_to(root).parts
        if any(part in IGNORE_DIRS for part in relative_parts[:-1]):
            continue
        if path.suffix.lower() not in extensions:
            continue
        stat = path.stat()
        if stat.st_size > max_bytes:
            print(f"SKIP too large ({stat.st_size:,} bytes): {path.relative_to(root)}", file=sys.stderr)
            continue
        found.append(
            FileRecord(
                relative_path=path.relative_to(root).as_posix(),
                absolute_path=path,
                sha256=sha256_file(path),
                size=stat.st_size,
                mtime_ns=stat.st_mtime_ns,
            )
        )
    return found


def extract_text(path: Path) -> str:
    suffix = path.suffix.lower()
    if suffix == ".pdf":
        try:
            from pypdf import PdfReader
        except ImportError as exc:
            raise RuntimeError("PDF support requires: pip install pypdf") from exc
        return "\n\n".join((page.extract_text() or "") for page in PdfReader(str(path)).pages)
    if suffix == ".docx":
        try:
            from docx import Document
        except ImportError as exc:
            raise RuntimeError("DOCX support requires: pip install python-docx") from exc
        return "\n".join(paragraph.text for paragraph in Document(str(path)).paragraphs)
    if suffix == ".pptx":
        try:
            from pptx import Presentation
        except ImportError as exc:
            raise RuntimeError("PPTX support requires: pip install python-pptx") from exc
        lines: list[str] = []
        for slide_number, slide in enumerate(Presentation(str(path)).slides, start=1):
            lines.append(f"[Slide {slide_number}]")
            for shape in slide.shapes:
                if hasattr(shape, "text") and shape.text.strip():
                    lines.append(shape.text.strip())
        return "\n".join(lines)
    raw = path.read_bytes()
    if b"\x00" in raw[:8192]:
        raise RuntimeError("file appears to be binary")
    return raw.decode("utf-8", errors="replace")


def chunk_text(text: str, chunk_chars: int = 1800, overlap_chars: int = 250) -> list[tuple[int, int, str]]:
    if chunk_chars < 200:
        raise ValueError("chunk_chars must be at least 200")
    if not 0 <= overlap_chars < chunk_chars:
        raise ValueError("overlap_chars must be >= 0 and smaller than chunk_chars")
    text = text.replace("\r\n", "\n").replace("\r", "\n").strip()
    if not text:
        return []
    chunks: list[tuple[int, int, str]] = []
    start = 0
    while start < len(text):
        target = min(len(text), start + chunk_chars)
        end = target
        if target < len(text):
            candidates = [text.rfind("\n\n", start + chunk_chars // 2, target), text.rfind("\n", start + chunk_chars // 2, target)]
            boundary = max(candidates)
            if boundary > start:
                end = boundary
        chunk = text[start:end].strip()
        if chunk:
            start_line = text.count("\n", 0, start) + 1
            end_line = start_line + chunk.count("\n")
            chunks.append((start_line, end_line, chunk))
        if end >= len(text):
            break
        next_start = max(0, end - overlap_chars)
        newline = text.find("\n", next_start, end)
        start = newline + 1 if newline != -1 else next_start
        if start <= 0 or start >= end:
            start = end
    return chunks


def batched(items: Sequence[str], size: int) -> Iterable[tuple[int, Sequence[str]]]:
    for start in range(0, len(items), size):
        yield start, items[start : start + size]


def set_setting(conn: sqlite3.Connection, key: str, value: str) -> None:
    conn.execute(
        "INSERT INTO settings(key, value) VALUES(?, ?) ON CONFLICT(key) DO UPDATE SET value=excluded.value",
        (key, value),
    )


def get_setting(conn: sqlite3.Connection, key: str) -> str | None:
    row = conn.execute("SELECT value FROM settings WHERE key=?", (key,)).fetchone()
    return row[0] if row else None


def rebuild_index(conn: sqlite3.Connection, index_path: Path) -> tuple[int, int]:
    rows = conn.execute("SELECT id, vector, vector_dim FROM chunks ORDER BY id").fetchall()
    dim = int(rows[0]["vector_dim"]) if rows else int(get_setting(conn, "vector_dim") or 768)
    base = faiss.IndexFlatIP(dim)
    index = faiss.IndexIDMap2(base)
    if rows:
        vectors = np.stack([np.frombuffer(row["vector"], dtype=np.float32) for row in rows]).astype(np.float32)
        if vectors.shape[1] != dim:
            raise RuntimeError(f"Vector dimension mismatch while rebuilding: {vectors.shape[1]} != {dim}")
        faiss.normalize_L2(vectors)
        ids = np.asarray([row["id"] for row in rows], dtype=np.int64)
        index.add_with_ids(vectors, ids)
    index_path.parent.mkdir(parents=True, exist_ok=True)
    temporary = index_path.with_suffix(index_path.suffix + ".tmp")
    faiss.write_index(index, str(temporary))
    os.replace(temporary, index_path)
    return len(rows), dim


def sync_repository(
    root: Path,
    db_path: Path,
    index_path: Path,
    embedder: OllamaEmbedder,
    chunk_chars: int,
    overlap_chars: int,
    batch_size: int,
    max_bytes: int,
    extensions: set[str],
    rebuild_all: bool = False,
) -> dict:
    root = root.resolve()
    conn = connect(db_path)
    existing_model = get_setting(conn, "embedding_model")
    chunk_count = conn.execute("SELECT COUNT(*) FROM chunks").fetchone()[0]
    if existing_model and existing_model != embedder.model and chunk_count and not rebuild_all:
        raise RuntimeError(
            f"Database uses embedding model {existing_model!r}, not {embedder.model!r}. "
            "Use --rebuild-all to re-embed everything."
        )
    if rebuild_all:
        conn.execute("DELETE FROM files")
        conn.execute("DELETE FROM chunks")
        conn.commit()

    discovered = discover_files(root, max_bytes=max_bytes, extensions=extensions)
    current = {record.relative_path: record for record in discovered}
    stored = {row["path"]: row for row in conn.execute("SELECT * FROM files")}
    deleted = sorted(set(stored) - set(current))
    changed = [
        record for path, record in current.items()
        if path not in stored or stored[path]["sha256"] != record.sha256
    ]
    unchanged = len(current) - len(changed)

    documents: list[tuple[FileRecord, int, int, int, str]] = []
    skipped: list[tuple[str, str]] = []
    for record in changed:
        try:
            text = extract_text(record.absolute_path)
            pieces = chunk_text(text, chunk_chars=chunk_chars, overlap_chars=overlap_chars)
            for chunk_index, (start_line, end_line, chunk) in enumerate(pieces):
                documents.append((record, chunk_index, start_line, end_line, chunk))
        except Exception as exc:
            skipped.append((record.relative_path, str(exc)))

    texts = [item[4] for item in documents]
    matrices: list[np.ndarray] = []
    for start, batch in batched(texts, batch_size):
        print(f"Embedding chunks {start + 1}-{start + len(batch)} of {len(texts)}...")
        matrices.append(embedder.embed(batch))
    vectors = np.vstack(matrices).astype(np.float32) if matrices else np.empty((0, 0), dtype=np.float32)
    if len(vectors) != len(documents):
        raise RuntimeError("Embedding count does not match chunk count")

    now = time.time()
    conn.execute("BEGIN")
    try:
        for path in deleted:
            conn.execute("DELETE FROM files WHERE path=?", (path,))
        for record in changed:
            conn.execute("DELETE FROM files WHERE path=?", (record.relative_path,))
            conn.execute(
                "INSERT INTO files(path, sha256, size, mtime_ns, indexed_at) VALUES(?,?,?,?,?)",
                (record.relative_path, record.sha256, record.size, record.mtime_ns, now),
            )
        for item, vector in zip(documents, vectors):
            record, chunk_index, start_line, end_line, chunk = item
            conn.execute(
                "INSERT INTO chunks(path, chunk_index, start_line, end_line, text, vector, vector_dim) VALUES(?,?,?,?,?,?,?)",
                (
                    record.relative_path,
                    chunk_index,
                    start_line,
                    end_line,
                    chunk,
                    vector.tobytes(),
                    int(vector.shape[0]),
                ),
            )
        set_setting(conn, "embedding_model", embedder.model)
        set_setting(conn, "repository_root", str(root))
        if len(vectors):
            set_setting(conn, "vector_dim", str(vectors.shape[1]))
        conn.commit()
    except Exception:
        conn.rollback()
        raise

    indexed_chunks, dim = rebuild_index(conn, index_path)
    result = {
        "root": str(root),
        "files_seen": len(current),
        "files_changed": len(changed),
        "files_deleted": len(deleted),
        "files_unchanged": unchanged,
        "chunks_embedded": len(documents),
        "chunks_total": indexed_chunks,
        "vector_dim": dim,
        "skipped": skipped,
        "db_path": str(db_path.resolve()),
        "index_path": str(index_path.resolve()),
    }
    conn.close()
    return result


def search_index(
    query: str,
    db_path: Path,
    index_path: Path,
    embedder: OllamaEmbedder,
    top_k: int,
    path_prefix: str | None = None,
) -> list[dict]:
    conn = connect(db_path)
    stored_model = get_setting(conn, "embedding_model")
    if stored_model and stored_model != embedder.model:
        raise RuntimeError(f"Query model {embedder.model!r} differs from indexed model {stored_model!r}")
    if not index_path.exists():
        raise RuntimeError(f"FAISS index does not exist: {index_path}. Run sync first.")
    index = faiss.read_index(str(index_path))
    if index.ntotal == 0:
        return []
    vector = embedder.embed([query]).astype(np.float32)
    if vector.shape[1] != index.d:
        raise RuntimeError(f"Query dimension {vector.shape[1]} does not match index dimension {index.d}")
    faiss.normalize_L2(vector)
    candidate_k = min(index.ntotal, max(top_k, top_k * 10 if path_prefix else top_k))
    scores, ids = index.search(vector, candidate_k)
    ordered = [(int(doc_id), float(score)) for doc_id, score in zip(ids[0], scores[0]) if doc_id != -1]
    rows_by_id = {}
    if ordered:
        placeholders = ",".join("?" for _ in ordered)
        for row in conn.execute(
            f"SELECT id, path, chunk_index, start_line, end_line, text FROM chunks WHERE id IN ({placeholders})",
            [doc_id for doc_id, _ in ordered],
        ):
            rows_by_id[int(row["id"])] = row
    results = []
    for doc_id, score in ordered:
        row = rows_by_id.get(doc_id)
        if not row:
            continue
        if path_prefix and not row["path"].startswith(path_prefix):
            continue
        results.append(
            {
                "id": doc_id,
                "score": score,
                "path": row["path"],
                "chunk_index": row["chunk_index"],
                "start_line": row["start_line"],
                "end_line": row["end_line"],
                "text": row["text"],
            }
        )
        if len(results) >= top_k:
            break
    conn.close()
    return results


def database_status(db_path: Path, index_path: Path) -> dict:
    conn = connect(db_path)
    status = {
        "database": str(db_path.resolve()),
        "index": str(index_path.resolve()),
        "repository_root": get_setting(conn, "repository_root"),
        "embedding_model": get_setting(conn, "embedding_model"),
        "vector_dim": get_setting(conn, "vector_dim"),
        "files": conn.execute("SELECT COUNT(*) FROM files").fetchone()[0],
        "chunks": conn.execute("SELECT COUNT(*) FROM chunks").fetchone()[0],
        "index_exists": index_path.exists(),
    }
    if index_path.exists():
        status["index_vectors"] = int(faiss.read_index(str(index_path)).ntotal)
    conn.close()
    return status


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("--db", type=Path, default=Path("vector_store/knowledge.sqlite3"))
    parser.add_argument("--index", type=Path, default=Path("vector_store/vectors.faiss"))
    parser.add_argument("--model", default="nomic-embed-text")
    parser.add_argument("--ollama", default="http://localhost:11434")
    sub = parser.add_subparsers(dest="command", required=True)

    sync = sub.add_parser("sync", help="Index new/changed files and remove deleted files")
    sync.add_argument("root", type=Path, help="Repository or folder to index")
    sync.add_argument("--chunk-chars", type=int, default=1800)
    sync.add_argument("--overlap-chars", type=int, default=250)
    sync.add_argument("--batch-size", type=int, default=32)
    sync.add_argument("--max-file-mb", type=float, default=10.0)
    sync.add_argument("--extensions", help="Comma-separated extension list; defaults to common text/document types")
    sync.add_argument("--rebuild-all", action="store_true", help="Delete existing records and re-embed every file")

    search = sub.add_parser("search", help="Search the indexed repository")
    search.add_argument("query")
    search.add_argument("--top-k", type=int, default=5)
    search.add_argument("--path-prefix", help="Only return paths beginning with this repository-relative prefix")
    search.add_argument("--full-text", action="store_true")

    sub.add_parser("status", help="Show database/index health")
    sub.add_parser("rebuild-index", help="Rebuild FAISS from vectors stored in SQLite")
    return parser


def main() -> int:
    args = build_parser().parse_args()
    embedder = OllamaEmbedder(model=args.model, endpoint=args.ollama)
    if args.command == "sync":
        extensions = DEFAULT_EXTENSIONS
        if args.extensions:
            extensions = {item.strip().lower() for item in args.extensions.split(",") if item.strip()}
            extensions = {item if item.startswith(".") else f".{item}" for item in extensions}
        result = sync_repository(
            root=args.root,
            db_path=args.db,
            index_path=args.index,
            embedder=embedder,
            chunk_chars=args.chunk_chars,
            overlap_chars=args.overlap_chars,
            batch_size=args.batch_size,
            max_bytes=int(args.max_file_mb * 1024 * 1024),
            extensions=extensions,
            rebuild_all=args.rebuild_all,
        )
        print(json.dumps(result, indent=2))
        return 0
    if args.command == "search":
        results = search_index(
            args.query, args.db, args.index, embedder, args.top_k, args.path_prefix
        )
        for rank, result in enumerate(results, start=1):
            print(f"\n#{rank} score={result['score']:.4f} {result['path']}:{result['start_line']}-{result['end_line']}")
            text = result["text"] if args.full_text else result["text"][:500]
            print(text)
        return 0
    if args.command == "status":
        print(json.dumps(database_status(args.db, args.index), indent=2))
        return 0
    if args.command == "rebuild-index":
        conn = connect(args.db)
        count, dim = rebuild_index(conn, args.index)
        conn.close()
        print(json.dumps({"vectors": count, "dimension": dim, "index": str(args.index.resolve())}, indent=2))
        return 0
    return 2


if __name__ == "__main__":
    raise SystemExit(main())