"""Unit tests for Federation Python bindings (T27.2 – T27.7 / T27.8).

Tests:
- T27.2: Federation __init__, context manager, __repr__
- T27.3: register_graph, register_document, unregister, list_stores
- T27.4: query, cross_query
- T27.5: write, batch_write
- T27.6: begin_transaction, commit, rollback
- T27.7: write_with_auth, query_with_auth, get_record_actions
"""

import os
import sys
import tempfile
import shutil

# Ensure the Python package is importable
pkg_dir = os.path.join(os.path.dirname(__file__), "..", "python")
if pkg_dir not in sys.path:
    sys.path.insert(0, pkg_dir)

from aifastdb import Federation, DocumentStore


def make_temp_dir(name="aifastdb_fed_test"):
    """Create a temporary directory for test data."""
    d = tempfile.mkdtemp(prefix=name + "_")
    return d


def cleanup(path):
    if os.path.exists(path):
        shutil.rmtree(path, ignore_errors=True)


# ========================================================================
# T27.2: Federation construction & lifecycle
# ========================================================================

def test_federation_init():
    """Test Federation construction and repr."""
    with Federation() as fed:
        r = repr(fed)
        assert "Federation" in r
        assert "stores=0" in r

    print("✓ test_federation_init passed")


def test_federation_with_permissions():
    """Test creating Federation with permissions enabled."""
    with Federation(enable_permissions=True, query_timeout_ms=5000) as fed:
        r = repr(fed)
        assert "Federation" in r

    print("✓ test_federation_with_permissions passed")


# ========================================================================
# T27.3: Store registration
# ========================================================================

def test_register_document_store():
    """Test registering and listing document stores."""
    tmp = make_temp_dir("fed_doc")
    try:
        with Federation() as fed:
            fed.register_document("docs", os.path.join(tmp, "docs.jsonl"))

            stores = fed.list_stores()
            assert len(stores) == 1
            alias, stype = stores[0]
            assert alias == "docs"
            assert stype == "document"

        print("✓ test_register_document_store passed")
    finally:
        cleanup(tmp)


def test_register_graph_store():
    """Test registering a graph store."""
    tmp = make_temp_dir("fed_graph")
    try:
        with Federation() as fed:
            fed.register_graph("social", tmp, shard_count=2, wal_enabled=False)

            stores = fed.list_stores()
            assert len(stores) == 1
            alias, stype = stores[0]
            assert alias == "social"
            assert stype == "graph"

        print("✓ test_register_graph_store passed")
    finally:
        cleanup(tmp)


def test_register_multiple_and_unregister():
    """Test registering multiple stores and unregistering."""
    tmp = make_temp_dir("fed_multi")
    try:
        with Federation() as fed:
            fed.register_document("docs", os.path.join(tmp, "docs.jsonl"))
            fed.register_graph("social", os.path.join(tmp, "social"), wal_enabled=False)

            stores = fed.list_stores()
            assert len(stores) == 2

            fed.unregister("docs")
            stores = fed.list_stores()
            assert len(stores) == 1
            assert stores[0][0] == "social"

        print("✓ test_register_multiple_and_unregister passed")
    finally:
        cleanup(tmp)


# ========================================================================
# T27.4: Unified query
# ========================================================================

def test_document_query():
    """Test querying a document store through Federation."""
    tmp = make_temp_dir("fed_query")
    try:
        with Federation() as fed:
            doc_path = os.path.join(tmp, "docs.jsonl")
            fed.register_document("docs", doc_path)

            # Write some documents first
            fed.write("docs", {
                "op": "CreateDocument",
                "id": "q1",
                "doc_type": "text",
                "content": {"text": "First document"},
                "tags": ["test"],
            })
            fed.write("docs", {
                "op": "CreateDocument",
                "id": "q2",
                "doc_type": "code",
                "content": {"text": "print('hello')"},
                "tags": ["python"],
            })
            fed.write("docs", {
                "op": "CreateDocument",
                "id": "q3",
                "doc_type": "text",
                "content": {"text": "Third document"},
                "tags": ["test"],
            })

            # Query all text documents
            result = fed.query("docs", {
                "target": "text",
            })
            assert "records" in result
            # Should find q1 and q3 (both text type)
            assert len(result["records"]) == 2, f"Expected 2 text docs, got {len(result['records'])}"

            # Query with limit
            result = fed.query("docs", {
                "target": "*",
                "limit": 2,
            })
            assert len(result["records"]) == 2

        print("✓ test_document_query passed")
    finally:
        cleanup(tmp)


