from chromadb import Search, K, Knn, And

def search_products(collection, user_query, min_price=None, max_price=None,
                   category=None, in_stock_only=True, page=0, page_size=20):
    """
    Search for products with semantic search and filters.

    Args:
        collection: Chroma collection
        user_query: Natural language search query (e.g., "wireless headphones")
        min_price: Minimum price filter
        max_price: Maximum price filter
        category: Product category filter
        in_stock_only: Only show in-stock items
        page: Page number (0-indexed)
        page_size: Results per page
    """

    # Build filter conditions
    from chromadb import And

    combined_filter = And([])

    if in_stock_only:
        combined_filter &= K("in_stock") == True

    if category:
        combined_filter &= K("category") == category

    if min_price is not None:
        combined_filter &= K("price") >= min_price

    if max_price is not None:
        combined_filter &= K("price") <= max_price

    # Build search
    search = Search().where(combined_filter)

    search = (search
        .rank(Knn(query=user_query))
        .limit(page_size, offset=page * page_size)
        .select(K.DOCUMENT, K.SCORE, "name", "price", "category", "rating", "image_url"))

    # Execute search
    results = collection.search(search)
    rows = results.rows()[0]

    # Format results for display
    products = []
    for row in rows:
        products.append({
            "id": row["id"],
            "name": row["metadata"]["name"],
            "description": row["document"][:200] + "...",
            "price": row["metadata"]["price"],
            "category": row["metadata"]["category"],
            "rating": row["metadata"]["rating"],
            "image_url": row["metadata"]["image_url"],
            "relevance_score": row["score"]
        })

    return products

# Example usage
products = search_products(
    collection,
    user_query="noise cancelling headphones for travel",
    min_price=50,
    max_price=300,
    category="electronics",
    page=0,
    page_size=20
)

for i, product in enumerate(products, 1):
    print(f"{i}. {product['name']}")
    print(f"   Price: ${product['price']:.2f} | Rating: {product['rating']}/5")
    print(f"   {product['description']}")
    print(f"   Relevance: {product['relevance_score']:.3f}")
    print()

from chromadb import Search, K, Knn, Rrf

def get_recommendations(collection, user_id, user_preferences,
                       seen_content_ids, num_recommendations=10):
    """
    Get personalized content recommendations for a user.

    Args:
        collection: Chroma collection
        user_id: User identifier
        user_preferences: Dict with user interests and preferences
        seen_content_ids: List of content IDs the user has already seen
        num_recommendations: Number of recommendations to return
    """

    # Build filter to exclude seen content and match preferences
    combined_filter = K.ID.not_in(seen_content_ids)

    # Filter by preferred categories
    if user_preferences.get("categories"):
        combined_filter &= K("category").is_in(user_preferences["categories"])

    # Filter by language preference
    if user_preferences.get("language"):
        combined_filter &= K("language") == user_preferences["language"]

    # Filter by minimum rating
    min_rating = user_preferences.get("min_rating", 3.5)
    combined_filter &= K("rating") >= min_rating

    # Only show published content
    combined_filter &= K("status") == "published"

    # Create hybrid search combining multiple signals
    # Signal 1: User interest embedding
    user_interest_query = " ".join(user_preferences.get("interests", ["general"]))

    # Signal 2: Similar to user's favorite content
    favorite_topics_query = " ".join(user_preferences.get("favorite_topics", []))

    # Use RRF to combine both signals
    hybrid_rank = Rrf(
        ranks=[
            Knn(query=user_interest_query, return_rank=True, limit=200),
            Knn(query=favorite_topics_query, return_rank=True, limit=200)
        ],
        weights=[0.6, 0.4],  # User interests weighted higher
        k=60
    )

    search = (Search()
        .where(combined_filter)
        .rank(hybrid_rank)
        .limit(num_recommendations)
        .select(K.DOCUMENT, K.SCORE, "title", "category", "author",
                "rating", "published_date", "thumbnail_url"))

    results = collection.search(search)
    rows = results.rows()[0]

    # Format recommendations
    recommendations = []
    for row in rows:
        recommendations.append({
            "id": row["id"],
            "title": row["metadata"]["title"],
            "description": row["document"][:150] + "...",
            "category": row["metadata"]["category"],
            "author": row["metadata"]["author"],
            "rating": row["metadata"]["rating"],
            "published_date": row["metadata"]["published_date"],
            "thumbnail_url": row["metadata"]["thumbnail_url"],
            "relevance_score": row["score"]
        })

    return recommendations

# Example usage
user_preferences = {
    "interests": ["machine learning", "artificial intelligence", "data science"],
    "favorite_topics": ["neural networks", "deep learning", "transformers"],
    "categories": ["technology", "science", "research"],
    "language": "en",
    "min_rating": 4.0
}

seen_content = ["content_001", "content_045", "content_123"]

recommendations = get_recommendations(
    collection,
    user_id="user_42",
    user_preferences=user_preferences,
    seen_content_ids=seen_content,
    num_recommendations=10
)

print("Personalized Recommendations:")
for i, rec in enumerate(recommendations, 1):
    print(f"\n{i}. {rec['title']}")
    print(f"   Category: {rec['category']} | Author: {rec['author']}")
    print(f"   Rating: {rec['rating']}/5 | Published: {rec['published_date']}")
    print(f"   {rec['description']}")
    print(f"   Match Score: {rec['relevance_score']:.3f}")

from chromadb import Search, K, Knn

def search_across_categories(collection, user_query, categories, results_per_category=5):
    """
    Search across multiple categories in parallel using batch operations.

    Args:
        collection: Chroma collection
        user_query: User's search query
        categories: List of categories to search
        results_per_category: Number of results per category
    """

    # Build a search for each category
    searches = []
    for category in categories:
        search = (Search()
            .where(K("category") == category)
            .rank(Knn(query=user_query))
            .limit(results_per_category)
            .select(K.DOCUMENT, K.SCORE, "title", "category", "date"))
        searches.append(search)

    # Execute all searches in one batch
    results = collection.search(searches)

    # Process results by category
    category_results = {}
    for i, category in enumerate(categories):
        rows = results.rows()[i]
        category_results[category] = [
            {
                "id": row["id"],
                "title": row["metadata"]["title"],
                "description": row["document"][:100] + "...",
                "date": row["metadata"]["date"],
                "score": row["score"]
            }
            for row in rows
        ]

    return category_results

# Example usage
query = "latest developments in renewable energy"
categories = ["technology", "science", "news", "research"]

results_by_category = search_across_categories(
    collection,
    user_query=query,
    categories=categories,
    results_per_category=3
)

# Display results
for category, results in results_by_category.items():
    print(f"\n{'='*60}")
    print(f"Category: {category.upper()}")
    print('='*60)

    if not results:
        print("  No results found")
        continue

    for i, result in enumerate(results, 1):
        print(f"\n  {i}. {result['title']}")
        print(f"     Date: {result['date']}")
        print(f"     {result['description']}")
        print(f"     Relevance: {result['score']:.3f}")