"""Search the persisted tmux embeddings by semantic similarity."""

import argparse
from pathlib import Path
from typing import cast

import chromadb
from chromadb.api.types import Embeddable, EmbeddingFunction
from chromadb.errors import ChromaError
from chromadb.utils.embedding_functions import DefaultEmbeddingFunction


def search(
    chroma_path: Path,
    query: str,
    collection_name: str = "tmux-key-bindings",
    results: int = 5,
) -> list[tuple[str, str, float]]:
    """Return (key, description, distance) tuples, closest first."""
    if not query.strip():
        raise ValueError("Enter a nonempty query.")
    if results < 1:
        raise ValueError("The number of results must be positive.")
    # PersistentClient otherwise creates a database when the path is wrong.
    if not (chroma_path / "chroma.sqlite3").is_file():
        raise ValueError("Chroma database not found. Run create_embeddings.py first.")

    embedding_function = DefaultEmbeddingFunction()
    client = chromadb.PersistentClient(path=str(chroma_path))
    # As in create_embeddings.py, the cast bridges Chroma's text-or-image
    # parameter type with the default model's text-only input type.
    collection = client.get_collection(
        name=collection_name,
        embedding_function=cast(EmbeddingFunction[Embeddable], embedding_function),
    )
    count = collection.count()
    if count == 0:
        return []

    # Embed the query with the same model used for the stored descriptions.
    # This is natural-language text, not SQL LIKE or FTS5 query syntax.
    query_embeddings = embedding_function([query])
    matches = collection.query(
        query_embeddings=query_embeddings,
        n_results=min(results, count),
        include=["documents", "metadatas", "distances"],
    )

    # Chroma returns one list of matches per query; we submitted just one.
    documents = matches["documents"]
    metadatas = matches["metadatas"]
    distances = matches["distances"]
    if documents is None or metadatas is None or distances is None:
        raise ValueError("Chroma did not return the requested result fields.")
    return [
        (str(metadata["key"]), document, distance)
        for metadata, document, distance in zip(
            metadatas[0], documents[0], distances[0], strict=True
        )
    ]


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("query", help="Natural-language description to search for")
    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="Maximum results (default: 5)")
    args = parser.parse_args()
    try:
        rows = search(args.chroma, args.query, args.collection, args.results)
    except (ChromaError, ValueError) as error:
        parser.exit(1, f"Search failed: {error}\n")
    if not rows:
        print("No results found.")
        return
    print("Distance\tKey\tDescription")
    for key, description, distance in rows:
        print(f"{distance:.6f}\t{key}\t{description}")


if __name__ == "__main__":
    main()

find-command ❯ uv run python search_chroma.py "close the current window"
Distance	Key	Description
0.395517	&	Kill the current window.
0.643967	w	Choose the current window interactively.
0.663402	,	Rename the current window.
0.679014	n	Change to the next window.
0.681442	i	Display some information about the current window.

find-command ❯ uv run python search_chroma.py "show tmux key bindings"
Distance	Key	Description
0.765333	?	List all key bindings.
0.794635	~	Show previous messages from tmux, if any.
0.939636	:	Enter the tmux command prompt.
1.030039	C-z	Suspend the tmux client.=
1.436135	z	Toggle zoom state of the current pane.

find-command ❯ uv run python search_chroma.py "show tmux key bindings"
Distance	Key	Description
0.765333	?	List all key bindings.
0.794635	~	Show previous messages from tmux, if any.
0.939636	:	Enter the tmux command prompt.
1.030039	C-z	Suspend the tmux client.
1.436135	z	Toggle zoom state of the current pane.

find-command ❯ uv run python search_chroma.py "copy text"
Distance	Key	Description
0.621593	]	Paste the most recently copied buffer of text.
0.795785	[	Enter copy mode to copy text or view the history.
0.807986	-	Delete the most recently copied buffer of text.
1.171978	Page Up	Enter copy mode and scroll one page up.
1.178153	=	Choose which buffer to paste interactively from a list.

find-command ❯ uv run python search_chroma.py "make the current pane smaller"
Distance	Key	Description
0.508896	"	Split the current pane into two, top and bottom.
0.586366	%	Split the current pane into two, left and right.
0.690726	}	Swap the current pane with the next pane.
0.705375	*	Create a new floating pane.
0.713714	C-Up, C-Down, C-Left, C-Right	Resize the current pane in steps of one cell.

find_command master
find-command ❯ uv run python search_chroma.py "change the size of the current pane"
Distance	Key	Description
0.528513	"	Split the current pane into two, top and bottom.
0.587517	%	Split the current pane into two, left and right.
0.635723	}	Swap the current pane with the next pane.
0.648525	{	Swap the current pane with the previous pane.
0.715728	C-Up, C-Down, C-Left, C-Right	Resize the current pane in steps of one cell.

find_command master
find-command ❯ uv run python search_chroma.py "resize pane"
Distance	Key	Description
0.666115	C-Up, C-Down, C-Left, C-Right	Resize the current pane in steps of one cell.
0.756851	M-Up, M-Down, M-Left, M-Right	Resize the current pane in steps of five cells.
0.778547	"	Split the current pane into two, top and bottom.
0.806779	%	Split the current pane into two, left and right.
0.807611	*	Create a new floating pane.