#!/usr/bin/env python3
"""Local CLI semantic search backed by Qwen3 embedding and reranking."""

from __future__ import annotations

import argparse
import ast
import hashlib
import json
import multiprocessing
from multiprocessing import AuthenticationError
from multiprocessing.connection import Client, Listener
import os
import queue
from pathlib import Path
import re
import sqlite3
import stat
import subprocess
import sys
import threading
import time


MODEL = "Qwen/Qwen3-Embedding-0.6B"
MODEL_REVISION = "97b0c614be4d77ee51c0cef4e5f07c00f9eb65b3"
RERANKER_MODEL = "Qwen/Qwen3-Reranker-0.6B"
RERANKER_MODEL_REVISION = "e61197ed45024b0ed8a2d74b80b4d909f1255473"
EMBEDDING_DIMENSIONS = 1024
CHUNK_TOKENS = 400
OVERLAP_TOKENS = 64
DENSE_CANDIDATES = 50
LEXICAL_CANDIDATES = 50
RERANK_CANDIDATES = 50
FINAL_RESULTS = 5
CLI_TIMEOUT_SECONDS = 300
RRF_K = 60
INDEX_VERSION = 3
RETRIEVAL_INSTRUCTION = "Given a codebase question, retrieve relevant code and documentation that answer the question."
# Qwen3-Embedding query format from the model card: no space after "Query:".
QUERY_PROMPT = f"Instruct: {RETRIEVAL_INSTRUCTION}\nQuery:"
TEXT_EXTENSIONS = {".c", ".cjs", ".cpp", ".cs", ".csproj", ".css", ".go", ".gradle", ".html", ".java", ".js", ".json", ".jsx", ".kt", ".kts",
                   ".md", ".mjs", ".props", ".ps1", ".psm1", ".py", ".razor", ".rs", ".scss", ".sh", ".sql", ".targets", ".toml", ".ts",
                   ".tsx", ".txt", ".xaml", ".xml", ".yaml", ".yml"}
DEFAULT_MAX_FILES = 2000
MAX_FILES_ENV = "SEMANTIC_RETRIEVAL_MAX_FILES"

# A daemon keeps the embedding/reranker models resident in memory across CLI
# invocations, since each `server.py search` call is otherwise a fresh process
# that would reload both models from disk before it can answer a single query.
DAEMON_IDLE_SECONDS = 1800
DAEMON_READY_TIMEOUT_SECONDS = 120
DAEMON_MESSAGE_LIMIT = 16 * 1024 * 1024


def runtime_python(script=__file__, executable=sys.executable, system=os.name):
    relative = Path('.venv') / ('Scripts/python.exe' if system == 'nt' else 'bin/python')
    candidate = (Path(script).resolve().parent / relative)
    if candidate.is_file() and candidate != Path(executable).resolve():
        return candidate
    return None


def use_runtime_python(system=os.name):
    runtime = runtime_python()
    if runtime is not None:
        command = [str(runtime), str(Path(__file__).resolve()), *sys.argv[1:]]
        if system == 'nt':
            # os.execv does not quote arguments on Windows, so a multi-word --query would be split.
            sys.exit(subprocess.run(command).returncode)
        os.execv(command[0], command)


