#!/usr/bin/env python3
"""
Initialize Sofiia Memory Collections
Creates collections for Sofiia agent in Qdrant
"""

import sys
import time
import requests
from typing import List, Dict

# Configuration
QDRANT_URL = "http://localhost:6333"
AGENT_ID = "sofiia"
EMBEDDING_DIMENSION = 1024  # Cohere embed-multilingual-v3.0

# Collections to create
COLLECTIONS = [
    {
        "name": f"{AGENT_ID}_messages",
        "description": "Chat message history for Sofiia",
        "vectors": {
            "size": EMBEDDING_DIMENSION,
            "distance": "Cosine"
        }
    },
    {
        "name": f"{AGENT_ID}_docs",
        "description": "Documents and knowledge base for Sofiia",
        "vectors": {
            "size": EMBEDDING_DIMENSION,
            "distance": "Cosine"
        }
    },
    {
        "name": f"{AGENT_ID}_memory_items",
        "description": "Long-term memory items for Sofiia",
        "vectors": {
            "size": EMBEDDING_DIMENSION,
            "distance": "Cosine"
        }
    },
    {
        "name": f"{AGENT_ID}_user_context",
        "description": "User context and preferences for Sofiia",
        "vectors": {
            "size": EMBEDDING_DIMENSION,
            "distance": "Cosine"
        }
    }
]


def check_qdrant_health() -> bool:
    """Check if Qdrant is running and healthy"""
    try:
        response = requests.get(f"{QDRANT_URL}/healthz", timeout=5)
        return response.status_code == 200
    except Exception as e:
        print(f"❌ Qdrant not accessible: {e}")
        return False


def collection_exists(collection_name: str) -> bool:
    """Check if collection already exists"""
    try:
        response = requests.get(f"{QDRANT_URL}/collections/{collection_name}", timeout=5)
        return response.status_code == 200
    except:
        return False


def create_collection(collection: Dict) -> bool:
    """Create a collection in Qdrant"""
    collection_name = collection["name"]
    
    if collection_exists(collection_name):
        print(f"   ℹ️  Collection '{collection_name}' already exists")
        return True
    
    payload = {
        "vectors": collection["vectors"]
    }
    
    try:
        response = requests.put(
            f"{QDRANT_URL}/collections/{collection_name}",
            json=payload,
            timeout=10
        )
        
        if response.status_code == 200:
            print(f"   ✅ Created collection '{collection_name}'")
            return True
        else:
            print(f"   ❌ Failed to create '{collection_name}': {response.text}")
            return False
    except Exception as e:
        print(f"   ❌ Error creating '{collection_name}': {e}")
        return False


def list_all_collections() -> List[str]:
    """List all collections in Qdrant"""
    try:
        response = requests.get(f"{QDRANT_URL}/collections", timeout=5)
        if response.status_code == 200:
            data = response.json()
            return [c["name"] for c in data["result"]["collections"]]
        return []
    except:
        return []


def main():
    print("🧠 Sofiia Memory Initialization")
    print("=" * 50)
    print()
    
    # Check Qdrant health
    print("🔍 Checking Qdrant connection...")
    if not check_qdrant_health():
        print("❌ Qdrant is not running or not accessible")
        print("   Please start Qdrant first: docker-compose -f docker-compose.memory-node2.yml up -d qdrant-node2")
        sys.exit(1)
    
    print("✅ Qdrant is healthy")
    print()
    
    # Create collections
    print(f"📦 Creating collections for '{AGENT_ID}'...")
    print()
    
    success_count = 0
    for collection in COLLECTIONS:
        if create_collection(collection):
            success_count += 1
    
    print()
    
    # Summary
    print("=" * 50)
    print(f"📊 Summary: {success_count}/{len(COLLECTIONS)} collections ready")
    print()
    
    # List all collections
    print("📋 All collections:")
    all_collections = list_all_collections()
    for coll in sorted(all_collections):
        prefix = "   ✅" if coll.startswith(f"{AGENT_ID}_") else "   •"
        print(f"{prefix} {coll}")
    
    print()
    print("✅ Sofiia memory initialization complete!")
    print()
    print("📝 Next steps:")
    print("   1. Verify collections: curl http://localhost:6333/collections")
    print("   2. Test memory service: curl http://localhost:8000/health")
    print("   3. Start using Sofiia with memory!")
    print()


if __name__ == "__main__":
    main()