#!/usr/bin/env python3
"""
Persistence Test for CyborgDB Python SDK

This script tests data persistence across service restarts:
- Phase 1 (setup): Creates index, inserts data, saves index info to file
- Phase 2 (verify): Loads index, verifies data still exists

Usage:
  python persistence_test_python.py setup    # Run before restart
  python persistence_test_python.py verify   # Run after restart
"""

import os
import sys
import json
import base64
import numpy as np

from cyborgdb import Client, IndexIVFFlat


STATE_FILE = "/tmp/persistence_test_state.json"
VECTOR_COUNT = 1000
VECTOR_DIM = 128


def setup_phase():
    """Create index and insert test data"""
    print("=" * 60)
    print("Persistence Test - Setup Phase")
    print("=" * 60)

    base_url = os.environ.get("CYBORGDB_URL", "http://localhost:8000")
    api_key = os.environ.get("CYBORGDB_API_KEY")

    if not api_key:
        print("ERROR: CYBORGDB_API_KEY not set")
        return 1

    # Initialize client
    print("\n[1/4] Initializing client...")
    client = Client(base_url=base_url, api_key=api_key)

    # Generate key and create index
    print("[2/4] Creating index...")
    index_key = Client.generate_key()
    index_name = "persistence_test_index"

    # Delete existing index if present
    try:
        existing = client.load_index(index_name, index_key)
        existing.delete_index()
        print(f"  Deleted existing index '{index_name}'")
    except Exception:
        pass

    index_config = IndexIVFFlat(dimension=VECTOR_DIM)
    index = client.create_index(
        index_name=index_name,
        index_key=index_key,
        index_config=index_config,
        metric="euclidean",
    )
    print(f"  Created index '{index_name}'")

    # Insert test vectors with known values
    print(f"[3/4] Inserting {VECTOR_COUNT} test vectors...")

    # Use seeded random for reproducible vectors
    np.random.seed(42)
    vectors = np.random.randn(VECTOR_COUNT, VECTOR_DIM).astype(np.float32)
    ids = [f"persist_vec_{i}" for i in range(VECTOR_COUNT)]

    index.upsert(ids, vectors)
    print(f"  Inserted {VECTOR_COUNT} vectors")

    # Verify data was inserted
    stored_ids = index.list_ids()
    print(f"  Verified {len(stored_ids)} vectors in index")

    # Save state for verification phase
    print("[4/4] Saving state...")
    state = {
        "index_name": index_name,
        "index_key": base64.b64encode(index_key).decode("utf-8"),
        "vector_count": VECTOR_COUNT,
        "vector_dim": VECTOR_DIM,
        "test_ids": ids[:10],  # Save first 10 IDs for verification
    }

    with open(STATE_FILE, "w") as f:
        json.dump(state, f)
    print(f"  State saved to {STATE_FILE}")

    print("\n" + "=" * 60)
    print("Setup phase completed successfully!")
    print("Service can now be restarted.")
    print("=" * 60)
    return 0


def verify_phase():
    """Verify data persisted after restart"""
    print("=" * 60)
    print("Persistence Test - Verify Phase")
    print("=" * 60)

    base_url = os.environ.get("CYBORGDB_URL", "http://localhost:8000")
    api_key = os.environ.get("CYBORGDB_API_KEY")

    if not api_key:
        print("ERROR: CYBORGDB_API_KEY not set")
        return 1

    # Load state from setup phase
    print("\n[1/4] Loading state...")
    if not os.path.exists(STATE_FILE):
        print(f"ERROR: State file not found: {STATE_FILE}")
        print("Did you run the setup phase first?")
        return 1

    with open(STATE_FILE, "r") as f:
        state = json.load(f)

    index_name = state["index_name"]
    index_key = base64.b64decode(state["index_key"])
    expected_count = state["vector_count"]
    test_ids = state["test_ids"]

    print(f"  Index: {index_name}")
    print(f"  Expected vectors: {expected_count}")

    # Initialize client
    print("[2/4] Initializing client...")
    client = Client(base_url=base_url, api_key=api_key)

    # Load existing index
    print("[3/4] Loading index...")
    try:
        index = client.load_index(index_name, index_key)
        print(f"  Index '{index_name}' loaded successfully")
    except Exception as e:
        print(f"ERROR: Failed to load index: {e}")
        return 1

    # Verify data
    print("[4/4] Verifying data...")

    # Check vector count
    stored_ids = index.list_ids()
    actual_count = len(stored_ids)
    print(f"  Found {actual_count} vectors (expected {expected_count})")

    if actual_count != expected_count:
        print("ERROR: Vector count mismatch!")
        return 1

    # Verify specific vectors exist
    try:
        items = index.get(test_ids)
        found_ids = [item["id"] for item in items]
        missing = set(test_ids) - set(found_ids)

        if missing:
            print(f"ERROR: Missing vectors: {missing}")
            return 1

        print(f"  Verified {len(test_ids)} test vectors exist")
    except Exception as e:
        print(f"ERROR: Failed to get test vectors: {e}")
        return 1

    # Test query still works
    np.random.seed(42)
    query_vector = np.random.randn(state["vector_dim"]).astype(np.float32)
    results = index.query(query_vectors=query_vector, top_k=10)

    if not results or len(results) == 0:
        print("ERROR: Query returned no results")
        return 1

    print(f"  Query returned {len(results)} results")

    # Cleanup
    print("\nCleaning up...")
    index.delete_index()
    os.remove(STATE_FILE)
    print("  Deleted index and state file")

    print("\n" + "=" * 60)
    print("PERSISTENCE TEST PASSED!")
    print("Data successfully persisted across service restart.")
    print("=" * 60)
    return 0


def main():
    if len(sys.argv) < 2:
        print("Usage: python persistence_test_python.py <setup|verify>")
        return 1

    phase = sys.argv[1].lower()

    if phase == "setup":
        return setup_phase()
    elif phase == "verify":
        return verify_phase()
    else:
        print(f"Unknown phase: {phase}")
        print("Usage: python persistence_test_python.py <setup|verify>")
        return 1


if __name__ == "__main__":
    sys.exit(main())
