import tempfile

import lance
import lancedb
import pyarrow as pa
import pyarrow.compute as pc

print(lance.__version__)
print(lancedb.__version__)

with tempfile.TemporaryDirectory() as tempdir:
    db = lancedb.connect(
        tempdir,
    )

    schema = pa.schema(
        [
            pa.field("vector", pa.list_(pa.float32(), 4), nullable=True),
            pa.field("path", pa.string(), nullable=False),
            pa.field("status", pa.utf8()),
            pa.field("file_size", pa.int64()),
        ]
    )

    # 1. Create table with 1000 rows
    tbl = db.create_table("test", schema=schema)
    records = [{"vector": None, "path": f"/img/{i}.jpg", "status": "pending", "file_size": 1000 + i} for i in range(1000)]
    print(tbl.add(pa.Table.from_pylist(records, schema=schema)))
    print(tbl.count_rows())

    # 2. Create scalar index on the merge key
    tbl.create_scalar_index("path", index_type="BTREE")

    # 3. merge_insert 128 rows
    updates = [{"path": f"/img/{i}.jpg", "vector": [0.1] * 4, "status": "indexed"} for i in range(128)]
    print(tbl.merge_insert("path").when_matched_update_all().execute(pa.Table.from_pylist(updates)))

    arrow = tbl.to_arrow().select(["status"])
    n_indexed = pc.sum(pc.equal(arrow.column("status"), "indexed")).as_py()
    print(f"Indexed: {n_indexed}")  # Expected: 128, Actual: 0

    # 4. Drop index will fix the problem
    tbl.drop_index("path_idx")
    print(tbl.merge_insert("path").when_matched_update_all().execute(pa.Table.from_pylist(updates)))

    arrow = tbl.to_arrow().select(["status"])
    n_indexed = pc.sum(pc.equal(arrow.column("status"), "indexed")).as_py()
    print(f"Indexed: {n_indexed}")  # Expected: 128, Actual: 128