"""Combine FTS5 term substitutions and semantic rankings with RRF."""

import argparse
import re
import sqlite3
from dataclasses import dataclass
from pathlib import Path

from chromadb.errors import ChromaError

from search_bm25 import search as search_bm25
from search_chroma import search as search_chroma

STOP_WORDS = {
    "a",
    "an",
    "and",
    "are",
    "as",
    "at",
    "be",
    "by",
    "for",
    "from",
    "in",
    "is",
    "it",
    "of",
    "on",
    "or",
    "that",
    "the",
    "to",
    "was",
    "with",
}


# Exact whole-word replacements, looked up case-insensitively.
# Values replace the original FTS terms; they do not expand the query with synonyms.
SUBSTITUTIONS = {
    "close": "kill",
    "bigger": "resize",
    "smaller": "resize",
    "2": "two",
    "go": "move",
}


@dataclass
class Result:
    key: str
    description: str
    score: float = 0.0
    fts_rank: int | None = None
    semantic_rank: int | None = None


def fts_query(query: str) -> str:
    """Remove stop words, substitute terms, then OR-join them for FTS5."""
    words = re.findall(r"[^\W_]+", query, flags=re.UNICODE)
    if not words:
        raise ValueError("Enter at least one word to search for.")
    # Match stop words case-insensitively, preserving the remaining words.
    # Filtering affects the FTS query only, not the index or semantic query.
    words = [word for word in words if word.casefold() not in STOP_WORDS]
    # Apply substitutions only here, after stop-word filtering and before
    # FTS5 tokenization/stemming. Unmapped words retain their original spelling.
    words = [SUBSTITUTIONS.get(word.casefold(), word) for word in words]
    # Quote terms so user-supplied words cannot become FTS operators.
    return " OR ".join(f'"{word}"' for word in words)


def fuse(
    lexical: list[tuple[str, str, float]],
    semantic: list[tuple[str, str, float]],
    k: int = 60,
) -> list[Result]:
    """Fuse ranked lists by key; use positions, not the original scores."""
    if k < 0:
        raise ValueError("RRF k must be nonnegative.")
    combined: dict[str, Result] = {}
    for rank, (key, description, _) in enumerate(lexical, start=1):
        result = combined.setdefault(key, Result(key, description))
        result.fts_rank = rank
        result.score += 1 / (k + rank)
    for rank, (key, description, _) in enumerate(semantic, start=1):
        kkresult = combined.setdefault(key, Result(key, description))
        result.semantic_rank = rank
        result.score += 1 / (k + rank)
    # Missing from a list means zero contribution. Higher fused scores win.
    # Case-sensitive key ordering makes exact score ties deterministic.
    return sorted(combined.values(), key=lambda result: (-result.score, result.key))


def search(
    database: Path,
    chroma_path: Path,
    query: str,
    collection_name: str = "tmux-key-bindings",
    results: int = 5,
    candidates: int = 20,
    k: int = 60,
) -> list[Result]:
    """Retrieve independently, fuse their top candidates, then limit output."""
    if results < 1 or candidates < 1:
        raise ValueError("Results and candidates must be positive.")
    if k < 0:
        raise ValueError("RRF k must be nonnegative.")
    lexical_query = fts_query(query)
    # MATCH cannot accept an empty expression. If every word was excluded,
    # skip lexical retrieval and let the semantic list contribute on its own.
    lexical = search_bm25(database, lexical_query)[:candidates] if lexical_query else []
    # The original query goes to the embedding model, without OR rewriting.
    # Semantic retrieval runs even when FTS5 returns no matches.
    semantic = search_chroma(chroma_path, query, collection_name, candidates)
    return fuse(lexical, semantic, k)[:results]


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("query", help="Text to search with both retrieval methods")
    parser.add_argument(
        "--database", type=Path, default=Path("tmux_key_bindings.sqlite3")
    )
    parser.add_argument("--chroma", type=Path, default=Path("data/chroma"))
    parser.add_argument("--collection", default="tmux-key-bindings")
    parser.add_argument(
        "--results", type=int, default=5, help="Final results (default: 5)"
    )
    parser.add_argument(
        "--candidates",
        type=int,
        default=20,
        help="Maximum entries from each ranked list to fuse (default: 20)",
    )
    parser.add_argument("--k", type=int, default=60, help="RRF constant (default: 60)")
    args = parser.parse_args()
    try:
        rows = search(
            args.database,
            args.chroma,
            args.query,
            args.collection,
            args.results,
            args.candidates,
            args.k,
        )
    except (sqlite3.Error, ChromaError, ValueError) as error:
        parser.exit(1, f"Search failed: {error}\n")
    lexical_query = fts_query(args.query)
    print(f"FTS5 query: {lexical_query or '(skipped: all terms are stop words)'}")
    if not rows:
        print("No results found.")
        return
    print("RRF score\tFTS rank\tSemantic rank\tKey\tDescription")
    for row in rows:
        print(
            f"{row.score:.9g}\t{row.fts_rank or '-'}\t"
            f"{row.semantic_rank or '-'}\t{row.key}\t{row.description}"
        )


