#!/usr/bin/env python3
"""
Semantic + keyword search over an Obsidian vault.

Indexes all markdown files by section (split on headers).
Uses sentence-transformers for embedding-based semantic search
and BM25 for keyword matching, fused via Reciprocal Rank Fusion.
Auto-rebuilds index when vault files change.

Usage:
    vault-search "query"                   # search, return top 5 chunks
    vault-search "query" --top 10          # return top 10 chunks
    vault-search "query" --full            # don't truncate results
    vault-search "query" --semantic        # semantic search only (no BM25)
    vault-search "query" --keyword         # keyword search only (no semantic)
    vault-search "query" --multi           # multi-query mode (3 phrasings)
    vault-search --reindex                 # force reindex
    vault-search --stats                   # show index stats
"""

import sys
import os
import re
import json
import math
import hashlib
import pickle
import numpy as np
from pathlib import Path
from collections import Counter

# --- Configuration ---
# These can be overridden via environment variables for portability.
VAULT_ROOT = Path(os.environ.get("VAULT_ROOT", Path(__file__).resolve().parent))
INDEX_DIR = VAULT_ROOT / ".vault-search"
INDEX_PATH = INDEX_DIR / "index.pkl"
HASH_PATH = INDEX_DIR / "hash.json"
MODEL_NAME = "all-MiniLM-L6-v2"  # Fast, good quality, 80MB

# Directories to skip during indexing. Note: agent-learnings/ is NOT skipped --
# these files are indexed and searchable.
SKIP_DIRS = {".obsidian", ".trash", ".git", "node_modules", ".claude",
             "obsidian-vault-plugin", ".vault-search", "__pycache__",
             "_github_imports", ".lab"}


# ============================================================
# Section Splitting
# ============================================================

def split_into_sections(filepath, content):
    """Split a markdown file into sections by headers. Preserves hierarchy."""
    sections = []
    lines = content.split("\n")
    current_header = filepath.stem
    current_lines = []
    header_level = 0
    parent_headers = {}  # level -> header text

    for line in lines:
        header_match = re.match(r'^(#{1,4})\s+(.+)', line)
        if header_match:
            body = "\n".join(current_lines).strip()
            if body and len(body) > 20:
                breadcrumb = _build_breadcrumb(parent_headers, header_level, current_header)
                sections.append({
                    "file": str(filepath.relative_to(VAULT_ROOT)),
                    "header": current_header,
                    "breadcrumb": breadcrumb,
                    "body": body,
                    "level": header_level,
                })

            level = len(header_match.group(1))
            current_header = header_match.group(2).strip()
            header_level = level
            parent_headers[level] = current_header
            # Clear child headers
            for l in list(parent_headers.keys()):
                if l > level:
                    del parent_headers[l]
            current_lines = []
        else:
            current_lines.append(line)

    body = "\n".join(current_lines).strip()
    if body and len(body) > 20:
        breadcrumb = _build_breadcrumb(parent_headers, header_level, current_header)
        sections.append({
            "file": str(filepath.relative_to(VAULT_ROOT)),
            "header": current_header,
            "breadcrumb": breadcrumb,
            "body": body,
            "level": header_level,
        })

    return sections


def _build_breadcrumb(parent_headers, current_level, current_header):
    """Build a breadcrumb path like 'Top Section > Sub Section > Current'."""
    parts = []
    for level in sorted(parent_headers.keys()):
        if level < current_level:
            parts.append(parent_headers[level])
    parts.append(current_header)
    return " > ".join(parts)


def collect_all_sections():
    """Walk the vault and collect all sections from all .md files."""
    all_sections = []

    for md_file in VAULT_ROOT.rglob("*.md"):
        parts = md_file.relative_to(VAULT_ROOT).parts
        if any(p in SKIP_DIRS or p.startswith(".") for p in parts):
            continue
        try:
            content = md_file.read_text(encoding="utf-8")
        except Exception:
            continue
        sections = split_into_sections(md_file, content)
        all_sections.extend(sections)

    return all_sections


# ============================================================
# BM25 Implementation
# ============================================================

def _tokenize(text):
    """Simple whitespace + punctuation tokenizer with lowercasing."""
    return re.findall(r'[a-z0-9_\-\.#/]+', text.lower())