def sections(text: str, path: str):
    lines = text.splitlines(keepends=True)
    boundaries = {0: "<module>", len(lines): "<module>"}
    suffix = Path(path).suffix.lower()
    if suffix == ".py":
        try:
            tree = ast.parse(text)
        except (SyntaxError, ValueError):
            tree = None

        def visit(node, parent):
            symbol = parent
            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
                symbol = f"{parent}.{node.name}" if parent != "<module>" else node.name
                start = min([node.lineno, *(item.lineno for item in node.decorator_list)]) - 1
                boundaries[start] = symbol
                boundaries.setdefault(node.end_lineno, parent)
            for child in ast.iter_child_nodes(node):
                visit(child, symbol)

        if tree is not None:
            visit(tree, "<module>")
    else:
        fence = None
        headings = []
        for number, line in enumerate(lines):
            if suffix == ".md":
                marker = re.match(r"^\s{0,3}(`{3,}|~{3,})", line)
                if marker:
                    value = marker[1]
                    if fence is None:
                        fence = value
                    elif value[0] == fence[0] and len(value) >= len(fence):
                        fence = None
                    continue
                heading = re.match(r"^\s{0,3}(#{1,6})\s+(.+?)\s*#*\s*$", line)
                if heading and fence is None:
                    level = len(heading[1])
                    headings = [(depth, title) for depth, title in headings if depth < level]
                    headings.append((level, heading[2]))
                    boundaries[number] = " / ".join(title for _, title in headings)
            elif suffix == ".toml":
                heading = re.match(r"^\s*\[\[?([^\]]+)\]\]?\s*(?:#.*)?$", line)
                if heading:
                    boundaries[number] = heading[1]
            elif suffix == ".sql":
                heading = re.match(r"^\s*create\s+(?:or\s+(?:alter|replace)\s+)?(?:procedure|proc|function|view|table|trigger|index)\s+([\w.\[\]\"]+)",
                                   line, re.IGNORECASE)
                if heading:
                    boundaries[number] = heading[1]
            elif suffix in {".js", ".jsx", ".mjs", ".cjs", ".ts", ".tsx", ".cs", ".java", ".kt", ".kts", ".go", ".rs", ".c", ".cpp", ".ps1", ".psm1", ".sh"}:
                # A closing brace in column 0 ends a top-level declaration; code after it belongs to the module again.
                if line.startswith("}"):
                    boundaries.setdefault(number + 1, "<module>")
                symbol = re.match(
                    r"^\s*(?:(?:export|default|public|private|protected|internal|static|abstract|async|sealed|partial|pub)\s+)*"
                    r"(?:(?:class|interface|struct|enum|record|function|fun|fn|def)\s+([\w$-]+)"
                    r"|func\s+(?:\([^)]*\)\s*)?([\w]+)"
                    r"|(?:const|let|var)\s+([\w$]+)\s*=.*=>"
                    r"|([\w$-]+)\s*\([^;]*\)\s*\{"
                    r"|[\w<>\[\],?]+\s+([\w]+)\s*\([^;]*\)\s*(?:\{|=>|$))", line)
                if symbol:
                    name = next(value for value in symbol.groups() if value)
                    if name not in {"if", "for", "foreach", "while", "switch", "catch", "using", "lock"}:
                        boundaries[number] = name
    positions = sorted(boundaries)
    for start, end in zip(positions, positions[1:]):
        content = "".join(lines[start:end])
        if content.strip():
            yield boundaries[start], content


def chunks(text: str, path: str, tokenizer):
    for symbol, content in sections(text, path):
        offsets = tokenizer(content, add_special_tokens=False, return_offsets_mapping=True)["offset_mapping"]
        start = 0
        while start < len(offsets):
            end = min(start + CHUNK_TOKENS, len(offsets))
            first = 0 if start == 0 else offsets[start][0]
            while end > start:
                last = len(content) if end == len(offsets) else offsets[end][0]
                value = content[first:last].strip()
                if len(tokenizer(value, add_special_tokens=False)["input_ids"]) <= CHUNK_TOKENS:
                    break
                end -= 1
            if value:
                yield symbol, value
            if end == len(offsets):
                break
            start = max(start + 1, end - OVERLAP_TOKENS)


def document_text(path: str, symbol: str, content: str):
    return f"Path: {path}\nSymbol: {symbol}\n\n{content}"


def run_action(action: str, root: Path, data_dir: Path, model_cache: Path, query: str | None, top_k: int, output):
    model_cache.mkdir(parents=True, exist_ok=True)
    os.environ.setdefault("HF_HOME", str(model_cache.resolve()))
    index = Index(root, data_dir)
    result = index.search(query, top_k) if action == "search" else {"chunks": index.rebuild()}
    output.put((True, result))


def preferred_device():
    import torch
    return "cuda" if torch.cuda.is_available() else "cpu"


def model_kwargs():
    """Use bfloat16 on GPUs that support it; CPUs keep float32, where bfloat16 is often slower."""
    import torch
    return {"dtype": torch.bfloat16} if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else {}


def read_max_files():
    value = os.environ.get(MAX_FILES_ENV)
    if value is None:
        return DEFAULT_MAX_FILES
    try:
        parsed = int(value)
    except ValueError:
        parsed = 0
    if parsed <= 0:
        raise ValueError(f"{MAX_FILES_ENV} must be a positive integer, got {value!r}")
    return parsed


