engine.py

import os
import json
import hashlib
import subprocess
import tempfile
import numpy as np
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor, as_completed
from google import genai
from google.genai import types
from dotenv import load_dotenv
import threading

MODEL = "gemini-embedding-2-preview"
DIMENSIONS = 768
TEXT_BATCH_SIZE = 20  # texts per API call
MAX_WORKERS = 10  # concurrent API calls

TEXT_EXTENSIONS = {
    ".txt", ".md", ".py", ".js", ".ts", ".html", ".css", ".json", ".yaml",
    ".yml", ".csv", ".xml", ".sh", ".rs", ".go", ".java", ".c", ".cpp",
    ".h", ".rb", ".php", ".swift", ".kt", ".sql", ".r", ".m", ".toml",
}
IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp"}
VIDEO_EXTENSIONS = {".mp4", ".mov"}
AUDIO_EXTENSIONS = {".mp3", ".wav"}
PDF_EXTENSIONS = {".pdf"}

MIME_MAP = {
    ".png": "image/png",
    ".jpg": "image/jpeg",
    ".jpeg": "image/jpeg",
    ".webp": "image/webp",
    ".mp4": "video/mp4",
    ".mov": "video/quicktime",
    ".mp3": "audio/mpeg",
    ".wav": "audio/wav",
    ".pdf": "application/pdf",
}


def _get_file_category(path: Path) -> str | None:
    ext = path.suffix.lower()
    if ext in TEXT_EXTENSIONS:
        return "text"
    if ext in IMAGE_EXTENSIONS:
        return "image"
    if ext in VIDEO_EXTENSIONS:
        return "video"
    if ext in AUDIO_EXTENSIONS:
        return "audio"
    if ext in PDF_EXTENSIONS:
        return "pdf"
    return None


def _normalize(vec: np.ndarray) -> np.ndarray:
    return vec / np.linalg.norm(vec)


def _file_hash(path: Path) -> str:
    """Fast change detection using size + mtime (just a stat call, no I/O)."""
    stat = path.stat()
    return f"{stat.st_size}:{stat.st_mtime_ns}"


