#coding:utf8

from BaseDataMaintenance.chat.ERNIE_utils import *

from BaseDataMaintenance.dataSource.source import getConnect_ots
from BaseDataMaintenance.chat.chatUtil import *

from tablestore import *
from BaseDataMaintenance.common.Utils import getRow_ots,getCurrent_date,timeAdd
from bs4 import BeautifulSoup
import json
import re
import pandas as pd
import time



def get_columns(ots_client,docid,columns):

    bool_query = BoolQuery(must_queries=[TermQuery("docid",docid)])
    rows,next_token,total_count,is_all_succeed = ots_client.search("document","document_index",
                                                                   SearchQuery(bool_query),
                                                                   ColumnsToGet(columns,return_type=ColumnReturnType.SPECIFIED))
    list_data = getRow_ots(rows)
    _dict = {}
    if len(list_data)==1:
        _dict = list_data[0]
    return _dict


def jsonchat(msg,try_times):

    try:
        print(msg)
    except Exception as e:
        pass
    while try_times>0:
        try:
            try_times -= 1
            resp = chat(msg)
            time.sleep(1)

            if resp.status_code == 200:
                result_dict = json.loads(resp.text)
                result = result_dict.get("result", "")
                error_msg = result_dict.get("error_msg")
                if error_msg is not None:
                    print("error_msg",error_msg)
                    time.sleep(10)
                    continue
                _pattern = "```json(?P<json>.*)```"
                _search = re.search(_pattern, result, re.DOTALL)
                if _search is not None:
                    _json = _search.groupdict().get("json")
                    _d = json.loads(_json)
                    return _json
        except Exception as e:
            pass


def extract_tenderee():
    filename = r'F:\Workspace2016\DataMining\data\2024-11-26_174430_数据导出.xlsx'
    df = pd.read_excel(filename)

    ots_client = getConnect_ots()

    list_data = []

    for docid in df["docid"]:
        docid = int(docid)
        # if docid!=559799502:
        #     continue
        _dict = get_columns(ots_client,docid,["doctextcon","attachmenttextcon","nlp_enterprise","nlp_enterprise_attachment"])
        doctextcon = _dict.get("doctextcon","")
        attachmenttextcon = _dict.get("attachmenttextcon","")
        nlp_enterprise = _dict.get("nlp_enterprise","")
        nlp_enterprise_attachment = _dict.get("nlp_enterprise_attachment","")

        pre_tenderee = ""
        if len(nlp_enterprise)>2:
            _ent = json.loads(nlp_enterprise)
            pre_tenderee = _ent[0]
        if len(nlp_enterprise_attachment)>2:
            _ent = json.loads(nlp_enterprise_attachment)
            pre_tenderee = _ent[0]

        msg = '''从内容中提取出招标人，招标人应该是公司实体，如果没有则返回"",返回结果为json格式{"tenderee":""}\n%s\n%s''' % (str(doctextcon),str(attachmenttextcon))
        _json = jsonchat(msg,3)
        new_tenderee = ""
        if _json is not None:
            _d = json.loads(_json)
            new_tenderee = _d.get("tenderee")
        new_d = {"docid":docid,"nlp_enterprise":nlp_enterprise,"nlp_enterprise_attachment":nlp_enterprise_attachment,
                 "pre_tenderee":pre_tenderee,"new_tenderee":new_tenderee}
        list_data.append(new_d)
        print(new_d)
    df1 = pd.DataFrame(list_data)
    df1.to_excel("tenderee_extract.xlsx",columns=["docid","nlp_enterprise","nlp_enterprise_attachment","pre_tenderee","new_tenderee"])

def prompt_tenderee():
    _prompt = '招标人，招标人应该是公司实体，如果没有则返回""'
    _ret = {"招标人":""}
    return _prompt,_ret

def prompt_budget():
    _prompt = "预算金额，如果没有则默认0"
    _ret = {"预算金额":0}
    return _prompt,_ret

def prompt_win_tenderer():
    _prompt = '中标人及其中标金额，中标人应该是公司实体，中标金额没有则默认0，中标人与中标金额放到一个字典中，如果有多个，则在数组中分别返回，如果没有则返回空数组'
    _ret = {"中标人及金额":[{"中标人":"","中标金额":0}]}
    return _prompt,_ret

def extract_bidding_budget():
    pass

def extract_win_tenderer():
    pass


def get_data_to_qualify(ots_client,count=-1):
    current_date = getCurrent_date('%Y-%m-%d')
    last_date = timeAdd(current_date,-1)
    bool_query = BoolQuery(
        must_queries=[
            RangeQuery("crtime",last_date,current_date),
            RangeQuery("status",201,301),
            TermsQuery("docchannel",[52,101,119,120])
        ]
    )

    list_data = []
    rows,next_token,total_count,is_all_succeed = ots_client.search("document","document_index",
                                                                   SearchQuery(bool_query,limit=100,get_total_count=True),
                                                                   ColumnsToGet(["extract_json"],return_type=ColumnReturnType.SPECIFIED))
    list_data.extend(getRow_ots(rows))
    while 1:
        if next_token is None or len(list_data)>=30*10000:
            break
        rows,next_token,total_count,is_all_succeed = ots_client.search("document","document_index",
                                                                   SearchQuery(bool_query,next_token=next_token,limit=100,get_total_count=True),
                                                                   ColumnsToGet(["extract_json"],return_type=ColumnReturnType.SPECIFIED))
        list_data.extend(getRow_ots(rows))
        if count>0 and list_data>=count:
            break
    return list_data


def quality_inspection():


def merge_extract_json():


if __name__ == '__main__':
    extract_tenderee()