_shared_model_lock = threading.Lock()
_shared_encoder = None
_shared_reranker = None


def shared_encoder():
    global _shared_encoder
    with _shared_model_lock:
        if _shared_encoder is None:
            from sentence_transformers import SentenceTransformer
            _shared_encoder = SentenceTransformer(MODEL, revision=MODEL_REVISION, truncate_dim=EMBEDDING_DIMENSIONS,
                                                   prompts={"query": QUERY_PROMPT}, device=preferred_device(),
                                                   model_kwargs=model_kwargs())
        return _shared_encoder


def shared_reranker():
    global _shared_reranker
    with _shared_model_lock:
        if _shared_reranker is None:
            from sentence_transformers import CrossEncoder
            _shared_reranker = CrossEncoder(RERANKER_MODEL, revision=RERANKER_MODEL_REVISION, prompts={"query": RETRIEVAL_INSTRUCTION},
                                             default_prompt_name="query", device=preferred_device(),
                                             model_kwargs=model_kwargs())
        return _shared_reranker


def preload_models(model_cache: Path):
    model_cache.mkdir(parents=True, exist_ok=True)
    os.environ["HF_HOME"] = str(model_cache.resolve())
    query = "Verify semantic retrieval models."
    embeddings = shared_encoder().encode([query], prompt_name="query")
    if len(embeddings) != 1 or len(embeddings[0]) != EMBEDDING_DIMENSIONS:
        raise ValueError("Embedding model returned unexpected dimensions.")
    scores = shared_reranker().predict([(query, query)], show_progress_bar=False)
    if len(scores) != 1:
        raise ValueError("Reranking model returned unexpected scores.")


def reciprocal_rank_fusion(*rankings):
    scores = {}
    for ranking in rankings:
        for rank, identifier in enumerate(ranking, 1):
            scores[identifier] = scores.get(identifier, 0.0) + 1.0 / (RRF_K + rank)
    return sorted(scores, key=lambda identifier: (-scores[identifier], identifier))


def code_version() -> str:
    """Hash of this file, so a daemon started by older code is never reused after an update."""
    return hashlib.sha256(Path(__file__).read_bytes()).hexdigest()[:12]


def daemon_lock_path(data_dir: Path) -> Path:
    return data_dir / f"daemon-{code_version()}.starting"


def daemon_address(data_dir: Path) -> str:
    identity = os.path.normcase(str(data_dir.resolve())) + "|" + code_version()
    digest = hashlib.sha256(identity.encode()).hexdigest()[:16]
    if os.name == "nt":
        return r"\\.\pipe\ai-config-retrieval-" + digest
    return str(data_dir / f"daemon-{digest}.sock")


def daemon_authkey(data_dir: Path) -> bytes:
    if data_dir.is_symlink():
        raise OSError("Semantic retrieval data directory must not be a link.")
    data_dir.mkdir(parents=True, exist_ok=True, mode=0o700)
    path = data_dir / "daemon.key"
    try:
        descriptor = os.open(path, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o600)
    except FileExistsError:
        pass
    else:
        with os.fdopen(descriptor, "wb") as stream:
            stream.write(os.urandom(32))
            stream.flush()
            os.fsync(stream.fileno())
    metadata = path.lstat()
    if not stat.S_ISREG(metadata.st_mode) or (os.name != "nt" and metadata.st_mode & 0o077):
        raise OSError("Semantic retrieval daemon credential is not private.")
    key = path.read_bytes()
    if len(key) != 32:
        raise OSError("Semantic retrieval daemon credential is incomplete.")
    return key


def request_daemon(address: str, payload: dict, data_dir: Path):
    with Client(address, authkey=daemon_authkey(data_dir)) as connection:
        connection.send_bytes(json.dumps(payload).encode("utf-8"))
        if not connection.poll(CLI_TIMEOUT_SECONDS):
            raise TimeoutError(f"Semantic retrieval daemon exceeded the {CLI_TIMEOUT_SECONDS} second timeout.")
        success, result = json.loads(connection.recv_bytes(DAEMON_MESSAGE_LIMIT))
    if not success:
        raise RuntimeError(str(result))
    return result


