import time
import json
import random

import pymochow
import logging
from pymochow.configuration import Configuration
from pymochow.auth.bce_credentials import BceCredentials
from pymochow.exception import ClientError, ServerError
from pymochow.model.schema import Schema, Field, SecondaryIndex, VectorIndex, HNSWParams
from pymochow.model.enum import FieldType, IndexType, MetricType, ServerErrCode
from pymochow.model.enum import TableState, IndexState
from pymochow.model.table import Partition, Row, AnnSearch, HNSWSearchParams

logging.basicConfig(filename='documentinsight.log', level=logging.DEBUG,
                format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

if __name__ == "__main__":
    account = '$您的账户名称'
    api_key = '$您的账户API密钥'
    endpoint = '$您的实例访问端点' # 例如：'http://127.0.0.1:5287'

    db_name = 'DocumentInsight'
    table_name = 'AutoIncrementTestTable'

    # 根据配置创建一个MochowClient对象
    config = Configuration(credentials=BceCredentials(account, api_key), endpoint=endpoint)
    mochow_client = pymochow.MochowClient(config)

    # 建库
    db = mochow_client.create_database(db_name);
    # 若库已存在，则直接使用下述代码获取一个库对象
    # db = mochow_client.database(db_name)

    # 建表，该表有2个分区
    fields = []
    fields.append(Field("DocId", FieldType.UINT64, primary_key=True, partition_key=True, auto_increment=True, not_null=True))
    fields.append(Field("Title", FieldType.STRING, not_null=True))
    table = db.create_table(
        table_name=table_name,
        replication=1,
        partition=Partition(partition_num=2), # 2个分区
        schema=Schema(fields=fields, indexes=[])
    )
    while True:
        time.sleep(1)
        table = db.describe_table(table_name)
        if table.state == TableState.NORMAL:
            break

    # 插入几条数据，无需指定主键，即DocId的值，该字段的值由数据库自动分配
    # 也因为没有指定主键，这些数据作为一个批次会被随机写入该表的某个分区中
    rows = [
        Row(Title="LLM技术详解"),
        Row(Title="深入浅出大语言模型技术"),
        Row(Title="基于大模型的RAG应用开发实践"),
        Row(Title="AI原生应用开发"),
        Row(Title="向量数据库在RAG开发中的应用与实践")
    ]
    table.insert(rows=rows)

    # 标量查询，分别查询两个分区的第1条数据，两个查询会返回1条数据
    primaryKey = {}
    primaryKey['DocId'] = 1  # 查询第0个分区（分区ID=0）的第一条数据
    results = table.query(primary_key=primaryKey)
    print("Query result from tablet 0: {}".format(results))

    primaryKey = {}
    primaryKey['DocId'] = (1 << 40) + 1  # 查询第1个分区（分区ID=1）的第一条数据
    results = table.query(primary_key=primaryKey)
    print("Query result from tablet 1: {}".format(results))

Query result from tablet 0: {metadata:{content__length:u'35',content__type:u'application/json',request_id:u'0fd95c45-00df-4060-9b18-ce231f7a6883'},row:{},code:0,msg:u'Success'}
Query result from tablet 1: {metadata:{content__length:u'82',content__type:u'application/json',request_id:u'bb72d085-da96-437c-8f54-92c9c5c5e6b6'},row:{'DocId': 1099511627777, 'Title': 'LLM技术详解'},code:0,msg:u'Success'}