class BM25Index:
    """Okapi BM25 ranking over document sections."""

    def __init__(self, sections, k1=1.5, b=0.75):
        self.k1 = k1
        self.b = b
        self.sections = sections

        # Tokenize all documents
        self.doc_tokens = []
        self.doc_freqs = []  # per-document term frequencies
        self.doc_lengths = []
        df = Counter()  # document frequency per term

        for s in sections:
            text = f"{s['breadcrumb']} {s['header']} {s['body']}"
            tokens = _tokenize(text)
            freq = Counter(tokens)
            self.doc_tokens.append(tokens)
            self.doc_freqs.append(freq)
            self.doc_lengths.append(len(tokens))
            for term in freq:
                df[term] += 1

        self.N = len(sections)
        self.avgdl = sum(self.doc_lengths) / max(self.N, 1)
        self.df = df

        # Pre-compute IDF
        self.idf = {}
        for term, freq in self.df.items():
            self.idf[term] = math.log((self.N - freq + 0.5) / (freq + 0.5) + 1.0)

    def search(self, query, top_k=50):
        """Return list of (section_index, score) tuples."""
        query_tokens = _tokenize(query)
        scores = []

        for i in range(self.N):
            score = 0.0
            dl = self.doc_lengths[i]
            freq = self.doc_freqs[i]

            for token in query_tokens:
                if token not in self.idf:
                    continue
                tf = freq.get(token, 0)
                idf = self.idf[token]
                numerator = tf * (self.k1 + 1)
                denominator = tf + self.k1 * (1 - self.b + self.b * dl / self.avgdl)
                score += idf * numerator / denominator

            if score > 0:
                scores.append((i, score))

        scores.sort(key=lambda x: -x[1])
        return scores[:top_k]


# ============================================================
# Index Building
# ============================================================

def compute_vault_hash():
    """Hash of all .md file paths + mtimes to detect changes."""
    entries = []
    for md_file in sorted(VAULT_ROOT.rglob("*.md")):
        parts = md_file.relative_to(VAULT_ROOT).parts
        if any(p in SKIP_DIRS or p.startswith(".") for p in parts):
            continue
        try:
            mtime = md_file.stat().st_mtime
            entries.append(f"{md_file}:{mtime}")
        except Exception:
            continue
    return hashlib.md5("\n".join(entries).encode()).hexdigest()


def get_model():
    """Load the sentence-transformer model (cached after first download)."""
    from sentence_transformers import SentenceTransformer
    import logging
    logging.getLogger("sentence_transformers").setLevel(logging.WARNING)
    return SentenceTransformer(MODEL_NAME)


def build_index(force=False):
    """Build or load the embedding + BM25 index."""
    INDEX_DIR.mkdir(exist_ok=True)
    current_hash = compute_vault_hash()

    # Try cached index
    if not force and INDEX_PATH.exists() and HASH_PATH.exists():
        try:
            saved_hash = json.loads(HASH_PATH.read_text())["hash"]
            if saved_hash == current_hash:
                with open(INDEX_PATH, "rb") as f:
                    data = pickle.load(f)
                # Ensure BM25 index exists (may be an old cache)
                if "bm25" not in data:
                    data["bm25"] = BM25Index(data["sections"])
                    with open(INDEX_PATH, "wb") as f:
                        pickle.dump(data, f)
                return data
        except Exception:
            pass

    print("Building vault index...", file=sys.stderr)
    sections = collect_all_sections()
    if not sections:
        print("ERROR: No sections found in vault.", file=sys.stderr)
        sys.exit(1)

    # Create search-optimized text for each section
    documents = []
    for s in sections:
        doc = f"{s['breadcrumb']}\n{s['header']}\n{s['body']}"
        documents.append(doc)

    print(f"Encoding {len(sections)} sections...", file=sys.stderr)
    model = get_model()
    embeddings = model.encode(documents, show_progress_bar=True, convert_to_numpy=True)

    # Normalize for cosine similarity via dot product
    norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
    norms[norms == 0] = 1
    embeddings = embeddings / norms

    # Build BM25 index
    print("Building BM25 index...", file=sys.stderr)
    bm25 = BM25Index(sections)

    index = {
        "sections": sections,
        "embeddings": embeddings,
        "bm25": bm25,
    }

    with open(INDEX_PATH, "wb") as f:
        pickle.dump(index, f)
    HASH_PATH.write_text(json.dumps({"hash": current_hash}))
    print(f"Index built: {len(sections)} sections from {len(set(s['file'] for s in sections))} files", file=sys.stderr)

    return index