def spawn_daemon(data_dir: Path, model_cache: Path) -> bool:
    lock_path = daemon_lock_path(data_dir)
    try:
        os.close(os.open(lock_path, os.O_CREAT | os.O_EXCL | os.O_WRONLY))
    except FileExistsError:
        try:
            if time.time() - lock_path.stat().st_mtime > DAEMON_READY_TIMEOUT_SECONDS:
                lock_path.unlink(missing_ok=True)
        except OSError:
            pass
        return True
    try:
        subprocess.Popen(
            [sys.executable, str(Path(__file__).resolve()), "daemon", "--data-dir", str(data_dir), "--model-cache", str(model_cache)],
            stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
            start_new_session=(os.name != "nt"),
            creationflags=(subprocess.DETACHED_PROCESS | subprocess.CREATE_NEW_PROCESS_GROUP) if os.name == "nt" else 0)
        return True
    except OSError:
        lock_path.unlink(missing_ok=True)
        return False


def try_daemon(root: Path, data_dir: Path, model_cache: Path, action: str, query: str | None, top_k: int):
    address = daemon_address(data_dir)
    payload = {"action": action, "root": str(root), "query": query, "top_k": top_k}
    try:
        return request_daemon(address, payload, data_dir)
    except AuthenticationError:
        return None
    except (OSError, EOFError):
        pass
    if not spawn_daemon(data_dir, model_cache):
        return None
    deadline = time.monotonic() + DAEMON_READY_TIMEOUT_SECONDS
    while time.monotonic() < deadline:
        try:
            return request_daemon(address, payload, data_dir)
        except (OSError, EOFError, AuthenticationError):
            time.sleep(0.5)
    return None


def run_daemon(data_dir: Path, model_cache: Path):
    address = daemon_address(data_dir)
    if os.name != "nt":
        Path(address).unlink(missing_ok=True)
    model_cache.mkdir(parents=True, exist_ok=True)
    os.environ.setdefault("HF_HOME", str(model_cache.resolve()))
    indexes: dict[str, Index] = {}
    activity_lock = threading.Lock()
    last_activity = time.monotonic()

    def watchdog():
        while True:
            time.sleep(5)
            with activity_lock:
                idle_for = time.monotonic() - last_activity
            if idle_for > DAEMON_IDLE_SECONDS:
                os._exit(0)

    threading.Thread(target=watchdog, daemon=True).start()
    listener = Listener(address, family="AF_PIPE" if os.name == "nt" else "AF_UNIX", authkey=daemon_authkey(data_dir))
    daemon_lock_path(data_dir).unlink(missing_ok=True)
    try:
        while True:
            connection = listener.accept()
            try:
                response = (True, None)
                try:
                    request = json.loads(connection.recv_bytes(DAEMON_MESSAGE_LIMIT))
                    if isinstance(request, dict) and request.get("action") == "stop":
                        connection.send_bytes(json.dumps((True, None)).encode("utf-8"))
                        os._exit(0)
                    if not isinstance(request, dict) or request.get("action") not in ("search", "rebuild") or not isinstance(request.get("root"), str):
                        raise ValueError("Invalid semantic retrieval daemon request.")
                    index = indexes.get(request["root"])
                    if index is None:
                        if len(indexes) >= 8:
                            indexes.pop(next(iter(indexes))).db.close()
                        index = Index(Path(request["root"]), data_dir)
                        indexes[request["root"]] = index
                    if request["action"] == "search":
                        response = (True, index.search(request["query"], request["top_k"]))
                    else:
                        response = (True, {"chunks": index.rebuild()})
                except EOFError:
                    continue
                except Exception as error:
                    response = (False, str(error))
                connection.send_bytes(json.dumps(response).encode("utf-8"))
            except (OSError, EOFError):
                pass
            finally:
                connection.close()
                with activity_lock:
                    last_activity = time.monotonic()
    finally:
        listener.close()


def index_database(root: Path, data: Path) -> Path:
    identity = json.dumps([os.path.normcase(str(root.resolve())), MODEL, MODEL_REVISION,
                           EMBEDDING_DIMENSIONS, RETRIEVAL_INSTRUCTION, INDEX_VERSION,
                           CHUNK_TOKENS, OVERLAP_TOKENS])
    return data.resolve() / (hashlib.sha256(identity.encode()).hexdigest() + ".sqlite3")


def daemon_running(data_dir: Path) -> bool:
    if not pipes_available():
        return False
    try:
        Client(daemon_address(data_dir), authkey=daemon_authkey(data_dir)).close()
        return True
    except (OSError, EOFError, AuthenticationError):
        return False


