#!/usr/bin/env python3
"""
Script to demonstrate querying ChromaDB instance
This shows how to connect and query your ChromaDB database
"""

import os
import chromadb
import json
from pathlib import Path

def connect_to_chromadb():
    """Connect to ChromaDB instance"""
    try:
        # Get connection details from environment or use defaults
        CHROMADB_HOST = os.getenv("CHROMADB_HOST", "example.com")
        CHROMADB_PORT = int(os.getenv("CHROMADB_PORT", "8000"))
        
        print(f"Connecting to ChromaDB at {CHROMADB_HOST}:{CHROMADB_PORT}")
        
        # Create client connection
        client = chromadb.HttpClient(host=CHROMADB_HOST, port=CHROMADB_PORT)
        
        # Test connection
        collections = client.list_collections()
        print(f"Successfully connected! Found {len(collections)} collections:")
        for collection in collections:
            print(f"  - {collection.name}")
            
        return client
        
    except Exception as e:
        print(f"Failed to connect to ChromaDB: {e}")
        return None

def query_facts_collection(client):
    """Query the facts collection"""
    try:
        # Get the facts collection
        facts_collection = client.get_collection("facts")
        
        print("\n=== Querying Facts Collection ===")
        
        # Example queries that would work with your data
        example_queries = [
            "Unrivaled attendance records",
            "Unrivaled league revenue",
            "David Levy on Unrivaled",
            "Fox Business coverage of Unrivaled"
        ]
        
        for i, query in enumerate(example_queries, 1):
            print(f"\n{i}. Query: '{query}'")
            
            # Perform similarity search
            results = facts_collection.query(
                query_texts=[query],
                n_results=2,
                include=["documents", "metadatas", "distances"]
            )
            
            if results['documents'] and len(results['documents'][0]) > 0:
                print("   Results:")
                for j, doc in enumerate(results['documents'][0], 1):
                    print(f"     {j}. {doc[:100]}...")
                    if j >= 2:  # Show only first 2 results
                        break
            else:
                print("   No results found")
                
    except Exception as e:
        print(f"Error querying facts collection: {e}")

def query_articles_collection(client):
    """Query the articles collection"""
    try:
        # Get the articles collection
        articles_collection = client.get_collection("articles")
        
        print("\n=== Querying Articles Collection ===")
        
        # Example query
        query = "Unrivaled women's basketball"
        print(f"Query: '{query}'")
        
        # Perform similarity search
        results = articles_collection.query(
            query_texts=[query],
            n_results=2,
            include=["documents", "metadatas", "distances"]
        )
        
        if results['documents'] and len(results['documents'][0]) > 0:
            print("Results:")
            for j, doc in enumerate(results['documents'][0], 1):
                print(f"  {j}. {doc[:100]}...")
        else:
            print("No results found")
            
    except Exception as e:
        print(f"Error querying articles collection: {e}")

def main():
    """Main function to demonstrate ChromaDB queries"""
    print("=== ChromaDB Query Demonstration ===")
    print("This script shows how to connect and query your ChromaDB instance")
    print()
    
    # Connect to ChromaDB
    client = connect_to_chromadb()
    
    if client:
        print("\n=== Available Collections ===")
        collections = client.list_collections()
        for collection in collections:
            print(f"  - {collection.name} ({collection.count()} items)")
        
        # Query each collection
        query_facts_collection(client)
        query_articles_collection(client)
        
        print("\n=== Query Capabilities ===")
        print("The system supports:")
        print("✓ Semantic similarity search across facts")
        print("✓ Entity-based filtering")
        print("✓ Source-based filtering")
        print("✓ Time-based filtering")
        print("✓ Multi-entity queries")
        print()
        print("All queries use the qwen3:8b embedding model for efficient searching")
        
    else:
        print("Cannot connect to ChromaDB. Please ensure:")
        print("1. ChromaDB server is running")
        print("2. Network connectivity to the server")
        print("3. Correct host/port configuration")

if __name__ == "__main__":
    main()