import pyarrow as pa
import os
from pypaimon.common.file_io import FileIO
from pypaimon.common.options.options import Options
from pypaimon import CatalogFactory, Schema

catalog = CatalogFactory.create({
    'metastore': 'rest',
    'uri': f"https://dlfnext.cn-hangzhou.aliyuncs.com/",
    'warehouse': "morax_test",
    'dlf.region': 'cn-hangzhou',
    "token.provider": "dlf",
    'dlf.access-key-id': os.environ['DLF_ACCESS_KEY_ID'],
    'dlf.access-key-secret': os.environ['DLF_ACCESS_KEY_SECRET'],
    # 'dlf.oss-endpoint': 'oss-cn-hangzhou.aliyuncs.com'
})

pa_schema = pa.schema([
    ('product_id', pa.int64()),
    ('name', pa.string()),
    ('description', pa.string()),
    ('category', pa.string()),
    ('image', pa.large_binary()),
    ('embedding', pa.list_(pa.float32())),
])

schema = Schema.from_pyarrow_schema(
    pa_schema,
    partition_keys=['category'],
    options={
        'row-tracking.enabled': 'true',
        'data-evolution.enabled': 'true',
        'blob-field': 'image',
        'morax.btree-index.enabled': 'true',
        'global-index.btree.index-column': 'product_id',
        'morax.lumina-vector-ann-index.enabled': 'true',
        'global-index.lumina-vector-ann.index-column': 'embedding',
        'lumina.index.dimension': '512',
    }
)

catalog.create_table('default.product_images', schema, True)

table = catalog.get_table('default.product_images')
write_builder = table.new_batch_write_builder()
table_write = write_builder.new_write()
table_commit = write_builder.new_commit()

os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'

oss_io = FileIO.get("oss://your-bucket/product_images/", Options({
    'fs.oss.accessKeyId': os.environ['DLF_ACCESS_KEY_ID'],
    'fs.oss.accessKeySecret': os.environ['DLF_ACCESS_KEY_SECRET'],
    'fs.oss.endpoint': 'oss-cn-hangzhou.aliyuncs.com',
}))

images = []
file_infos = oss_io.list_status("oss://chengli-hz-dlf/blob_test/index_test/")
for info in file_infos:
    if info.type.name == 'File':
        oss_path = info.path if info.path.startswith("oss://") else f"oss://{info.path}"
        with oss_io.new_input_stream(oss_path) as f:
            images.append(f.read())

data = pa.Table.from_pydict({
    'product_id': list(range(len(images))),
    'name': [f'product_{i}' for i in range(len(images))],
    'description': [''] * len(images),
    'category': ['general'] * len(images),
    'image': pa.array(images, type=pa.large_binary()),
    'embedding': pa.array([None] * len(images), type=pa.list_(pa.float32())),
})
table_write.write_arrow(data)
table_commit.commit(table_write.prepare_commit())
table_write.close()
table_commit.close()