def index_status(root: Path, data_dir: Path) -> dict:
    """Report what the next search costs without loading models or starting the daemon."""
    database = index_database(root, data_dir)
    previous = {}
    if database.is_file():
        connection = sqlite3.connect(database.as_uri() + "?mode=ro", uri=True)
        try:
            previous = {path: (size, mtime, ctime) for path, size, mtime, ctime
                        in connection.execute("select path, size, mtime, ctime from files")}
        except sqlite3.OperationalError:
            pass
        finally:
            connection.close()
    files = indexable_files(root.resolve(), data_dir.resolve(), read_max_files())
    stale = 0
    for path in files:
        metadata = path.stat()
        stale += previous.get(path.relative_to(root.resolve()).as_posix()) != (metadata.st_size, metadata.st_mtime_ns, metadata.st_ctime_ns)
    return {"indexable_files": len(files), "indexed_files": len(previous), "stale_files": stale,
            "models_loaded": daemon_running(data_dir), "in_process": not pipes_available()}


def indexable_files(root: Path, data: Path, max_files: int):
    ignored = {".git", ".venv", "generated", "node_modules", ".my-ai-configuration", "bin", "obj"}
    output = subprocess.run(["git", "-C", str(root), "ls-files", "--cached", "--others",
                             "--exclude-standard", "-z", "--", "."],
                            check=True, capture_output=True, timeout=30).stdout
    candidates = [root / name for name in
                  output.decode("utf-8", errors="surrogateescape").split("\0") if name
                  and not any(part in ignored for part in Path(name).parts)]
    result = []
    for candidate in sorted(candidates):
        if candidate.is_symlink():
            continue
        path = candidate.resolve()
        if (not path.is_relative_to(root) or path.is_relative_to(data)
                or path.suffix.lower() not in TEXT_EXTENSIONS or not path.is_file() or path.stat().st_size > 524288):
            continue
        result.append(path)
        if len(result) == max_files:
            print(f"Semantic retrieval: stopped scanning after {max_files} indexable files; "
                  f"some files under {root} were not indexed. Set {MAX_FILES_ENV} to raise this limit.",
                  file=sys.stderr)
            return result
    return result


