import chromadb
from sentence_transformers import SentenceTransformer
import os

# Initialize
model = SentenceTransformer('all-MiniLM-L6-v2')
chroma_client = chromadb.Client()
collection = chroma_client.get_or_create_collection("lumio_support")

# Company data
documents = [
    # Products
    "Product: Lumio X1 Wireless Headphones. Price: $79.99. Features: 30hr battery, noise cancellation, Bluetooth 5.0. Colors: Black, White, Navy.",
    "Product: Lumio S3 Smart Speaker. Price: $49.99. Features: 360 sound, voice assistant compatible, WiFi and Bluetooth. Colors: Charcoal, Sand.",
    "Product: Lumio ProWatch 2. Price: $199.99. Features: Heart rate monitor, GPS, 7 day battery, waterproof to 50m. Colors: Silver, Midnight Black.",
    "Product: Lumio Tab Mini. Price: $299.99. Features: 8 inch display, 128GB storage, 12hr battery, stylus support.",
    "Product: Lumio Buds Air. Price: $39.99. Features: True wireless, 6hr battery, IPX4 water resistant, touch controls.",
    "Product: Lumio ChargeHub Pro. Price: $34.99. Features: 6 port USB charging station, 60W total output, surge protection.",
    "Product: Lumio CamDoor. Price: $89.99. Features: 1080p doorbell camera, night vision, two-way audio, motion alerts.",
    "Product: Lumio SleepLight. Price: $59.99. Features: Smart bedside lamp, sunrise alarm, 16 million colors, app controlled.",

    # FAQs
    "Return Policy: Lumio accepts returns within 30 days of purchase. Items must be unused and in original packaging. Refunds are processed within 5-7 business days to the original payment method.",
    "Shipping Policy: Standard shipping takes 3-5 business days and is free on orders over $50. Express shipping (1-2 days) costs $12.99. International shipping available to 30 countries.",
    "Warranty: All Lumio products come with a 1 year manufacturer warranty covering defects. ProWatch 2 and CamDoor come with 2 year warranty. Warranty does not cover physical damage or water damage beyond rated specs.",
    "How to contact support: You can reach Lumio support at support@lumio.com or call 1-800-LUMIO-01 Monday to Friday 9am-6pm EST. Live chat is available on the website 24/7.",
    "Payment methods: Lumio accepts Visa, Mastercard, American Express, PayPal, and Apple Pay. Buy now pay later available through Klarna.",
    "Order cancellation: Orders can be cancelled within 2 hours of placement. After that the order may have already been processed for shipping.",

    # Sample orders
    "Order #LM-10482: Customer Sarah Johnson. Product: Lumio X1 Wireless Headphones. Status: Delivered on February 20 2026. Tracking: FedEx 7489234823.",
    "Order #LM-10483: Customer James Lee. Product: Lumio ProWatch 2 in Midnight Black. Status: Out for delivery, expected March 2 2026. Tracking: UPS 1Z9238471.",
    "Order #LM-10484: Customer Priya Sharma. Product: Lumio Buds Air x2. Status: Processing, estimated ship date March 3 2026.",
    "Order #LM-10485: Customer Mike Torres. Product: Lumio Tab Mini 128GB. Status: Shipped February 28 2026. Tracking: USPS 9400111899.",
    "Order #LM-10486: Customer Emma Wilson. Product: Lumio SleepLight + ChargeHub Pro. Status: Delivered February 25 2026. Tracking: FedEx 7489234900.",
]

# Load into ChromaDB
def initialize_knowledge_base():
    existing = collection.count()
    if existing == 0:
        embeddings = model.encode(documents).tolist()
        ids = [f"doc_{i}" for i in range(len(documents))]
        collection.add(documents=documents, embeddings=embeddings, ids=ids)
        print(f"Knowledge base loaded with {len(documents)} documents.")
    else:
        print(f"Knowledge base already loaded with {existing} documents.")

def search_knowledge_base(query: str, n_results: int = 5) -> str:
    query_embedding = model.encode([query]).tolist()
    results = collection.query(query_embeddings=query_embedding, n_results=n_results)
    docs = results['documents'][0] if results['documents'] else []
    return "\n".join(docs)

initialize_knowledge_base()