
from BaseDataMaintenance.dataSource.source import getConnect_ots,is_internal
from BaseDataMaintenance.common.multiThread import MultiThreadHandler
from BaseDataMaintenance.common.Utils import *
from BaseDataMaintenance.model.ots.document import *
from BaseDataMaintenance.model.ots.document_extract2 import *
from tablestore import *
from apscheduler.schedulers.blocking import BlockingScheduler
import requests
from queue import Queue
# from multiprocessing import Queue

class EASEtract():

    def __init__(self):

        if is_internal:
            self.extract_url = "http://1255640119316927.vpc.cn-hangzhou.pai-eas.aliyuncs.com/api/predict/content_extract"
            self.industy_url = "http://1255640119316927.vpc.cn-hangzhou.pai-eas.aliyuncs.com/api/predict/industry_extract"
            self.other_url = "http://1255640119316927.vpc.cn-hangzhou.pai-eas.aliyuncs.com/api/predict/other_extract"
        else:
            self.extract_url = "http://1255640119316927.cn-hangzhou.pai-eas.aliyuncs.com/api/predict/content_extract"
            self.industy_url = "http://1255640119316927.cn-hangzhou.pai-eas.aliyuncs.com/api/predict/industry_extract"
            self.other_url = "http://1255640119316927.cn-hangzhou.pai-eas.aliyuncs.com/api/predict/other_extract"
        self.header = {'Content-Type': 'application/json',"Authorization":"NzZmOWZlMmU2MGY3YmQ4MDBjM2E5MDAyZjhjNjQ0MzZlMmE0NTMwZg=="}

        self.ots_client = getConnect_ots()
        self.task_queue = Queue(3000)
        self.list_docid = []
        self.last_docids = set()
        self.current_docids = set()


    def setInputQueue(self,row):
        _docid = row.get("docid")
        if self.list_docid.count(_docid)>0:
            if self.list_docid.index(_docid)<5000:
                pass
        else:
            self.task_queue.put(row)
        self.list_docid.append(_docid)
        self.list_docid = self.list_docid[:6000]


    def producer(self):
        columns = ["dochtmlcon","page_time","doctitle","docchannel",document_status]

        boolQuery = BoolQuery(must_queries=[
            RangeQuery("status",1,60,True,True)
            # TermQuery("docid",218804448)
            # TermQuery("page_time",'2021-12-20')
            # TermQuery("docid",213518099)
        ])
        res = self.ots_client.search("document","document_index",
                                     SearchQuery(boolQuery,sort=Sort(sorters=[FieldSort("status",SortOrder.ASC),FieldSort("docid",SortOrder.DESC)]),get_total_count=True,limit=100),
                                                                            ColumnsToGet(column_names=columns,return_type=ColumnReturnType.SPECIFIED))

        if isinstance(res,(tuple,list)):
            rows,next_token,total_count,is_all_succeed = res
        else:
            rows,next_token,total_count,is_all_succeed = (res.rows,res.next_token,res.total_count,res.is_all_succeed)
        log("prducer total_count:%d"%total_count)
        dict_rows = getRow_ots(rows)
        for _row in dict_rows:
            _docid = _row.get(document_docid)
            _status = _row.get(document_status)
            if _docid not in self.last_docids or (_status>=1 and _status<=50):
                self.task_queue.put(_row)
            self.current_docids.add(_docid)
            # self.setInputQueue(_row)
        process_count = len(dict_rows)
        while next_token:
            res = self.ots_client.search("document","document_index",
                                         SearchQuery(boolQuery,next_token=next_token,get_total_count=True,limit=100),
                                                                                ColumnsToGet(column_names=columns,return_type=ColumnReturnType.SPECIFIED))
            if isinstance(res,(tuple,list)):
                rows,next_token,total_count,is_all_succeed = res
            else:
                rows,next_token,total_count,is_all_succeed = (res.rows,res.next_token,res.total_count,res.is_all_succeed)
            dict_rows = getRow_ots(rows)
            for _row in dict_rows:
                _docid = _row.get(document_docid)
                if _docid not in self.last_docids or (_status>=1 and _status<=50):
                    self.task_queue.put(_row)
                self.current_docids.add(_docid)
                # self.setInputQueue(_row)
            process_count += len(dict_rows)
            if process_count>1000:
                break
        self.last_docids = self.current_docids
        self.current_docids = set()




    def comsumer(self):
        def _handle(item,result_queue):

            _extract = Document_extract2({})
            _extract.setValue(document_extract2_partitionkey,item.get(document_partitionkey))
            _extract.setValue(document_extract2_docid,item.get(document_docid))
            all_done = 1
            if all_done:
                data = item
                resp = requests.post(self.other_url,json=data,headers=self.header)
                if (resp.status_code >=200 and resp.status_code<=210):
                    _extract.setValue(document_extract2_other_json,resp.content.decode("utf8"),True)
                else:
                    all_done = -1

            data = {}
            for k,v in item.items():
                data[k] = v
            data["timeout"] = 120
            data["doc_id"] = data.get(document_docid)
            data["content"] = data.get(document_dochtmlcon,"")
            if document_dochtmlcon in data:
                data.pop(document_dochtmlcon)
            data["title"] = data.get(document_doctitle,"")
            if all_done:
                resp = requests.post(self.extract_url,json=data,headers=self.header)
                if (resp.status_code >=200 and resp.status_code<=210):
                    _extract.setValue(document_extract2_extract_json,resp.content.decode("utf8"),True)
                else:
                    all_done = -2
            if all_done:
                resp = requests.post(self.industy_url,json=data,headers=self.header)
                if (resp.status_code >=200 and resp.status_code<=210):
                    _extract.setValue(document_extract2_industry_json,resp.content.decode("utf8"),True)
                else:
                    all_done = -3
            # print(_extract.getProperties().get(document_extract2_extract_json))
            # return
            _d = Document(item)
            if all_done==1:
                _extract.setValue(document_extract2_status,random.randint(1,50),True)
                _d.fromEas2Extract()
                _d.update_row(self.ots_client)
                _extract.update_row(self.ots_client)
            else:
                _d.fromEasFailed()
                _d.update_row(self.ots_client)
            log("process docid:%d %s"%(data["doc_id"],str(all_done)))


        mt = MultiThreadHandler(task_queue=self.task_queue,task_handler=_handle,result_queue=None,thread_count=70)
        mt.run()

    def producerAndComsumer(self):
        self.producer()
        self.comsumer()

    def start(self):
        _schduler = BlockingScheduler()
        # _schduler.add_job(self.producer,"cron",second="*/5")
        # _schduler.add_job(self.comsumer,"cron",second="*/5")
        _schduler.add_job(self.producerAndComsumer,"cron",second="*/5")
        _schduler.start()
        print("start_time:%s"%getCurrent_date())

def start_eas_extract():
    _extract = EASEtract()
    _extract.start()
    # _extract.producerAndComsumer()

def rollbackStatus(docids,ots_client=None):
    if ots_client is None:
        ots_client = getConnect_ots()
    for docid in docids:
        partitionkey = docid%500+1
        _dict = {document_extract2_partitionkey:partitionkey,
                 document_extract2_docid:docid,
                 document_extract2_save:1,
                 document_extract2_already_storage:1,
                 document_extract2_status:69}
        _extract2 = Document_extract2(_dict)
        _extract2.update_row(ots_client)
        log("process %d done"%docid)

def rollbackdocuments():
    import pandas as pd
    df = pd.read_csv()

if __name__ == '__main__':
    # start_eas_extract
    rollbackStatus([140463748])


