import os

from pypaimon import CatalogFactory, Schema
import maxframe.dataframe as md
import maxframe
import pyarrow as pa
from odps import ODPS

if __name__ == '__main__':
    # 配置 ODPS entry
    # 优先从环境变量读取，如果没有则尝试从全局获取，最后才手动创建
    odps_entry = (
        ODPS.from_environments() or 
        ODPS.from_global() or 
        ODPS(
            access_id=os.environ['DLF_ACCESS_KEY_ID'],
            secret_access_key=os.environ['DLF_ACCESS_KEY_SECRET'],
            project='mc2dlf_test',  # 替换为实际的 project 名称
            endpoint='http://service.cn-hangzhou.maxcompute.aliyun.com/api'
        )
    )
    
    if odps_entry is None:
        raise ValueError("无法获取 ODPS entry，请配置 ODPS 环境变量或手动创建 ODPS 对象")
    
    # 初始化 maxframe session（必须在创建 DataFrame 之前）
    maxframe.new_session(odps_entry=odps_entry)
    
    # Paimon catalog 配置
    options = {
        'metastore': 'rest',
        'uri': f"http://pre1-test-bennett.aliyun-inc.com",
        'warehouse': "ray_poc",
        '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"
    }
    rest_catalog = CatalogFactory.create(options)

    # 建表（如果不存在）
    table_name = 'demo.bitdemo'
    pa_schema = pa.schema([
        ('userid', pa.string()),
        ('username', pa.string()),
        ('userbyte', pa.binary()),
    ])
    schema = Schema.from_pyarrow_schema(pa_schema, primary_keys=['userid'])
    rest_catalog.create_database('demo', True)
    rest_catalog.drop_table(table_name, True)
    rest_catalog.create_table(table_name, schema, False)
    
    # 使用 maxframe DataFrame
    new_data = md.DataFrame({
        'userid': ['3', '4'],
        'username': ['mf1', 'mf2'],
        'userbyte': [b'true', b'false'],
    })
    table = rest_catalog.get_table('demo.bitdemo')
    write_builder = table.new_batch_write_builder()
    table_write = write_builder.new_write()
    table_write.write_maxframe(new_data)