class Index:
    def __init__(self, root: Path, data: Path):
        self.root = root.resolve()
        self.data = data.resolve()
        self.data.mkdir(parents=True, exist_ok=True)
        self.db = sqlite3.connect(index_database(self.root, self.data), check_same_thread=False, timeout=30)
        self.db.execute("create table if not exists chunks (path text, symbol text, text text, vector blob)")
        self.db.execute("create index if not exists chunks_path on chunks (path)")
        self.db.execute("create virtual table if not exists lexical using fts5(document)")
        self.db.execute("create table if not exists files (path text primary key, digest text, size integer, mtime integer, ctime integer)")
        columns = {row[1] for row in self.db.execute("pragma table_info(files)")}
        for name in ("size", "mtime", "ctime"):
            if name not in columns:
                self.db.execute(f"alter table files add column {name} integer")
        self.db.execute("create table if not exists query_cache (cache_key text primary key, result text not null)")
        self.db.commit()
        self.lock = threading.RLock()
        self.max_files = read_max_files()
        self.vector_index = None

    def files(self):
        return indexable_files(self.root, self.data, self.max_files)

    def encoder(self):
        return shared_encoder()

    def reranker(self):
        return shared_reranker()

    def rerank(self, query: str, rows: list[dict], limit: int):
        candidates = rows[:RERANK_CANDIDATES]
        scores = self.reranker().predict([(query, document_text(item["path"], item["symbol"], item["text"]))
                                          for item in candidates], batch_size=8)
        for item, score in zip(candidates, scores):
            item["score"] = round(float(score), 6)
        return sorted(candidates, key=lambda item: item["score"], reverse=True)[:min(limit, FINAL_RESULTS)]

    def delete_path(self, path):
        self.db.execute("delete from lexical where rowid in (select rowid from chunks where path = ?)", (path,))
        self.db.execute("delete from chunks where path = ?", (path,))

    def rebuild(self):
        with self.lock:
            self.refresh(force=True)
            return self.db.execute("select count(*) from chunks").fetchone()[0]

    def refresh(self, force=False):
        previous = {path: (digest, size, mtime, ctime)
                    for path, digest, size, mtime, ctime in self.db.execute("select path, digest, size, mtime, ctime from files")}
        encoder = None
        changed = []
        documents = []
        document_offsets = []
        for path in self.files():
            relative = path.relative_to(self.root).as_posix()
            metadata = path.stat()
            size, mtime, ctime = metadata.st_size, metadata.st_mtime_ns, metadata.st_ctime_ns
            old_digest, old_size, old_mtime, old_ctime = previous.pop(relative, (None, None, None, None))
            if not force and old_digest and (size, mtime, ctime) == (old_size, old_mtime, old_ctime):
                continue
            content = path.read_bytes()
            digest = hashlib.sha256(content).hexdigest()
            if not force and digest == old_digest:
                with self.db:
                    self.db.execute("update files set size = ?, mtime = ?, ctime = ? where path = ?",
                                    (size, mtime, ctime, relative))
                continue
            if encoder is None:
                encoder = self.encoder()
            values = list(chunks(content.decode("utf-8", errors="ignore"), relative, encoder.tokenizer))
            file_documents = [document_text(relative, symbol, value) for symbol, value in values]
            document_offsets.append(len(documents))
            documents.extend(file_documents)
            changed.append((relative, digest, size, mtime, ctime, values, file_documents))
        vectors = encoder.encode(documents, normalize_embeddings=True) if documents else []
        with self.db:
            for number, (relative, digest, size, mtime, ctime, values, file_documents) in enumerate(changed):
                start = document_offsets[number]
                file_vectors = vectors[start:start + len(file_documents)]
                self.delete_path(relative)
                for (symbol, value), document, vector in zip(values, file_documents, file_vectors):
                    if len(vector) != EMBEDDING_DIMENSIONS:
                        raise ValueError("Embedding must have 1024 dimensions")
                    cursor = self.db.execute("insert into chunks values (?, ?, ?, ?)",
                                             (relative, symbol, value, vector.tobytes()))
                    self.db.execute("insert into lexical(rowid, document) values (?, ?)", (cursor.lastrowid, document))
                self.db.execute("insert or replace into files values (?, ?, ?, ?, ?)",
                                (relative, digest, size, mtime, ctime))
            for relative in previous:
                self.delete_path(relative)
                self.db.execute("delete from files where path = ?", (relative,))
            if changed or previous:
                self.db.execute("delete from query_cache")
                self.vector_index = None

    def dense_candidates(self, vector):
        import numpy as np
        from usearch.index import Index as VectorIndex

        if self.vector_index is None:
            index = VectorIndex(ndim=EMBEDDING_DIMENSIONS, metric="cos", dtype="f32",
                                connectivity=16, expansion_add=128, expansion_search=64)
            cursor = self.db.execute("select rowid, vector from chunks")
            while rows := cursor.fetchmany(4096):
                vectors = np.frombuffer(b"".join(row[1] for row in rows), dtype=np.float32).reshape(-1, EMBEDDING_DIMENSIONS)
                index.add([row[0] for row in rows], vectors)
            self.vector_index = index
        return [int(match.key) for match in self.vector_index.search(vector, DENSE_CANDIDATES)]

    def cache_key(self, query: str, top_k: int):
        identity = json.dumps([
            query, top_k,
            RERANKER_MODEL, RERANKER_MODEL_REVISION, RETRIEVAL_INSTRUCTION, QUERY_PROMPT,
            DENSE_CANDIDATES, LEXICAL_CANDIDATES, RERANK_CANDIDATES, FINAL_RESULTS, RRF_K,
        ], ensure_ascii=False, separators=(",", ":"))
        return hashlib.sha256(identity.encode("utf-8")).hexdigest()

    def search(self, query: str, top_k: int = FINAL_RESULTS):
        if not query.strip():
            raise ValueError("query must not be empty")
        with self.lock:
            self.refresh()
            if not self.db.execute("select 1 from chunks limit 1").fetchone():
                return []
            key = self.cache_key(query, top_k)
            cached = self.db.execute("select result from query_cache where cache_key = ?", (key,)).fetchone()
            if cached:
                try:
                    return json.loads(cached[0])
                except (TypeError, json.JSONDecodeError):
                    with self.db:
                        self.db.execute("delete from query_cache where cache_key = ?", (key,))
            vector = self.encoder().encode([query], prompt_name="query", normalize_embeddings=True)[0]
            if len(vector) != EMBEDDING_DIMENSIONS:
                raise ValueError("Embedding must have 1024 dimensions")
            dense = self.dense_candidates(vector)
            terms = list(dict.fromkeys(re.findall(r"[^\W_]+", query)))[:128]
            expression = " OR ".join(f'"{term}"' for term in terms)
            lexical = [row[0] for row in self.db.execute(
                "select rowid from lexical where lexical match ? order by bm25(lexical), rowid limit ?",
                (expression, LEXICAL_CANDIDATES))] if terms else []
            identifiers = reciprocal_rank_fusion(dense, lexical)[:RERANK_CANDIDATES]
            placeholders = ",".join("?" for _ in identifiers)
            rows = {identifier: {"path": path, "symbol": symbol, "text": text}
                    for identifier, path, symbol, text in self.db.execute(
                        f"select rowid, path, symbol, text from chunks where rowid in ({placeholders})", identifiers)}
            candidates = [rows[identifier] for identifier in identifiers]
            result = self.rerank(query, candidates, max(1, min(top_k, FINAL_RESULTS)))
            with self.db:
                self.db.execute("insert or replace into query_cache values (?, ?)",
                                (key, json.dumps(result, ensure_ascii=False)))
            return result