class OmnisearchEngine:
    def __init__(self):
        load_dotenv()
        api_key = os.getenv("GEMINI_API_KEY")
        if not api_key:
            raise ValueError("GEMINI_API_KEY not found in environment")
        self.client = genai.Client(api_key=api_key)
        self.embeddings: list[np.ndarray] = []
        self.metadata: list[dict] = []
        self._file_hashes: dict[str, str] = {}  # file_path -> content hash
        self._lock = threading.Lock()

    # ── Embedding helpers ──────────────────────────────────────────

    def _embed_text_batch(self, texts: list[str], task_type: str = "RETRIEVAL_DOCUMENT") -> list[np.ndarray]:
        """Embed multiple texts in a single API call."""
        result = self.client.models.embed_content(
            model=MODEL,
            contents=texts,
            config=types.EmbedContentConfig(
                task_type=task_type,
                output_dimensionality=DIMENSIONS,
            ),
        )
        return [
            _normalize(np.array(e.values, dtype=np.float32))
            for e in result.embeddings
        ]

    def _embed_text(self, text: str, task_type: str = "RETRIEVAL_DOCUMENT") -> np.ndarray:
        return self._embed_text_batch([text], task_type)[0]

    def _embed_bytes(self, data: bytes, mime_type: str) -> np.ndarray:
        result = self.client.models.embed_content(
            model=MODEL,
            contents=[types.Part.from_bytes(data=data, mime_type=mime_type)],
            config=types.EmbedContentConfig(
                task_type="RETRIEVAL_DOCUMENT",
                output_dimensionality=DIMENSIONS,
            ),
        )
        vec = np.array(result.embeddings[0].values, dtype=np.float32)
        return _normalize(vec)

    # ── Chunking ──────────────────────────────────────────────────

    def _chunk_text(self, text: str, chunk_size: int = 2000, overlap: int = 200) -> list[str]:
        chunks = []
        for i in range(0, len(text), chunk_size - overlap):
            chunk = text[i : i + chunk_size]
            if chunk.strip():
                chunks.append(chunk)
        return chunks

    # ── Per-type indexers (return metadata + embeddings) ──────────

    def _prepare_text_file(self, path: Path) -> tuple[list[str], list[dict]]:
        """Return (texts_to_embed, metadata_entries) without calling the API."""
        text = path.read_text(errors="ignore")
        if not text.strip():
            return [], []
        chunks = self._chunk_text(text)
        if not chunks:
            return [], []
        metas = [
            {
                "file_path": str(path),
                "file_name": path.name,
                "chunk_type": "text",
                "chunk_index": i,
                "total_chunks": len(chunks),
                "preview": c[:300],
            }
            for i, c in enumerate(chunks)
        ]
        return chunks, metas

    def _index_binary_file(self, path: Path, category: str) -> list[tuple[np.ndarray, dict]]:
        """Embed a binary file, chunking if needed. Returns list of (embedding, metadata)."""
        if category == "pdf":
            return self._index_pdf_chunked(path)
        if category == "video":
            return self._index_video_chunked(path)
        if category == "audio":
            return self._index_audio_chunked(path)
        # Images: no chunking needed
        data = path.read_bytes()
        mime = MIME_MAP[path.suffix.lower()]
        vec = self._embed_bytes(data, mime)
        meta = {
            "file_path": str(path),
            "file_name": path.name,
            "chunk_type": category,
            "chunk_index": 0,
            "total_chunks": 1,
            "preview": path.name,
        }
        return [(vec, meta)]

    def _index_pdf_chunked(self, path: Path) -> list[tuple[np.ndarray, dict]]:
        """Split PDFs into 6-page chunks (API limit)."""
        import fitz  # pymupdf

        doc = fitz.open(str(path))
        total_pages = len(doc)
        chunk_size = 6
        results = []
        num_chunks = (total_pages + chunk_size - 1) // chunk_size

        if total_pages <= chunk_size:
            data = path.read_bytes()
            vec = self._embed_bytes(data, "application/pdf")
            results.append((vec, {
                "file_path": str(path),
                "file_name": path.name,
                "chunk_type": "pdf",
                "chunk_index": 0,
                "total_chunks": 1,
                "preview": f"{path.name} (pp. 1-{total_pages})",
            }))
        else:
            for chunk_i in range(num_chunks):
                start = chunk_i * chunk_size
                end = min(start + chunk_size, total_pages)
                chunk_doc = fitz.open()
                chunk_doc.insert_pdf(doc, from_page=start, to_page=end - 1)
                chunk_bytes = chunk_doc.tobytes()
                chunk_doc.close()

                vec = self._embed_bytes(chunk_bytes, "application/pdf")
                results.append((vec, {
                    "file_path": str(path),
                    "file_name": path.name,
                    "chunk_type": "pdf",
                    "chunk_index": chunk_i,
                    "total_chunks": num_chunks,
                    "preview": f"{path.name} (pp. {start+1}-{end})",
                }))

        doc.close()
        return results

    def _get_media_duration(self, path: Path) -> float:
        """Get duration of audio/video file in seconds using ffprobe."""
        try:
            result = subprocess.run(
                ["ffprobe", "-v", "quiet", "-print_format", "json",
                 "-show_format", str(path)],
                capture_output=True, text=True,
            )
            info = json.loads(result.stdout)
            return float(info["format"]["duration"])
        except Exception:
            return 0.0

    def _split_media(self, path: Path, max_duration: float, ext: str) -> list[Path]:
        """Split a media file into chunks of max_duration seconds. Returns temp file paths."""
        duration = self._get_media_duration(path)
        if duration <= max_duration:
            return [path]

        tmp_dir = tempfile.mkdtemp()
        chunks = []
        start = 0.0
        i = 0
        while start < duration:
            chunk_path = Path(tmp_dir) / f"chunk_{i:03d}{ext}"
            subprocess.run(
                ["ffmpeg", "-y", "-ss", str(start), "-i", str(path),
                 "-t", str(max_duration), "-c", "copy", str(chunk_path)],
                capture_output=True,
            )
            if chunk_path.exists() and chunk_path.stat().st_size > 0:
                chunks.append(chunk_path)
            start += max_duration
            i += 1
        return chunks

    def _index_video_chunked(self, path: Path) -> list[tuple[np.ndarray, dict]]:
        """Split videos into 120s chunks (API limit)."""
        chunk_paths = self._split_media(path, max_duration=115.0, ext=path.suffix)
        results = []
        for i, cp in enumerate(chunk_paths):
            data = cp.read_bytes()
            mime = MIME_MAP[path.suffix.lower()]
            vec = self._embed_bytes(data, mime)
            start_sec = int(i * 115)
            results.append((vec, {
                "file_path": str(path),
                "file_name": path.name,
                "chunk_type": "video",
                "chunk_index": i,
                "total_chunks": len(chunk_paths),
                "preview": f"{path.name} ({start_sec}s-{start_sec+115}s)",
            }))
            # Clean up temp files
            if cp != path:
                cp.unlink(missing_ok=True)
        return results

    def _index_audio_chunked(self, path: Path) -> list[tuple[np.ndarray, dict]]:
        """Split audio into 75s chunks (API limit is 80s, use 75s for safety)."""
        chunk_paths = self._split_media(path, max_duration=75.0, ext=path.suffix)
        results = []
        for i, cp in enumerate(chunk_paths):
            data = cp.read_bytes()
            mime = MIME_MAP[path.suffix.lower()]
            vec = self._embed_bytes(data, mime)
            start_sec = int(i * 75)
            results.append((vec, {
                "file_path": str(path),
                "file_name": path.name,
                "chunk_type": "audio",
                "chunk_index": i,
                "total_chunks": len(chunk_paths),
                "preview": f"{path.name} ({start_sec}s-{start_sec+75}s)",
            }))
            if cp != path:
                cp.unlink(missing_ok=True)
        return results

    # ── Main indexing ─────────────────────────────────────────────

    def discover_files(self, directory: str) -> list[tuple[Path, str]]:
        dir_path = Path(directory).resolve()
        files = []
        for path in sorted(dir_path.rglob("*")):
            if not path.is_file():
                continue
            if any(part.startswith(".") for part in path.relative_to(dir_path).parts):
                continue
            cat = _get_file_category(path)
            if cat:
                files.append((path, cat))
        return files

    def index_file(self, path: Path, category: str):
        if category == "text":
            chunks, metas = self._prepare_text_file(path)
            if chunks:
                vecs = self._embed_text_batch(chunks)
                with self._lock:
                    self.embeddings.extend(vecs)
                    self.metadata.extend(metas)
        else:
            results = self._index_binary_file(path, category)
            with self._lock:
                for vec, meta in results:
                    self.embeddings.append(vec)
                    self.metadata.append(meta)

    def _remove_file_entries(self, file_path: str):
        """Remove all chunks for a given file path from the index."""
        indices_to_keep = [
            i for i, m in enumerate(self.metadata) if m["file_path"] != file_path
        ]
        self.metadata = [self.metadata[i] for i in indices_to_keep]
        self.embeddings = [self.embeddings[i] for i in indices_to_keep]
        self._file_hashes.pop(file_path, None)

    def index_directory(self, directory: str, progress_callback=None):
        files = self.discover_files(directory)
        if not files:
            return

        # Determine which files need (re-)indexing
        current_paths = {str(p) for p, _ in files}
        new_files: list[tuple[Path, str]] = []
        skipped = 0

        for path, category in files:
            fp = str(path)
            h = _file_hash(path)
            if fp in self._file_hashes and self._file_hashes[fp] == h:
                skipped += 1
                continue
            # File is new or changed — remove old entries if any
            if fp in self._file_hashes:
                self._remove_file_entries(fp)
            self._file_hashes[fp] = h
            new_files.append((path, category))

        # Remove entries for deleted files
        stale = [fp for fp in self._file_hashes if fp not in current_paths]
        for fp in stale:
            self._remove_file_entries(fp)

        if not new_files:
            if progress_callback:
                progress_callback(0, 0, f"skipped {skipped} unchanged files")
            return

        # Separate text files (batchable) from binary files (concurrent)
        text_files = [(p, c) for p, c in new_files if c == "text"]
        binary_files = [(p, c) for p, c in new_files if c != "text"]

        done = 0
        total = len(new_files)

        # 1) Batch all text files: collect chunks, batch-embed
        all_text_chunks: list[str] = []
        all_text_metas: list[dict] = []
        for path, _ in text_files:
            try:
                chunks, metas = self._prepare_text_file(path)
                all_text_chunks.extend(chunks)
                all_text_metas.extend(metas)
            except Exception as e:
                print(f"Error reading {path}: {e}")

            done += 1
            if progress_callback:
                progress_callback(done, total, f"{path.name} (skipped {skipped})")

        # Embed text chunks in batches
        for i in range(0, len(all_text_chunks), TEXT_BATCH_SIZE):
            batch = all_text_chunks[i : i + TEXT_BATCH_SIZE]
            batch_metas = all_text_metas[i : i + TEXT_BATCH_SIZE]
            try:
                vecs = self._embed_text_batch(batch)
                self.embeddings.extend(vecs)
                self.metadata.extend(batch_metas)
            except Exception as e:
                print(f"Error embedding text batch: {e}")

        # 2) Embed binary files concurrently
        def _process_binary(path_cat):
            path, category = path_cat
            return self._index_binary_file(path, category), path.name

        with ThreadPoolExecutor(max_workers=MAX_WORKERS) as pool:
            futures = {pool.submit(_process_binary, pc): pc for pc in binary_files}
            for future in as_completed(futures):
                done += 1
                try:
                    results, name = future.result()
                    for vec, meta in results:
                        self.embeddings.append(vec)
                        self.metadata.append(meta)
                    if progress_callback:
                        progress_callback(done, total, f"{name} (skipped {skipped})")
                except Exception as e:
                    path, _ = futures[future]
                    print(f"Error indexing {path}: {e}")
                    if progress_callback:
                        progress_callback(done, total, path.name)

    # ── Search ────────────────────────────────────────────────────

    def search(self, query: str, top_k: int = 20) -> list[dict]:
        """Search with per-file dedup. Returns one result per file, ranked by best chunk."""
        if not self.embeddings:
            return []

        query_vec = self._embed_text(query, task_type="RETRIEVAL_QUERY")
        matrix = np.stack(self.embeddings)
        scores = matrix @ query_vec

        # Group by file, keep the best-scoring chunk per file
        best_per_file: dict[str, dict] = {}
        for idx in range(len(scores)):
            meta = self.metadata[idx]
            fp = meta["file_path"]
            score = float(scores[idx])
            if fp not in best_per_file or score > best_per_file[fp]["score"]:
                best_per_file[fp] = {**meta, "score": score}

        results = sorted(best_per_file.values(), key=lambda r: r["score"], reverse=True)
        return results[:top_k]

    def search_chunks(self, query: str, top_k: int = 20) -> list[dict]:
        """Raw chunk-level search without dedup (for drilling into a specific file)."""
        if not self.embeddings:
            return []

        query_vec = self._embed_text(query, task_type="RETRIEVAL_QUERY")
        matrix = np.stack(self.embeddings)
        scores = matrix @ query_vec
        top_indices = np.argsort(scores)[::-1][:top_k]

        results = []
        for idx in top_indices:
            results.append({**self.metadata[idx], "score": float(scores[idx])})
        return results

    # ── Persistence ───────────────────────────────────────────────

    def save(self, directory: str):
        os.makedirs(directory, exist_ok=True)
        np.save(os.path.join(directory, "embeddings.npy"), np.stack(self.embeddings))
        with open(os.path.join(directory, "metadata.json"), "w") as f:
            json.dump(self.metadata, f, indent=2)
        with open(os.path.join(directory, "file_hashes.json"), "w") as f:
            json.dump(self._file_hashes, f, indent=2)

    def load(self, directory: str):
        emb_path = os.path.join(directory, "embeddings.npy")
        meta_path = os.path.join(directory, "metadata.json")
        hash_path = os.path.join(directory, "file_hashes.json")
        if not os.path.exists(emb_path):
            return False
        self.embeddings = list(np.load(emb_path))
        with open(meta_path) as f:
            self.metadata = json.load(f)
        if os.path.exists(hash_path):
            with open(hash_path) as f:
                self._file_hashes = json.load(f)
        return True

    @property
    def indexed_count(self) -> int:
        return len(self.metadata)

    @property
    def file_type_counts(self) -> dict[str, int]:
        counts: dict[str, int] = {}
        seen_files: dict[str, str] = {}
        for m in self.metadata:
            fp = m["file_path"]
            if fp not in seen_files:
                seen_files[fp] = m["chunk_type"]
                counts[m["chunk_type"]] = counts.get(m["chunk_type"], 0) + 1
        return counts