import time
from fallback import fallback
from config import simplemodel
from template import xuanxiang
from fallback import fallback
zhijing=['公称直径（mm以内）50', '公称直径（mm以内）80', '公称直径（mm以内）100', '公称直径（mm以内）150' , '公称直径（mm以内）200', '公称直径（mm以内）250', '公称直径（mm以内）300']
def select_zhijing(
       B, #data
       aiclient,
             qwclient,
              sfclient,
             dw):
    options=[]
    letters = "ABCDEFGHIJKLMN"
    for i in range(len(zhijing)):
        options.append("给定选项" + letters[i]+"，内容为"+zhijing[i] )

    completion = aiclient.chat.completions.create(
        model="glm-4.5-air",
        #model="THUDM/GLM-Z1-9B-0414",
        #model="ernie-speed-128k",
        messages=[
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": "问题描述： 给定一段工作内容: " + B['label'] + " " +  B['mc'] + " " + B['tz'] +  "，".join(options) + "。请选择最合适的选项，并返回字母代号，例如，返回 A"},
        ],
        extra_body={"thinking": {"type": "disabled"}},
        #extra_body={"enable_thinking": True},
        #stream=True
    )
    json_string = completion.choices[0].message.content
    print(json_string)
    if len([x for x in json_string if x != ',' and x != '[' and x != ']' and x != ' ' and (x < 'A' or x > 'M')]) < 5:
        answer=[]
        if 'A' in json_string and len(zhijing) > 0:
            answer.append('A')
        if 'B' in json_string and len(zhijing) > 1:
            answer.append('B')
        if 'C' in json_string and len(zhijing) > 2:
            answer.append('C')
        if 'D' in json_string and len(zhijing) > 3:
            answer.append('D')
        if 'E' in json_string and len(zhijing) > 4:
            answer.append('E')
        if 'F' in json_string and len(zhijing) > 5:
            answer.append('F')
        if 'G' in json_string and len(zhijing) > 6:
            answer.append('G')
        return answer
    completion = sfclient.chat.completions.create(
        #model="glm-4.5-flash",
        model=simplemodel(),
        messages=xuanxiang(json_string),
        extra_body={"thinking": {"type": "disabled"}},
        #extra_body={"enable_thinking": False},
    )
    json_string = completion.choices[0].message.content
    print(json_string)
    answer=[]
    if 'A' in json_string and len(zhijing) > 0:
        answer.append('A')
    if 'B' in json_string and len(zhijing) > 1:
        answer.append('B')
    if 'C' in json_string and len(zhijing) > 2:
        answer.append('C')
    if 'D' in json_string and len(zhijing) > 3:
        answer.append('D')
    if 'E' in json_string and len(zhijing) > 4:
        answer.append('E')
    if 'F' in json_string and len(zhijing) > 5:
        answer.append('F')
    if 'G' in json_string and len(zhijing) > 6:
        answer.append('G')
    
    return answer