# ============================================================
# Search Functions
# ============================================================

def _reciprocal_rank_fusion(ranked_lists, k=60):
    """Fuse multiple ranked lists using Reciprocal Rank Fusion (RRF).

    Each ranked_list is a list of (section_index, score) tuples.
    Returns a list of (section_index, fused_score) sorted by fused_score desc.
    """
    fused_scores = Counter()
    for ranked in ranked_lists:
        for rank, (idx, _score) in enumerate(ranked):
            fused_scores[idx] += 1.0 / (k + rank + 1)

    return sorted(fused_scores.items(), key=lambda x: -x[1])


def search_semantic(query, index, top_k=50):
    """Pure semantic (vector) search. Returns [(section_index, score), ...]."""
    embeddings = index["embeddings"]
    model = get_model()
    query_emb = model.encode([query], convert_to_numpy=True)
    query_emb = query_emb / np.linalg.norm(query_emb)

    with np.errstate(divide='ignore', over='ignore', invalid='ignore'):
        scores = (embeddings @ query_emb.T).flatten()
    scores = np.nan_to_num(scores, nan=0.0)

    ranked = sorted(enumerate(scores), key=lambda x: -x[1])
    return [(i, float(score)) for i, score in ranked if score > 0.15][:top_k]


def search_keyword(query, index, top_k=50):
    """Pure BM25 keyword search. Returns [(section_index, score), ...]."""
    bm25 = index["bm25"]
    return bm25.search(query, top_k)


def search_hybrid(query, index, top_k=5):
    """Hybrid search: semantic + BM25 fused via RRF. Returns [(section, score), ...]."""
    semantic_results = search_semantic(query, index, top_k=50)
    keyword_results = search_keyword(query, index, top_k=50)
    fused = _reciprocal_rank_fusion([semantic_results, keyword_results])
    sections = index["sections"]
    return [(sections[idx], score) for idx, score in fused[:top_k]]


def search_multi_query(query, index, top_k=5):
    """Multi-query search: generate 3 phrasings, fuse results.

    Uses simple query reformulation (no LLM needed):
    1. Original query
    2. Key nouns/verbs extracted
    3. Query with synonyms/related terms appended
    """
    # Phrasing 1: original
    q1 = query

    # Phrasing 2: extract key terms (remove common words)
    stop_words = {'the', 'a', 'an', 'is', 'are', 'was', 'were', 'be', 'been',
                  'being', 'have', 'has', 'had', 'do', 'does', 'did', 'will',
                  'would', 'could', 'should', 'may', 'might', 'can', 'shall',
                  'to', 'of', 'in', 'for', 'on', 'with', 'at', 'by', 'from',
                  'as', 'into', 'through', 'during', 'before', 'after', 'and',
                  'but', 'or', 'not', 'no', 'nor', 'so', 'yet', 'both', 'each',
                  'how', 'what', 'where', 'when', 'which', 'who', 'whom', 'why',
                  'this', 'that', 'these', 'those', 'i', 'we', 'you', 'it',
                  'they', 'my', 'our', 'your', 'its', 'their'}
    tokens = _tokenize(query)
    key_terms = [t for t in tokens if t not in stop_words]
    q2 = " ".join(key_terms) if key_terms else query

    # Phrasing 3: broaden with file/section context words
    context_words = []
    for t in key_terms[:3]:
        # Add common related terms for dev contexts
        if t in ('agent', 'agents'): context_words.extend(['pipeline', 'subagent'])
        elif t in ('hook', 'hooks'): context_words.extend(['event', 'lifecycle'])
        elif t in ('skill', 'skills'): context_words.extend(['command', 'workflow'])
        elif t in ('search', 'find'): context_words.extend(['query', 'index'])
        elif t in ('test', 'testing'): context_words.extend(['verify', 'check'])
        elif t in ('build', 'builder'): context_words.extend(['implement', 'create'])
        elif t in ('error', 'bug'): context_words.extend(['fix', 'issue', 'debug'])
        elif t in ('design', 'ui'): context_words.extend(['component', 'screen', 'layout'])
        elif t in ('color', 'colours'): context_words.extend(['hex', 'shade', 'theme'])
        elif t in ('database', 'db'): context_words.extend(['schema', 'json', 'data'])
    q3 = f"{query} {' '.join(context_words)}" if context_words else query

    # Run all three queries
    queries = list(set([q1, q2, q3]))  # deduplicate if same
    all_ranked = []
    for q in queries:
        semantic = search_semantic(q, index, top_k=30)
        keyword = search_keyword(q, index, top_k=30)
        all_ranked.extend([semantic, keyword])

    fused = _reciprocal_rank_fusion(all_ranked)
    sections = index["sections"]
    return [(sections[idx], score) for idx, score in fused[:top_k]]