def test_graph_query():
    """Test querying a graph store through Federation."""
    tmp = make_temp_dir("fed_gquery")
    try:
        with Federation() as fed:
            fed.register_graph("social", os.path.join(tmp, "social"), shard_count=2, wal_enabled=False)

            # Create some entities
            fed.write("social", {
                "op": "CreateEntity",
                "id": "p1",
                "entity_type": "person",
                "name": "Alice",
                "attributes": {"age": 30},
            })
            fed.write("social", {
                "op": "CreateEntity",
                "id": "p2",
                "entity_type": "person",
                "name": "Bob",
                "attributes": {"age": 25},
            })

            # Query all person entities
            result = fed.query("social", {
                "target": "person",
            })
            assert "records" in result
            assert len(result["records"]) == 2, f"Expected 2 persons, got {len(result['records'])}"

        print("✓ test_graph_query passed")
    finally:
        cleanup(tmp)


# ========================================================================
# T27.5: Unified write
# ========================================================================

def test_document_write():
    """Test write and batch_write on document store."""
    tmp = make_temp_dir("fed_write")
    try:
        with Federation() as fed:
            doc_path = os.path.join(tmp, "docs.jsonl")
            fed.register_document("docs", doc_path)

            # Single write
            result = fed.write("docs", {
                "op": "CreateDocument",
                "id": "w1",
                "doc_type": "text",
                "content": {"text": "Write test"},
                "tags": ["write"],
            })
            assert "Created" in result, f"Expected 'Created' in result, got: {result}"

            # Update
            result = fed.write("docs", {
                "op": "UpdateDocument",
                "id": "w1",
                "content": {"text": "Updated content"},
            })
            assert "Updated" in result

            # Delete
            result = fed.write("docs", {
                "op": "DeleteDocument",
                "id": "w1",
            })
            assert "Deleted" in result

        print("✓ test_document_write passed")
    finally:
        cleanup(tmp)


def test_batch_write():
    """Test batch_write on document store."""
    tmp = make_temp_dir("fed_batch")
    try:
        with Federation() as fed:
            doc_path = os.path.join(tmp, "docs.jsonl")
            fed.register_document("docs", doc_path)

            ops = [
                {
                    "op": "CreateDocument",
                    "id": f"b{i}",
                    "doc_type": "text",
                    "content": {"text": f"Batch doc {i}"},
                    "tags": ["batch"],
                }
                for i in range(5)
            ]

            result = fed.batch_write("docs", ops)
            assert result["success_count"] == 5, f"Expected 5 successes, got {result['success_count']}"
            assert result["failed_count"] == 0

        print("✓ test_batch_write passed")
    finally:
        cleanup(tmp)


def test_entity_write():
    """Test entity CreateEntity/UpdateEntity/DeleteEntity through Federation."""
    tmp = make_temp_dir("fed_entity")
    try:
        with Federation() as fed:
            fed.register_graph("social", os.path.join(tmp, "social"), shard_count=2, wal_enabled=False)

            # Create
            result = fed.write("social", {
                "op": "CreateEntity",
                "id": "e1",
                "entity_type": "person",
                "name": "Charlie",
                "attributes": {"city": "Shanghai"},
            })
            assert "Created" in result

            # Update
            result = fed.write("social", {
                "op": "UpdateEntity",
                "id": "e1",
                "name": "Charlie Updated",
            })
            assert "Updated" in result

            # Delete
            result = fed.write("social", {
                "op": "DeleteEntity",
                "id": "e1",
            })
            assert "Deleted" in result

        print("✓ test_entity_write passed")
    finally:
        cleanup(tmp)


def test_relation_write():
    """Test relation CreateRelation/DeleteRelation through Federation."""
    tmp = make_temp_dir("fed_rel")
    try:
        with Federation() as fed:
            fed.register_graph("social", os.path.join(tmp, "social"), shard_count=2, wal_enabled=False)

            # Create two entities first
            fed.write("social", {
                "op": "CreateEntity",
                "id": "r1",
                "entity_type": "person",
                "name": "Alice",
                "attributes": {},
            })
            fed.write("social", {
                "op": "CreateEntity",
                "id": "r2",
                "entity_type": "person",
                "name": "Bob",
                "attributes": {},
            })

            # Create relation
            result = fed.write("social", {
                "op": "CreateRelation",
                "source_id": "r1",
                "target_id": "r2",
                "relation_type": "friend",
                "attributes": {},
            })
            assert "Created" in result

        print("✓ test_relation_write passed")
    finally:
        cleanup(tmp)


# ========================================================================
# T27.6: Transaction support
# ========================================================================