def pipes_available() -> bool:
    """Sandboxes such as Codex on Windows forbid the pipes that the daemon and worker process need."""
    try:
        first, second = multiprocessing.Pipe()
    except PermissionError:
        return False
    first.close()
    second.close()
    return True


def run_direct(action: str, root: Path, data_dir: Path, model_cache: Path, query: str | None, top_k: int):
    if not pipes_available():
        # ponytail: in-process fallback has no worker timeout; the client's own command timeout still applies.
        output = queue.Queue(1)
        run_action(action, root, data_dir, model_cache, query, top_k, output)
        return output.get()[1]
    output = multiprocessing.Queue(1)
    worker = multiprocessing.Process(target=run_action, args=(action, root, data_dir, model_cache, query, top_k, output))
    worker.start()
    worker.join(CLI_TIMEOUT_SECONDS)
    if worker.is_alive():
        worker.terminate()
        worker.join()
        raise TimeoutError(f"Qwen retrieval exceeded the {CLI_TIMEOUT_SECONDS} second timeout.")
    if worker.exitcode != 0 or output.empty():
        raise RuntimeError("Qwen retrieval worker failed.")
    success, result = output.get()
    if not success:
        raise RuntimeError(str(result))
    return result


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("action", choices=("search", "rebuild", "status", "stop", "daemon", "preload"))
    parser.add_argument("--query")
    parser.add_argument("--top-k", type=int, default=FINAL_RESULTS)
    parser.add_argument("--root", default=os.getcwd(), type=Path)
    parser.add_argument("--data-dir", type=Path)
    parser.add_argument("--model-cache", required=True, type=Path)
    args = parser.parse_args()
    model_cache = args.model_cache.resolve()
    if args.action == "preload":
        preload_models(model_cache)
        return
    if args.data_dir is None:
        parser.error(f"{args.action} requires --data-dir")
    data_dir = args.data_dir.resolve()
    if args.action == "daemon":
        data_dir.mkdir(parents=True, exist_ok=True)
        run_daemon(data_dir, model_cache)
        return
    if args.action == "status":
        print(json.dumps(index_status(args.root, data_dir)))
        return
    if args.action == "stop":
        try:
            request_daemon(daemon_address(data_dir), {"action": "stop"}, data_dir)
            stopped = True
        except (OSError, EOFError, AuthenticationError):
            stopped = False
        print(json.dumps({"stopped": stopped}))
        return
    if args.action == "search" and not args.query:
        parser.error("search requires --query")
    root = args.root.resolve()
    data_dir.mkdir(parents=True, exist_ok=True)
    result = try_daemon(root, data_dir, model_cache, args.action, args.query, args.top_k) if pipes_available() else None
    if result is None:
        result = run_direct(args.action, root, data_dir, model_cache, args.query, args.top_k)
    print(json.dumps(result, ensure_ascii=False))


if __name__ == "__main__":
    use_runtime_python()
    main()