"""Phase-32 verification: SocialGraphV2 Vector Search - Python bindings test.

Covers:
  T32.5: Constructor with vector_dimension + is_vector_search_enabled / vector_count
  T32.6: index_entity / remove_entity_vector / search_entities_by_vector
  T32.7: hybrid_search (vector + text fusion)
  T32.8: End-to-end integration test
"""

import tempfile
import shutil

from aifastdb import SocialGraphV2


def make_embedding(dim: int, base: float) -> list:
    """Create a simple embedding vector."""
    return [base + i * 0.001 for i in range(dim)]


def test_all():
    tmpdir = tempfile.mkdtemp(prefix="aifastdb_vector_test_")
    print(f"Test dir: {tmpdir}")

    try:
        DIM = 4

        # ==================================================================
        # 1. Vector search disabled by default
        # ==================================================================
        print("\n=== 1. Vector search disabled by default ===")
        with SocialGraphV2(tmpdir + "/no-vector", wal_enabled=True) as sg:
            assert sg.is_vector_search_enabled() is False, "Should be disabled by default"
            assert sg.vector_count() == 0, "Should be 0"
            print("  [OK] is_vector_search_enabled=False, vector_count=0")

        # ==================================================================
        # 2. Vector search enabled via constructor
        # ==================================================================
        print("\n=== 2. Enable vector search ===")
        with SocialGraphV2(
            tmpdir + "/with-vector",
            wal_enabled=True,
            vector_dimension=DIM,
        ) as sg:
            assert sg.is_vector_search_enabled() is True
            assert sg.vector_count() == 0
            print("  [OK] is_vector_search_enabled=True, vector_count=0")

            # ==============================================================
            # 3. Index entities with vectors
            # ==============================================================
            print("\n=== 3. Index entities ===")

            alice = sg.add_person({"name": "Alice"})
            bob = sg.add_person({"name": "Bob"})
            charlie = sg.add_person({"name": "Charlie"})

            emb_alice = make_embedding(DIM, 0.1)
            emb_bob = make_embedding(DIM, 0.5)
            emb_charlie = make_embedding(DIM, 0.9)

            sg.index_entity(alice["id"], emb_alice)
            assert sg.vector_count() == 1
            print(f"  [OK] index Alice, vector_count={sg.vector_count()}")

            sg.index_entity(bob["id"], emb_bob)
            sg.index_entity(charlie["id"], emb_charlie)
            assert sg.vector_count() == 3
            print(f"  [OK] index all, vector_count={sg.vector_count()}")

            # ==============================================================
            # 4. indexEntity on non-existent entity should raise
            # ==============================================================
            print("\n=== 4. Error cases ===")
            try:
                sg.index_entity("non-existent-id", make_embedding(DIM, 0.0))
                assert False, "Should have raised"
            except Exception as e:
                print(f"  [OK] index_entity on non-existent entity raised: {e}")

            # ==============================================================
            # 5. search_entities_by_vector - basic k-NN
            # ==============================================================
            print("\n=== 5. search_entities_by_vector (k-NN) ===")

            # Query close to Alice
            query = make_embedding(DIM, 0.11)
            results = sg.search_entities_by_vector(query, 3)
            assert len(results) == 3, f"Expected 3 results, got {len(results)}"
            assert results[0]["entity_id"] == alice["id"], "Closest should be Alice"
            assert isinstance(results[0]["score"], float), "Score should be float"
            assert results[0]["entity"] is not None, "Entity data should be present"
            print(f"  [OK] Top result: {results[0]['entity']['name']} (score={results[0]['score']:.4f})")

            # Query close to Charlie
            query2 = make_embedding(DIM, 0.89)
            results2 = sg.search_entities_by_vector(query2, 1)
            assert len(results2) == 1
            assert results2[0]["entity_id"] == charlie["id"]
            print(f"  [OK] k=1 closest to Charlie: {results2[0]['entity']['name']}")

            # ==============================================================
            # 6. Entity type filtering
            # ==============================================================
            print("\n=== 6. Entity type filtering ===")

            tag = sg.add_tag({
                "name": "Tech",
                "tag_type": "skill",
                "description": "",
                "color": "",
                "parent_tags": [],
                "properties": {},
            })
            sg.index_entity(tag["id"], make_embedding(DIM, 0.12))

            person_results = sg.search_entities_by_vector(query, 10, entity_type_filter="person")
            for r in person_results:
                assert r["entity"]["entity_type"] == "person", \
                    f"Expected person, got {r['entity']['entity_type']}"
            print(f"  [OK] person-only results: {len(person_results)}")

            # ==============================================================
            # 7. remove_entity_vector
            # ==============================================================
            print("\n=== 7. remove_entity_vector ===")

            count_before = sg.vector_count()
            sg.remove_entity_vector(alice["id"])
            results_after = sg.search_entities_by_vector(emb_alice, 10)
            assert all(r["entity_id"] != alice["id"] for r in results_after), \
                "Removed entity should not appear in search"
            print(f"  [OK] vector_count: {count_before} -> {sg.vector_count()}, Alice removed from search")

            # ==============================================================
            # 8. Idempotent re-index
            # ==============================================================
            print("\n=== 8. Idempotent re-index ===")

            sg.index_entity(bob["id"], make_embedding(DIM, 0.6))
            reindex_results = sg.search_entities_by_vector(make_embedding(DIM, 0.6), 1)
            assert len(reindex_results) >= 1
            print("  [OK] Re-index Bob is idempotent, still searchable")

            # ==============================================================
            # 9. hybrid_search
            # ==============================================================
            print("\n=== 9. hybrid_search ===")

            # Vector + text
            hybrid = sg.hybrid_search(5, embedding=emb_bob, text_query="Bob")
            assert len(hybrid) > 0, "hybrid_search should return results"
            print(f"  [OK] hybrid (vec+text): {len(hybrid)} results")
            for r in hybrid[:3]:
                print(f"      {r['entity_id']}: score={r['score']:.4f}, name={r['entity'].get('name')}")

            # Text-only
            text_only = sg.hybrid_search(5, text_query="Charlie")
            assert len(text_only) > 0
            print(f"  [OK] text-only: {len(text_only)} results")

            # Vector-only
            vec_only = sg.hybrid_search(5, embedding=emb_charlie)
            assert len(vec_only) > 0
            print(f"  [OK] vector-only: {len(vec_only)} results")

        # ==================================================================
        # 10. WAL recovery for vectors
        # ==================================================================
        print("\n=== 10. WAL recovery ===")
        with SocialGraphV2(
            tmpdir + "/with-vector",
            wal_enabled=True,
            vector_dimension=DIM,
        ) as sg2:
            sg2.recover()
            count_after_recovery = sg2.vector_count()
            print(f"  vector_count after recovery: {count_after_recovery}")
            if count_after_recovery > 0:
                print("  [OK] Vector WAL recovery: vectors persisted")
            else:
                print("  [WARN] Vector WAL recovery: 0 vectors (may depend on flush timing)")

        print()
        print("=" * 60)
        print("All Phase-32 Python vector search tests PASSED!")
        print("=" * 60)

    finally:
        shutil.rmtree(tmpdir, ignore_errors=True)


if __name__ == "__main__":
    test_all()
