#!/opt/homebrew/bin/python3.12
"""
Loki Mode Hybrid Codebase Search

Combines lexical (ripgrep / grep) and semantic (ChromaDB) retrieval over the
same files the indexer covers, then fuses the two ranked lists with reciprocal
rank fusion (RRF). Results are deduped by file:line and truncated to a token
budget (greedy, highest fused score first).

Design notes:
    - Pure logic (RRF, dedup, budget truncation) is separated from I/O so it can
      be unit tested without a live ChromaDB or ripgrep.
    - chromadb is imported lazily inside the semantic path so this module loads
      and the grep-only fallback works even when chromadb is not installed.
    - When ChromaDB / docker is unreachable, search degrades to grep-only so it
      still returns results instead of erroring.

Usage:
    python tools/hybrid_search.py "rate limit detection"
    python tools/hybrid_search.py "council vote" --grep-only
    python tools/hybrid_search.py "memory retrieval" --semantic-only
    python tools/hybrid_search.py "build prompt" --budget 4000 --top 15
"""

import argparse
import importlib.util
import json
import os
import shutil
import subprocess
import sys
from pathlib import Path

# Project root (tools/hybrid_search.py -> repo root).
PROJECT_ROOT = Path(__file__).parent.parent.resolve()

CHROMA_HOST = os.environ.get("LOKI_CHROMA_HOST", "localhost")
CHROMA_PORT = int(os.environ.get("LOKI_CHROMA_PORT", "8100"))
COLLECTION_NAME = os.environ.get("LOKI_CHROMA_COLLECTION", "loki-codebase")

# Conservative default token budget for the merged result set. Embeddings and
# large dumps cost context, so we keep this small and let callers override.
DEFAULT_TOKEN_BUDGET = 3000
RRF_K = 60


# -----------------------------------------------------------------------------
# Token estimation (reuse memory.token_economics; fall back if unavailable)
# -----------------------------------------------------------------------------