def test_transaction_commit():
    """Test begin_transaction + commit."""
    tmp = make_temp_dir("fed_tx")
    try:
        with Federation() as fed:
            doc_path = os.path.join(tmp, "docs.jsonl")
            fed.register_document("docs", doc_path)

            tx = fed.begin_transaction(["docs"])
            tx.write("docs", {
                "op": "CreateDocument",
                "id": "tx1",
                "doc_type": "text",
                "content": {"text": "Transaction doc"},
                "tags": ["tx"],
            })
            tx.commit()

            # Verify data was committed
            result = fed.query("docs", {"target": "*"})
            ids = {r["id"] for r in result["records"]}
            assert "tx1" in ids, f"tx1 should be committed; found {ids}"

        print("✓ test_transaction_commit passed")
    finally:
        cleanup(tmp)


def test_transaction_rollback():
    """Test begin_transaction + rollback."""
    tmp = make_temp_dir("fed_rb")
    try:
        with Federation() as fed:
            doc_path = os.path.join(tmp, "docs.jsonl")
            fed.register_document("docs", doc_path)

            # First create a doc normally
            fed.write("docs", {
                "op": "CreateDocument",
                "id": "rb_base",
                "doc_type": "text",
                "content": {"text": "Base doc"},
                "tags": [],
            })

            tx = fed.begin_transaction(["docs"])
            tx.write("docs", {
                "op": "CreateDocument",
                "id": "rb_tx",
                "doc_type": "text",
                "content": {"text": "Should be rolled back"},
                "tags": [],
            })
            tx.rollback()

            # Verify rb_tx was NOT committed
            result = fed.query("docs", {"target": "*"})
            ids = {r["id"] for r in result["records"]}
            assert "rb_tx" not in ids, f"rb_tx should not exist after rollback; found {ids}"
            assert "rb_base" in ids, "rb_base should still exist"

        print("✓ test_transaction_rollback passed")
    finally:
        cleanup(tmp)


def test_transaction_double_use():
    """Test that using a consumed transaction raises an error."""
    tmp = make_temp_dir("fed_dbl")
    try:
        with Federation() as fed:
            doc_path = os.path.join(tmp, "docs.jsonl")
            fed.register_document("docs", doc_path)

            tx = fed.begin_transaction(["docs"])
            tx.commit()

            # Try to commit again — should raise
            try:
                tx.commit()
                assert False, "Should have raised RuntimeError"
            except RuntimeError as e:
                assert "consumed" in str(e).lower()

        print("✓ test_transaction_double_use passed")
    finally:
        cleanup(tmp)


# ========================================================================
# T27.7: Permission-checked operations
# ========================================================================

def test_write_with_auth_no_permissions():
    """Test write_with_auth when permissions are disabled (should pass through)."""
    tmp = make_temp_dir("fed_auth")
    try:
        with Federation(enable_permissions=False) as fed:
            doc_path = os.path.join(tmp, "docs.jsonl")
            fed.register_document("docs", doc_path)

            result = fed.write_with_auth(
                "docs",
                {
                    "op": "CreateDocument",
                    "id": "auth1",
                    "doc_type": "text",
                    "content": {"text": "Auth doc"},
                    "tags": [],
                },
                caller="user1",
                roles=["reader"],
            )
            assert "Created" in result

        print("✓ test_write_with_auth_no_permissions passed")
    finally:
        cleanup(tmp)


def test_query_with_auth():
    """Test query_with_auth."""
    tmp = make_temp_dir("fed_qauth")
    try:
        with Federation(enable_permissions=False) as fed:
            doc_path = os.path.join(tmp, "docs.jsonl")
            fed.register_document("docs", doc_path)

            # Write some docs
            fed.write("docs", {
                "op": "CreateDocument",
                "id": "qa1",
                "doc_type": "text",
                "content": {"text": "Public doc"},
                "tags": [],
            })

            result = fed.query_with_auth(
                "docs",
                {"target": "*"},
                caller="user1",
                roles=["reader"],
            )
            assert "records" in result
            assert len(result["records"]) >= 1

        print("✓ test_query_with_auth passed")
    finally:
        cleanup(tmp)


# ========================================================================
# Integration: Federation + DocumentStore
# ========================================================================

def test_federation_and_standalone_doc_store():
    """Test that Federation's document store and standalone DocumentStore can coexist."""
    tmp = make_temp_dir("fed_integ")
    try:
        doc_path = os.path.join(tmp, "standalone.jsonl")

        # Write via standalone (disable file lock for same-process re-open)
        with DocumentStore(doc_path, enable_file_lock=False) as store:
            store.put("s1", "Standalone doc", tags=["standalone"])
            store.flush()

        import time; time.sleep(0.5)  # wait for background writer

        # Register the same path in Federation and query
        with Federation() as fed:
            fed.register_document("docs", doc_path, enable_file_lock=False)

            result = fed.query("docs", {"target": "*"})
            ids = {r["id"] for r in result["records"]}
            assert "s1" in ids, "Should find standalone doc via Federation"

        print("✓ test_federation_and_standalone_doc_store passed")
    finally:
        cleanup(tmp)


