from tablestore import BoolQuery, TermQuery, Sort, FieldSort, SortOrder, ColumnReturnType, SearchQuery, ColumnsToGet

from BaseDataMaintenance.common.Utils import getRow_ots
from BaseDataMaintenance.dataSource.source import getConnect_ots
from BaseDataMaintenance.maintenance.proposedBuilding.pb_project_production import run_flow, PBProduction
from BaseDataMaintenance.model.ots.designed_project import designed_project


def repair_pb_project(pb_project_id, pb_project_pk):
    # 修改原来错的拟在建项目status为404, update_time为now
    ots_client = getConnect_ots()
    obj = designed_project({'id': pb_project_id, 'partitionkey': pb_project_pk, 'status': '404'})
    obj.update_row(ots_client, delete=0)

    # 把错的拟在建的docids,spids拿回来，用单进程重新跑一遍
    bool_query = BoolQuery(must_queries=[TermQuery("id", pb_project_id)])
    data_list = search_util(ots_client,
                            'designed_project', 'designed_project_index',
                            bool_query, ['docids', 'spids'], 'update_time',
                            limit=1, while_next=0)
    id_list = []
    if data_list:
        docids = data_list[0].get('docids')
        spids = data_list[0].get('spids')
        if docids:
            id_list += docids.split(',')
        if spids:
            id_list += spids.split(',')
    id_list = list(set(id_list))
    print('id_list', id_list)

    # 重跑
    for docid in id_list:
        ds = PBProduction(test_mode=1, update=1)
        ds.test_doc_ids = [docid]
        tmp_data_list = ds.pb_data_start()
        print('len(tmp_data_list)', len(tmp_data_list))
        if tmp_data_list:
            run_flow(tmp_data_list[0], test_mode=1, update=1)

    # 修改update_time，同步document
    # obj = designed_project({'id': _id, 'partitionkey': pk,
    #                         'status': '404', 'update_time': '2024-07-31 14:35:31'})
    # obj.update_row(ots_client, delete=0)
    return


def search_util(ots_client, table, table_index, bool_query, columns, sort_col,
                limit=999999, show_total=0, while_next=1, asc=0):
    if sort_col:
        if asc:
            sort = Sort(sorters=[FieldSort(sort_col, SortOrder.ASC)])
        else:
            sort = Sort(sorters=[FieldSort(sort_col, SortOrder.DESC)])
    else:
        sort = None
    return_type = ColumnReturnType.SPECIFIED
    rows, next_token, total_count, is_all_succeed = ots_client.search(table, table_index,
                                                                      SearchQuery(bool_query, sort=sort,
                                                                                  limit=100,
                                                                                  get_total_count=True),
                                                                      ColumnsToGet(columns, return_type))
    list_data = getRow_ots(rows)
    if show_total:
        print('search total_count', total_count)
    if len(list_data) >= limit:
        print('limit ', limit, len(list_data))
        return list_data[:limit]

    if while_next:
        while next_token and len(list_data) < limit:
            rows, next_token, total_count, is_all_succeed = ots_client.search(table, table_index,
                                                                              SearchQuery(bool_query,
                                                                                          next_token=next_token,
                                                                                          limit=100,
                                                                                          get_total_count=True),
                                                                              ColumnsToGet(columns,
                                                                                           return_type))
            list_data += getRow_ots(rows)
    return list_data


if __name__ == '__main__':
    _id = 4915783839244091393
    pk = 194

    repair_pb_project(_id, pk)