def _load_estimate_tokens():
    try:
        if str(PROJECT_ROOT) not in sys.path:
            sys.path.insert(0, str(PROJECT_ROOT))
        from memory.token_economics import estimate_tokens as _et
        return _et
    except Exception:
        def _fallback(text: str) -> int:
            if not text:
                return 0
            return max(1, len(text) // 4)
        return _fallback


estimate_tokens = _load_estimate_tokens()


# -----------------------------------------------------------------------------
# Pure logic: reciprocal rank fusion, dedup, budget truncation
# -----------------------------------------------------------------------------

def _result_key(item: dict) -> str:
    """Dedup key: a result is identified by its file:line location."""
    return f"{item.get('file', '')}:{item.get('line', 0)}"


def reciprocal_rank_fusion(grep_ranked: list, semantic_ranked: list,
                           k: int = RRF_K) -> list:
    """Fuse two ranked result lists with reciprocal rank fusion.

    RRF score for an item = sum over each list it appears in of 1 / (k + rank),
    where rank is 1-based position in that list. Items are deduped by file:line;
    when the same location appears in both lists the scores add (so locations
    found by both retrievers rank higher). Output is sorted descending by fused
    score with a deterministic tiebreak on (file, line).

    Args:
        grep_ranked: lexical results, best first. Each is a dict with at least
            "file" and "line"; may carry "snippet", "name", etc.
        semantic_ranked: semantic results, best first, same shape.
        k: RRF constant (default 60, the standard value).

    Returns:
        A list of merged dicts, each annotated with "_fused_score" and
        "_sources" (sorted list of which retrievers found it).
    """
    merged: dict = {}

    def _absorb(ranked: list, source: str):
        for rank, item in enumerate(ranked, start=1):
            key = _result_key(item)
            contribution = 1.0 / (k + rank)
            if key not in merged:
                entry = dict(item)
                entry["_fused_score"] = 0.0
                entry["_sources"] = set()
                merged[key] = entry
            merged[key]["_fused_score"] += contribution
            merged[key]["_sources"].add(source)
            # Prefer a snippet if this list has one and we do not yet.
            if not merged[key].get("snippet") and item.get("snippet"):
                merged[key]["snippet"] = item["snippet"]

    _absorb(grep_ranked, "grep")
    _absorb(semantic_ranked, "semantic")

    results = list(merged.values())
    for entry in results:
        entry["_sources"] = sorted(entry["_sources"])
    # Deterministic order: highest fused score, then file, then line.
    results.sort(key=lambda e: (-e["_fused_score"],
                                str(e.get("file", "")),
                                int(e.get("line", 0))))
    return results


def truncate_to_budget(results: list, budget: int) -> list:
    """Greedily keep highest-scored results until the token budget is reached.

    The token cost of a result is estimated from its snippet plus a small fixed
    overhead for the file:line header line. If a single result is larger than
    the whole budget it is skipped and smaller later results are still packed
    (skip-and-continue), so the budget is never exceeded and the function does
    not get stuck on one oversized hit. Assumes results are already sorted by
    desired priority (RRF output is).

    Returns the kept subset (a new list), preserving input order.
    """
    if budget <= 0:
        return []
    kept: list = []
    used = 0
    for item in results:
        header = f"{item.get('file', '')}:{item.get('line', 0)} "
        snippet = item.get("snippet", "") or ""
        cost = estimate_tokens(header) + estimate_tokens(snippet)
        if used + cost > budget:
            # Skip this one; a smaller later result may still fit.
            continue
        kept.append(item)
        used += cost
    return kept


# -----------------------------------------------------------------------------
# I/O: file scope, lexical search, semantic search
# -----------------------------------------------------------------------------

def _indexer_files() -> list:
    """The set of files the indexer covers, as absolute paths.

    Imports the indexer's collect_files() when possible (single source of truth);
    falls back to a directory glob if the indexer cannot be imported (e.g.
    chromadb missing). The fallback keeps grep-only search working.
    """
    try:
        if str(PROJECT_ROOT / "tools") not in sys.path:
            sys.path.insert(0, str(PROJECT_ROOT / "tools"))
        import importlib
        spec = importlib.util.spec_from_file_location(
            "loki_index_codebase", str(PROJECT_ROOT / "tools" / "index-codebase.py"))
        mod = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(mod)
        return [str(fp) for fp, _ in mod.collect_files()]
    except Exception:
        # Fallback: scan common code dirs without importing chromadb.
        paths = []
        for sub in ("autonomy", "providers", "memory", "dashboard", "mcp",
                    "swarm", "learning", "events", "state", "skills", "tests"):
            d = PROJECT_ROOT / sub
            if not d.is_dir():
                continue
            for ext in ("*.sh", "*.py", "*.md"):
                paths.extend(str(p) for p in d.rglob(ext))
        for top in ("SKILL.md", "CLAUDE.md", "autonomy/loki"):
            p = PROJECT_ROOT / top
            if p.exists():
                paths.append(str(p))
        return sorted(set(paths))


def _have_ripgrep() -> bool:
    return shutil.which("rg") is not None


def grep_search(query: str, files: list, top: int = 30) -> tuple:
    """Lexical search over the given files. Returns (results, tool_used).

    Prefers ripgrep; falls back to grep -rn; falls back to a pure-python scan.
    Results are ranked by per-file match count (more matches first), then by
    line number, then file path for determinism. Each result is a dict with
    file (rel path), line, snippet, name, source="grep".
    """
    if not query.strip() or not files:
        return [], "none"

    matches: list = []  # (abs_file, line_no, text)
    tool_used = "none"

    if _have_ripgrep():
        tool_used = "ripgrep"
        try:
            cmd = ["rg", "--no-heading", "--line-number", "--color", "never",
                   "--fixed-strings", "-e", query] + files
            proc = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
            for line in proc.stdout.splitlines():
                # format: path:line:content
                parts = line.split(":", 2)
                if len(parts) < 3:
                    continue
                fpath, lno, text = parts[0], parts[1], parts[2]
                if not lno.isdigit():
                    continue
                matches.append((fpath, int(lno), text))
        except Exception:
            matches = []
            tool_used = "none"

    if tool_used == "none" and shutil.which("grep"):
        tool_used = "grep"
        try:
            cmd = ["grep", "-rnF", "--", query] + files
            proc = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
            for line in proc.stdout.splitlines():
                parts = line.split(":", 2)
                if len(parts) < 3:
                    continue
                fpath, lno, text = parts[0], parts[1], parts[2]
                if not lno.isdigit():
                    continue
                matches.append((fpath, int(lno), text))
        except Exception:
            matches = []
            tool_used = "none"

    if tool_used == "none":
        # Pure-python scan as a last resort (no external tools at all).
        tool_used = "python-scan"
        needle = query
        for fpath in files:
            try:
                with open(fpath, "r", errors="replace") as fh:
                    for i, text in enumerate(fh, start=1):
                        if needle in text:
                            matches.append((fpath, i, text.rstrip("\n")))
            except Exception:
                continue

    # Rank by per-file match count desc, then line asc, then path.
    counts: dict = {}
    for fpath, _, _ in matches:
        counts[fpath] = counts.get(fpath, 0) + 1

    def _rel(p: str) -> str:
        try:
            return str(Path(p).resolve().relative_to(PROJECT_ROOT))
        except Exception:
            return p

    matches.sort(key=lambda m: (-counts[m[0]], m[1], m[0]))

    results = []
    for fpath, lno, text in matches[:top]:
        results.append({
            "file": _rel(fpath),
            "line": lno,
            "name": "",
            "snippet": text.strip()[:300],
            "source": "grep",
        })
    return results, tool_used


def semantic_search(query: str, top: int = 30) -> tuple:
    """Semantic search via ChromaDB. Returns (results, available: bool).

    Imports chromadb lazily so a missing dependency or stopped container does
    not break the module. On any failure returns ([], False) so callers can
    fall back to grep-only.
    """
    try:
        import chromadb
        client = chromadb.HttpClient(host=CHROMA_HOST, port=CHROMA_PORT)
        client.heartbeat()
        collection = client.get_collection(name=COLLECTION_NAME)
        res = collection.query(
            query_texts=[query],
            n_results=top,
            include=["documents", "metadatas", "distances"],
        )
    except Exception:
        return [], False

    results = []
    ids = res.get("ids", [[]])
    if not ids or not ids[0]:
        return [], True
    for i in range(len(ids[0])):
        meta = res["metadatas"][0][i]
        doc = res["documents"][0][i]
        dist = res["distances"][0][i]
        results.append({
            "file": meta.get("file", ""),
            "line": meta.get("line", 0),
            "name": meta.get("name", ""),
            "snippet": (doc[:300] if doc else ""),
            "relevance": round(max(0.0, 1.0 - dist / 2.0), 4),
            "source": "semantic",
        })
    return results, True


# -----------------------------------------------------------------------------
# Orchestration
# -----------------------------------------------------------------------------

def hybrid_search(query: str, top: int = 10, budget: int = DEFAULT_TOKEN_BUDGET,
                  grep_only: bool = False, semantic_only: bool = False) -> dict:
    """Run hybrid (or single-mode) search and return a structured result dict.

    Returns:
        {
          "query": str,
          "results": [ {file, line, snippet, sources, fused_score}, ... ],
          "grep_tool": str,            # ripgrep | grep | python-scan | none
          "semantic_available": bool,
          "mode": str,                 # hybrid | grep-only | semantic-only
          "fallback": bool,            # True if semantic was requested but down
          "budget": int,
        }
    """
    files = _indexer_files()
    fallback = False

    grep_results: list = []
    grep_tool = "none"
    if not semantic_only:
        grep_results, grep_tool = grep_search(query, files, top=max(top, 30))

    semantic_results: list = []
    semantic_available = False
    if not grep_only:
        semantic_results, semantic_available = semantic_search(query, top=max(top, 30))
        if not semantic_available and not semantic_only:
            # Semantic requested as part of hybrid but unavailable -> grep-only.
            fallback = True

    if semantic_only:
        fused = reciprocal_rank_fusion([], semantic_results)
        mode = "semantic-only"
    elif grep_only:
        fused = reciprocal_rank_fusion(grep_results, [])
        mode = "grep-only"
    else:
        fused = reciprocal_rank_fusion(grep_results, semantic_results)
        mode = "hybrid"

    fused = fused[:top]
    kept = truncate_to_budget(fused, budget)

    out = []
    for item in kept:
        out.append({
            "file": item.get("file", ""),
            "line": item.get("line", 0),
            "name": item.get("name", ""),
            "snippet": item.get("snippet", ""),
            "sources": item.get("_sources", []),
            "fused_score": round(item.get("_fused_score", 0.0), 6),
        })

    return {
        "query": query,
        "results": out,
        "grep_tool": grep_tool,
        "semantic_available": semantic_available,
        "mode": mode,
        "fallback": fallback,
        "budget": budget,
    }


def _render_text(payload: dict) -> str:
    """Render a hybrid_search payload as plain text (no emojis, no dashes)."""
    lines = []
    mode = payload["mode"]
    note = ""
    if payload.get("fallback"):
        note = " (semantic index unavailable, grep-only fallback)"
    grep_tool = payload.get("grep_tool", "none")
    if grep_tool in ("grep", "python-scan") and mode != "semantic-only":
        note += f" (ripgrep not found, using {grep_tool})"
    lines.append(f"hybrid search: {payload['query']!r} [{mode}]{note}")
    lines.append(f"budget: {payload['budget']} tokens, "
                 f"{len(payload['results'])} result(s)")
    lines.append("")
    if not payload["results"]:
        lines.append("no matches.")
        return "\n".join(lines)
    for i, r in enumerate(payload["results"], start=1):
        src = ",".join(r.get("sources", [])) or "?"
        lines.append(f"[{i}] {r['file']}:{r['line']}  "
                     f"(match: {src}, score: {r['fused_score']})")
        snip = (r.get("snippet") or "").strip()
        if snip:
            lines.append(f"    {snip}")
    return "\n".join(lines)


def main():
    parser = argparse.ArgumentParser(
        description="Hybrid (grep + semantic) codebase search")
    parser.add_argument("query", help="Search query")
    parser.add_argument("--top", type=int, default=10, help="Max results")
    parser.add_argument("--budget", type=int, default=DEFAULT_TOKEN_BUDGET,
                        help="Token budget for the merged result set")
    parser.add_argument("--grep-only", action="store_true",
                        help="Lexical search only (skip semantic)")
    parser.add_argument("--semantic-only", action="store_true",
                        help="Semantic search only (skip grep)")
    parser.add_argument("--json", action="store_true", help="Output JSON")
    args = parser.parse_args()

    if args.grep_only and args.semantic_only:
        print("error: --grep-only and --semantic-only are mutually exclusive",
              file=sys.stderr)
        return 2

    payload = hybrid_search(
        args.query, top=args.top, budget=args.budget,
        grep_only=args.grep_only, semantic_only=args.semantic_only)

    if args.json:
        print(json.dumps(payload, indent=2))
    else:
        print(_render_text(payload))
    return 0


if __name__ == "__main__":
    sys.exit(main())