# ============================================================================
# Phase-30: Federation register_vector & register_concurrent_document
# ============================================================================

def test_federation_register_vector():
    """T30.7: register_vector() registers a ConcurrentHNSWIndex-backed store."""
    with Federation() as fed:
        fed.register_vector("vecs", dimension=32, shard_count=2)
        stores = fed.list_stores()
        # list_stores returns [(alias, store_type), ...]
        found = [(a, t) for a, t in stores if a == "vecs"]
        assert len(found) == 1
        assert found[0][1] == "vector"
        print("  ✓ register_vector OK")


def test_federation_register_concurrent_document():
    """T30.7: register_concurrent_document() registers a high-concurrency doc store."""
    tmp = make_temp_dir("fed_cdoc")
    try:
        cds_path = os.path.join(tmp, "concurrent_docs")
        with Federation() as fed:
            fed.register_concurrent_document("fast_docs", cds_path,
                                             shard_count=2, enable_wal=False)
            stores = fed.list_stores()
            # list_stores returns [(alias, store_type), ...]
            found = [(a, t) for a, t in stores if a == "fast_docs"]
            assert len(found) == 1
            assert found[0][1] == "document"
            print("  ✓ register_concurrent_document OK")
    finally:
        shutil.rmtree(tmp, ignore_errors=True)


def test_federation_vector_write_and_query():
    """T30.7: Write and query vectors through Federation."""
    import time

    with Federation() as fed:
        fed.register_vector("vecs", dimension=8, shard_count=1, max_elements=1000)

        # Write a vector document
        fed.write("vecs", {
            "op": "UpsertVector",
            "id": "v1",
            "vector": [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
            "metadata": {"label": "unit_x"},
        })
        fed.write("vecs", {
            "op": "UpsertVector",
            "id": "v2",
            "vector": [0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
            "metadata": {"label": "unit_y"},
        })

        time.sleep(0.2)

        # Query with vector search
        result = fed.query("vecs", {
            "target": "vector",
            "vector_search": {
                "query_vector": [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
                "top_k": 2,
            },
        })
        records = result.get("records", [])
        assert len(records) >= 1
        # The closest should be v1
        ids = [r["id"] for r in records]
        assert "v1" in ids
        print(f"  ✓ vector write & query OK  (found {len(records)} records)")


def test_federation_concurrent_doc_write_and_query():
    """T30.7: Write and query documents through Federation via concurrent store."""
    import time

    tmp = make_temp_dir("fed_cdoc_wq")
    try:
        cds_path = os.path.join(tmp, "concurrent_docs")
        with Federation() as fed:
            fed.register_concurrent_document("fast_docs", cds_path,
                                             shard_count=2, enable_wal=False)

            # Write documents (ConcurrentDocStoreAdapter supports CreateDocument)
            fed.write("fast_docs", {
                "op": "CreateDocument",
                "id": "fd1",
                "doc_type": "text",
                "content": {"text": "Fast document one"},
                "tags": ["fast"],
            })
            fed.write("fast_docs", {
                "op": "CreateDocument",
                "id": "fd2",
                "doc_type": "text",
                "content": {"text": "Fast document two"},
                "tags": ["fast", "second"],
            })

            time.sleep(0.5)

            # Query all documents
            result = fed.query("fast_docs", {"target": "*"})
            records = result.get("records", [])
            assert len(records) >= 1, f"Expected >=1 records, got {len(records)}"
            ids = [r["id"] for r in records]
            assert "fd1" in ids or "fd2" in ids, f"Expected fd1 or fd2 in {ids}"
            print(f"  ✓ concurrent doc write & query OK  (found {len(records)} records)")
    finally:
        shutil.rmtree(tmp, ignore_errors=True)


if __name__ == "__main__":
    # T27.2
    test_federation_init()
    test_federation_with_permissions()

    # T27.3
    test_register_document_store()
    test_register_graph_store()
    test_register_multiple_and_unregister()

    # T27.4
    test_document_query()
    test_graph_query()

    # T27.5
    test_document_write()
    test_batch_write()
    test_entity_write()
    test_relation_write()

    # T27.6
    test_transaction_commit()
    test_transaction_rollback()
    test_transaction_double_use()

    # T27.7
    test_write_with_auth_no_permissions()
    test_query_with_auth()

    # Integration
    test_federation_and_standalone_doc_store()

    # Phase-30: High-concurrency stores via Federation
    test_federation_register_vector()
    test_federation_register_concurrent_document()
    test_federation_vector_write_and_query()
    test_federation_concurrent_doc_write_and_query()

    print("\n✅ All Federation tests passed!")