if __name__ == "__main__":
    main()

find-command ❯ uv run python search_rrf.py "go to window 3"
FTS5 query: "go" OR "to" OR "window" OR "3"
RRF score	FTS rank	Semantic rank	Key	Description
0.0322580645	2	2	n	Change to the next window.
0.0320184426	1	4	0 to 9	Select windows 0 to 9.
0.0317460317	3	3	p	Change to the previous window.
0.0303308824	4	8	l	Move to the previously selected window.
0.0300904977	8	5	M-n	Move to the next window with a bell or activity marker.

find-command ❯ uv run python search_rrf_substitutions.py "change the size of the pane"
FTS5 query: "change" OR "size" OR "pane"
RRF score	FTS rank	Semantic rank	Key	Description
0.0312805474	6	2	}	Swap the current pane with the next pane.
0.0310096154	5	4	{	Swap the current pane with the previous pane.
0.030798389	3	7	Up, Down, Left, Right	Change to the pane above, below, to the left, or to the right of the current pane.
0.0296703297	10	5	*	Create a new floating pane.
0.0296312555	7	8	M	Clear the marked pane.

find-command ❯ uv run python search_rrf_substitutions.py "change the size of the pane" --results 25
FTS5 query: "change" OR "size" OR "pane"
RRF score	FTS rank	Semantic rank	Key	Description
0.0312805474	6	2	}	Swap the current pane with the next pane.
0.0310096154	5	4	{	Swap the current pane with the previous pane.
0.030798389	3	7	Up, Down, Left, Right	Change to the pane above, below, to the left, or to the right of the current pane.
0.0296703297	10	5	*	Create a new floating pane.
0.0296312555	7	8	M	Clear the marked pane.
0.0293804556	17	1	"	Split the current pane into two, top and bottom.
0.0292360222	11	6	;	Move to the previously active pane.
0.0287828947	4	16	m	Mark the current pane (see select-pane -m).
0.0287784679	9	10	x	Kill the current pane.
0.0286935287	18	3	%	Split the current pane into two, left and right.
0.0277777778	12	12	z	Toggle zoom state of the current pane.
0.0276506484	16	9	M-o	Rotate the panes in the current window backwards.
0.0273972603	13	13	C-o	Rotate the panes in the current window forwards.
0.0273641102	8	19	q	Briefly display pane indexes.
0.027027027	14	14	!	Break the current pane out of the window.
0.0267427349	19	11	C-Up, C-Down, C-Left, C-Right	Resize the current pane in steps of one cell.
0.0263203463	15	17	o	Select the next pane in the current window.
0.0258333333	20	15	M-Up, M-Down, M-Left, M-Right	Resize the current pane in steps of five cells.
0.0163934426	1	-	n	Change to the next window.
0.0161290323	2	-	p	Change to the previous window.
0.0128205128	-	18	M-1 to M-7	Arrange panes in one of the seven preset layouts: even-horizontal, even-vertical, main-horizontal, main-horizontal-mirrored, main-vertical, main-vertical-mirrored, or tiled.