#!/usr/bin/env python3
"""
Command-line semantic search using a pre-built index.
"""

import argparse
import pickle
from pathlib import Path
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity
from pygeai import Client


def load_index(index_path: Path):
    """Load search index from disk."""
    with open(index_path, "rb") as f:
        return pickle.load(f)


def search(query: str, index_data: dict, client: Client, model: str, top_k: int = 5):
    """Search for similar documents."""
    
    # Generate query embedding
    query_emb = client.embeddings.create(input=query, model=model).data[0].embedding
    
    # Calculate similarities
    similarities = cosine_similarity([query_emb], index_data["embeddings"])[0]
    
    # Get top results
    top_indices = np.argsort(similarities)[::-1][:top_k]
    
    results = []
    for idx in top_indices:
        results.append({
            "rank": len(results) + 1,
            "document": index_data["documents"][idx],
            "score": float(similarities[idx])
        })
    
    return results


def main():
    parser = argparse.ArgumentParser(description="Semantic search CLI")
    parser.add_argument("index", type=Path, help="Search index file (.pkl)")
    parser.add_argument("query", help="Search query")
    parser.add_argument("--model", default="text-embedding-3-small", help="Embedding model")
    parser.add_argument("--top-k", type=int, default=5, help="Number of results")
    parser.add_argument("--json", action="store_true", help="Output as JSON")
    
    args = parser.parse_args()
    
    # Load index
    print(f"📂 Loading index from {args.index}...")
    index_data = load_index(args.index)
    print(f"   {index_data['count']} documents in index\n")
    
    # Search
    print(f"🔍 Searching for: \"{args.query}\"\n")
    client = Client()
    results = search(args.query, index_data, client, args.model, args.top_k)
    
    # Display results
    if args.json:
        import json
        print(json.dumps(results, indent=2))
    else:
        print("Results:")
        print("=" * 80)
        for result in results:
            print(f"\n{result['rank']}. Score: {result['score']:.4f}")
            print(f"   {result['document']}")
        print("\n" + "=" * 80)


if __name__ == "__main__":
    main()