def aifilter1(A, #options
       B, #data
       aiclient,
             qwclient,
              sfclient,
             dw):
    options=[]
    letters = "ABCDEFGHIJKLMN"
    for i in range(len(A)):
        options.append("给定选项" + letters[i]+"，内容为"+A[i] )

    completion = aiclient.chat.completions.create(
        model="glm-4.5-air",
        #model="THUDM/GLM-Z1-9B-0414",
        #model="ernie-speed-128k",
        messages=[
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": " 特殊处理要求一：如果工作内容描述是配电箱安装，并且安装方式是悬挂而非落地式，则去掉所有基础槽钢、角钢的选项"},
            {"role": "user", "content": " 重要提示：选项指的是给定的A、B、C之类的选项，不是指的工作内容中的可能的1、2、3这样罗列的特征"},
            {"role": "user", "content": " 重要提示：除特殊处理要求提及的内容外，不需考虑选项内容与工作内容是否符合，只需要根据特殊处理要求做出处理"},
            {"role": "user", "content": "问题描述： 给定一段工作内容: " + B['label'] + " " +  B['mc'] + " " + B['tz'] +  "，".join(options) + "。请根据处理要求做出处理，并返回结果, 删除选项必须对应到明确的特殊处理要求，不要擅自删除选项。例如，如果处理完后剩余A,B,C三个选项，请返回[A,B,C]"},
        ],
        extra_body={"thinking": {"type": "disabled"}},
        #extra_body={"enable_thinking": True},
        #stream=True
    )
    json_string = completion.choices[0].message.content
    print(json_string)
    if len([x for x in json_string if x != ',' and x != '[' and x != ']' and x != ' ' and (x < 'A' or x > 'M')]) < 5:
        answer=[]
        if 'A' in json_string and len(A) > 0:
            answer.append(A[0])
        if 'B' in json_string and len(A) > 1:
            answer.append(A[1])
        if 'C' in json_string and len(A) > 2:
            answer.append(A[2])
        if 'D' in json_string and len(A) > 3:
            answer.append(A[3])
        if 'E' in json_string and len(A) > 4:
            answer.append(A[4])
        if 'F' in json_string and len(A) > 5:
            answer.append(A[5])
        if 'G' in json_string and len(A) > 6:
            answer.append(A[6])
        if 'H' in json_string and len(A) > 7:
            answer.append(A[7])
        if 'I' in json_string and len(A) > 8:
            answer.append(A[8])
        if 'J' in json_string and len(A) > 9:
            answer.append(A[9])
        return answer
    completion = sfclient.chat.completions.create(
        #model="glm-4.5-flash",
        model=simplemodel(),
        messages=xuanxiang(json_string),
        extra_body={"thinking": {"type": "disabled"}},
        #extra_body={"enable_thinking": False},
    )
    json_string = completion.choices[0].message.content
    print(json_string)
    answer=[]
    if 'A' in json_string and len(A) > 0:
        answer.append(A[0])
    if 'B' in json_string and len(A) > 1:
        answer.append(A[1])
    if 'C' in json_string and len(A) > 2:
        answer.append(A[2])
    if 'D' in json_string and len(A) > 3:
        answer.append(A[3])
    if 'E' in json_string and len(A) > 4:
        answer.append(A[4])
    if 'F' in json_string and len(A) > 5:
        answer.append(A[5])
    if 'G' in json_string and len(A) > 6:
        answer.append(A[6])
    if 'H' in json_string and len(A) > 7:
        answer.append(A[7])
    if 'I' in json_string and len(A) > 8:
        answer.append(A[8])
    if 'J' in json_string and len(A) > 9:
        answer.append(A[9])
    return answer

def postprocess0304(selected, data, aiclient, qwclient, sfclient, label_name, name_dw, candidates):
    prime = aifilter1(selected, data, aiclient, qwclient, sfclient, name_dw)
    if data['bianma'].startswith('030411001'):#配管
        if '翼环' in data['tz']:
            if len([x for x in prime if '刚性防水套管' in x]) == 0:
                t = select_zhijing(data, aiclient, qwclient, sfclient, name_dw)
                if 'A' in t:
                    prime.append('第十册 给排水、采暖、燃气工程 第二章 支架及其它 2.3 刚性防水套管制作、安装 刚性防水套管制作、安装 公称直径(mm以内) 50')
                elif 'B' in t:
                    prime.append('第十册 给排水、采暖、燃气工程 第二章 支架及其它 2.3 刚性防水套管制作、安装 刚性防水套管制作、安装 公称直径(mm以内) 80')
                elif 'C' in t:
                    prime.append('第十册 给排水、采暖、燃气工程 第二章 支架及其它 2.3 刚性防水套管制作、安装 刚性防水套管制作、安装 公称直径(mm以内) 100')
                elif 'D' in t:
                    prime.append('第十册 给排水、采暖、燃气工程 第二章 支架及其它 2.3 刚性防水套管制作、安装 刚性防水套管制作、安装 公称直径(mm以内) 150')
                elif 'E' in t:
                    prime.append('第十册 给排水、采暖、燃气工程 第二章 支架及其它 2.3 刚性防水套管制作、安装 刚性防水套管制作、安装 公称直径(mm以内) 200')
                elif 'F' in t:
                    prime.append('第十册 给排水、采暖、燃气工程 第二章 支架及其它 2.3 刚性防水套管制作、安装 刚性防水套管制作、安装 公称直径(mm以内) 250')
                elif 'G' in t:
                    prime.append('第十册 给排水、采暖、燃气工程 第二章 支架及其它 2.3 刚性防水套管制作、安装 刚性防水套管制作、安装 公称直径(mm以内) 300')

    return prime