# ============================================================
# Output
# ============================================================

def display_results(results, full_output=False, explain=False, query=None, index=None):
    """Display search results."""
    if not results:
        print("No results found.")
        return

    for rank, (section, score) in enumerate(results, 1):
        print(f"--- Result {rank} (score: {score:.3f}) ---")
        print(f"File: {section['file']}")
        print(f"Section: {section['breadcrumb']}")

        if explain and query and index:
            _print_explanation(section, query, index)

        print()
        body = section["body"]
        if not full_output and len(body) > 1500:
            body = body[:1500] + f"\n[... truncated at 1500 chars, full section is {len(section['body'])} chars. Use --full or Read tool for complete content ...]"
        print(body)
        print()


def _print_explanation(section, query, index):
    """Explain why a result matched."""
    query_tokens = _tokenize(query)
    body_tokens = set(_tokenize(section['body']))
    header_tokens = set(_tokenize(section['header'] + ' ' + section['breadcrumb']))

    # Keyword matches
    keyword_hits = [t for t in query_tokens if t in body_tokens or t in header_tokens]
    if keyword_hits:
        print(f"  Keywords matched: {', '.join(set(keyword_hits))}")

    # Semantic similarity hint
    print(f"  Match type: semantic similarity + keyword overlap (hybrid)")


def search(query, top_k=5, full_output=False, mode="hybrid", explain=False):
    """Main search entry point."""
    index = build_index()
    sections = index["sections"]

    if mode == "semantic":
        raw = search_semantic(query, index, top_k=top_k)
        results = [(sections[i], score) for i, score in raw]
    elif mode == "keyword":
        raw = search_keyword(query, index, top_k=top_k)
        results = [(sections[i], score) for i, score in raw]
    elif mode == "multi":
        results = search_multi_query(query, index, top_k=top_k)
    else:  # hybrid (default)
        results = search_hybrid(query, index, top_k=top_k)

    display_results(results, full_output, explain=explain, query=query, index=index)


def show_stats():
    """Show index statistics."""
    index = build_index()
    sections = index["sections"]
    files = set(s["file"] for s in sections)
    print(f"Indexed files: {len(files)}")
    print(f"Total sections: {len(sections)}")
    print(f"Embedding model: {MODEL_NAME}")
    print(f"BM25 index: {index['bm25'].N} documents, {len(index['bm25'].df)} unique terms")
    print(f"Index cached: {INDEX_PATH.exists()}")
    print()
    print("Files indexed:")
    for f in sorted(files):
        count = sum(1 for s in sections if s["file"] == f)
        print(f"  {f} ({count} sections)")


# ============================================================
# CLI
# ============================================================

def main():
    args = sys.argv[1:]

    if not args:
        print(__doc__)
        sys.exit(0)

    if "--reindex" in args:
        build_index(force=True)
        print("Index rebuilt.")
        show_stats()
        return

    if "--stats" in args:
        show_stats()
        return

    # Parse flags
    full_output = "--full" in args
    semantic_only = "--semantic" in args
    keyword_only = "--keyword" in args
    multi_query = "--multi" in args
    explain = "--explain" in args
    flags = {"--full", "--semantic", "--keyword", "--multi", "--explain"}

    top_k = 5
    if "--top" in args:
        idx = args.index("--top")
        if idx + 1 < len(args):
            top_k = int(args[idx + 1])
            args = [a for i, a in enumerate(args) if i != idx and i != idx + 1]

    query = " ".join(a for a in args if a not in flags and not a.startswith("--"))
    if not query:
        print("Usage: vault-search \"your query\"")
        sys.exit(1)

    if semantic_only:
        mode = "semantic"
    elif keyword_only:
        mode = "keyword"
    elif multi_query:
        mode = "multi"
    else:
        mode = "hybrid"

    search(query, top_k, full_output, mode, explain=explain)


if __name__ == "__main__":
    main()