#!/usr/bin/env python3
"""
Cleanup script to remove all test data from Neo4j graph database.
This prepares the system for production by removing:
- SearchQuery nodes (HybridRAG search history)
- WebResult nodes (external web search results)
- Any orphaned nodes without relationships
"""
import requests
import sys

# Configuration
BASE_URL = "http://192.168.86.149:8089"
API_KEY = "af88ed8f44bed81bdb20d0534f1c4547340b29e2aba4963f61a71b993d7eb6e5"
USER = "jpmschweitzer"

headers = {
    "Authorization": f"Bearer {API_KEY}",
    "Content-Type": "application/json"
}

def execute_query(query_text, description):
    """Execute a Cypher query via API."""
    url = f"{BASE_URL}/graph/query"
    payload = {
        "user": USER,
        "query": query_text
    }

    print(f"🔄 {description}...")
    response = requests.post(url, headers=headers, json=payload)

    if response.status_code != 200:
        print(f"   ❌ Failed: {response.status_code}")
        print(f"   {response.text}")
        return None

    data = response.json()
    return data

def count_nodes(node_type):
    """Count nodes of a specific type."""
    if node_type == "SearchQuery":
        query = f'MATCH (n:SearchQuery) WHERE n.user = "{USER}" RETURN count(n) as count'
    elif node_type == "WebResult":
        query = f'MATCH (n:User_{USER}_WebResult:WebResult) RETURN count(n) as count'
    elif node_type == "Document":
        query = f'MATCH (n:User_{USER}_Document:Document) RETURN count(n) as count'
    else:
        query = f'MATCH (n:User_{USER}) WHERE NOT n:Document AND NOT n:SearchQuery AND NOT n:WebResult RETURN count(n) as count'

    result = execute_query(query, f"Counting {node_type} nodes")
    if result and result.get("results"):
        return result["results"][0]["count"]
    return 0

def delete_search_queries():
    """Delete all SearchQuery nodes and their relationships."""
    query = f'''
    MATCH (sq:SearchQuery)
    WHERE sq.user = "{USER}"
    DETACH DELETE sq
    RETURN count(*) as deleted
    '''

    result = execute_query(query, "Deleting SearchQuery nodes")
    if result and result.get("results"):
        return result["results"][0].get("deleted", 0)
    return 0

def delete_web_results():
    """Delete all WebResult nodes."""
    query = f'''
    MATCH (wr:User_{USER}_WebResult:WebResult)
    DETACH DELETE wr
    RETURN count(*) as deleted
    '''

    result = execute_query(query, "Deleting WebResult nodes")
    if result and result.get("results"):
        return result["results"][0].get("deleted", 0)
    return 0

def delete_orphaned_nodes():
    """Delete any orphaned entity nodes without relationships."""
    query = f'''
    MATCH (n:User_{USER})
    WHERE NOT n:Document
      AND NOT n:SearchQuery
      AND NOT n:WebResult
      AND NOT (n)--()
    DETACH DELETE n
    RETURN count(*) as deleted
    '''

    result = execute_query(query, "Deleting orphaned nodes")
    if result and result.get("results"):
        return result["results"][0].get("deleted", 0)
    return 0

def get_all_node_stats():
    """Get statistics on all node types."""
    query = f'''
    MATCH (n:User_{USER})
    WITH labels(n) as labels, count(n) as count
    RETURN labels, count
    ORDER BY count DESC
    '''

    result = execute_query(query, "Getting node statistics")
    if result and result.get("results"):
        return result["results"]
    return []

def main():
    print("=" * 70)
    print("🧹 Neo4j Graph Database Cleanup Script")
    print("=" * 70)
    print()

    # Show current state
    print("📊 Current Database State:")
    print("-" * 70)

    search_count = count_nodes("SearchQuery")
    web_count = count_nodes("WebResult")
    doc_count = count_nodes("Document")
    other_count = count_nodes("Other")

    print(f"   SearchQuery nodes:  {search_count}")
    print(f"   WebResult nodes:    {web_count}")
    print(f"   Document nodes:     {doc_count}")
    print(f"   Other entity nodes: {other_count}")
    print()

    if search_count == 0 and web_count == 0 and doc_count == 0 and other_count == 0:
        print("✨ Database is already clean!")
        return

    total_to_delete = search_count + web_count + other_count
    if total_to_delete == 0 and doc_count > 0:
        print("⚠️  Only Document nodes found (legitimate data)")
        print("   No cleanup needed!")
        return

    # Ask for confirmation
    print("⚠️  WARNING: About to delete:")
    if search_count > 0:
        print(f"   - {search_count} SearchQuery nodes (search history)")
    if web_count > 0:
        print(f"   - {web_count} WebResult nodes (cached web results)")
    if other_count > 0:
        print(f"   - {other_count} orphaned entity nodes")

    print()
    print("   Document nodes will NOT be deleted (they represent wiki pages)")
    print()

    response = input("Continue with cleanup? (yes/no): ").strip().lower()
    if response != "yes":
        print("❌ Cancelled")
        return

    print()
    print("🚀 Starting cleanup...")
    print()

    # Delete SearchQuery nodes
    if search_count > 0:
        deleted_sq = delete_search_queries()
        print(f"   ✅ Deleted {deleted_sq} SearchQuery nodes")

    # Delete WebResult nodes
    if web_count > 0:
        deleted_wr = delete_web_results()
        print(f"   ✅ Deleted {deleted_wr} WebResult nodes")

    # Delete orphaned nodes
    if other_count > 0:
        deleted_orphans = delete_orphaned_nodes()
        print(f"   ✅ Deleted {deleted_orphans} orphaned nodes")

    print()
    print("=" * 70)
    print("📊 Final Database State")
    print("=" * 70)

    # Show final state
    search_count_final = count_nodes("SearchQuery")
    web_count_final = count_nodes("WebResult")
    doc_count_final = count_nodes("Document")
    other_count_final = count_nodes("Other")

    print(f"   SearchQuery nodes:  {search_count_final}")
    print(f"   WebResult nodes:    {web_count_final}")
    print(f"   Document nodes:     {doc_count_final}")
    print(f"   Other entity nodes: {other_count_final}")
    print()

    if search_count_final == 0 and web_count_final == 0 and other_count_final == 0:
        print("✨ Graph database is now production-ready!")
        print("   Only legitimate wiki-derived data remains.")
    else:
        print("⚠️  Some nodes remain - this might be expected if you have active wiki pages")

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n\n❌ Interrupted by user")
        sys.exit(1)
    except Exception as e:
        print(f"\n\n❌ Error: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)