from config import simplemodel
from template import xuanxiang
import json
with open("0103basic_rule", "r") as f:
    content = f.read()

rule = json.loads(content)
def aifilter(A, #options
       B, #data
       aiclient, sfclient):
    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",
        messages=[
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": " 处理要求：如果清单描述的工作内容是管桩清理，则去除给定选项中的管桩填芯的选项; 如果清单描述的工作内容不是管桩清理，则不做处理"},
            {"role": "user", "content": "问题描述： 给定一段工程量清单描述: " + B['mc'] + " " + B['tz'] +   "," + "，".join(options) + "。请根据处理要求，处理选项，并返回结果,删除选项必须对应到明确的特殊处理要求，不要擅自删除选项。例如，如果处理完后剩余A,B,C三个选项，请返回[A,B,C]"},
        ],
        extra_body={"thinking": {"type": "disabled"}},
    )
    json_string = completion.choices[0].message.content
    print(json_string)
    completion = sfclient.chat.completions.create(
        model=simplemodel(),
        messages=xuanxiang(json_string),
        extra_body={"thinking": {"type": "disabled"}},
    )
    json_string = completion.choices[0].message.content
    print(json_string)
    answer=[]
    if 'A' in json_string:
        answer.append(A[0])
    if 'B' in json_string:
        answer.append(A[1])
    if 'C' in json_string:
        answer.append(A[2])
    if 'D' in json_string:
        answer.append(A[3])
    if 'E' in json_string:
        answer.append(A[4])
    if 'F' in json_string:
        answer.append(A[5])
    if 'G' in json_string:
        answer.append(A[6])
    if 'H' in json_string:
        answer.append(A[7])
    return answer
def associate_jiezhuang(answer):
    hit = False
    for entry in answer:
        if entry in ['第三章  桩基工程 3.1  打桩工程 3.1.5  电焊接桩 电焊接桩 方桩包角钢','第三章  桩基工程 3.1  打桩工程 3.1.5  电焊接桩 电焊接桩 方桩包钢板','第三章  桩基工程 3.1  打桩工程 3.1.5  电焊接桩 电焊接桩 螺栓+电焊']:
            hit = True
            return answer
    if not hit:
        return answer + ['第三章  桩基工程 3.1  打桩工程 3.1.5  电焊接桩 电焊接桩 螺栓+电焊']

def associate(answer):
    hit = False
    for entry in answer:
        if entry in rule['3.1.1.1']:
            hit = "3.1.1.2"
        elif entry in rule['3.1.2.1']:
            hit = "3.1.2.2"
        elif entry in rule['3.1.3.1']:
            hit = "3.1.3.2"
        elif entry in rule["3.1.4.1"]:
            hit = "3.1.4.2"
    if hit:
        already = False
        for entry in answer:
            if entry in rule[hit]:
                already = True
        if not already:
            return answer + [rule[hit][0]]
        else:
            return answer
    else:
        return answer


def jiezhuang(
       data, #data
       aiclient, sfclient):

    completion = aiclient.chat.completions.create(
        model="glm-4.5-air",
        messages=[
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": " 背景知识：已知预应力高强混凝土管桩（PHC）型号定义为PHC-AAA(BB)CC-DDD-E1,E2,E3,E4,其中AAA代表管桩外径，BB代表管桩壁厚，CC表示型号，DDD表示混凝土强度等级，E1/E2/E3/E4表示分段桩长。例如，PHC-500(125)-AB-C80-9,7 表示外径500mm,壁厚125mm,型号AB，混凝土强度C80, 分段桩长分别为9米和7米，总桩长16米,施工时需要将两个分段接桩;再例如，PHC-500(125)-AB-C80-9 表示外径500mm,壁厚125mm,型号AB，混凝土强度C80, 为一整段桩,长9米, 施工时不需要接桩"},
            {"role": "user", "content": " 背景知识：在施工过程中，若管桩长度不足，则需在原有管桩基础上连接新的管桩，这一过程即为接桩。 通过接桩，可以将多节管桩连续打入地基，直至达到设计深度"},
        {"role": "user", "content": "问题描述：  给定一段工作内容描述，内容为" + data['mc'] +data['tz'] + "。请判断完成该工作内容是否需要接桩。请回答是或者否"},
        ],
        extra_body={"thinking": {"type": "disabled"}},
    )
    json_string = completion.choices[0].message.content
    print(json_string)
    if "是" in json_string and "否" not in json_string:
        return True
    if "是" not in json_string and "否" in json_string:
        return False
    completion = sfclient.chat.completions.create(
        model=simplemodel(),
        messages=[
            {"role": "system", "content": "You are a helpful assistant.请将最终答案以JSON格式输出"},
            {"role": "user", "content": " 给你一段文字如下， " + json_string + ",其中给出了一个是或者否的答案，请将该最终答案输出"},                       
        ],
        extra_body={"thinking": {"type": "disabled"}},
    )
    json_string = completion.choices[0].message.content
    print(json_string)
    if "是" in json_string and "否" not in json_string:
        return True
    if "是" not in json_string and "否" in json_string:
        return False
def songzhuang(
       data, #data
       aiclient, sfclient):

    completion = aiclient.chat.completions.create(
        model="glm-4.5-flash",
        messages=[
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": " 背景知识：送桩 是指当桩顶设计标高低于地面时,需要把桩顶打入到设计标高。"},
        {"role": "user", "content": "问题描述：  给定一段工作内容描述，内容为" + data['mc'] +data['tz'] + "。请判断完成该工作内容是否需要送桩。请回答是或者否"},
        ],
        extra_body={"thinking": {"type": "disabled"}},
    )
    json_string = completion.choices[0].message.content
    print(json_string)
    if "是" in json_string and "否" not in json_string:
        return True
    if "是" not in json_string and "否" in json_string:
        return False
    completion = sfclient.chat.completions.create(
        model=simplemodel(),
        messages=[
            {"role": "system", "content": "You are a helpful assistant.请将最终答案以JSON格式输出"},
            {"role": "user", "content": " 给你一段文字如下， " + json_string + ",其中给出了一个是或者否的答案，请将该最终答案输出"},                       
        ],
        extra_body={"thinking": {"type": "disabled"}},
    )
    json_string = completion.choices[0].message.content
    print(json_string)
    if "是" in json_string and "否" not in json_string:
        return True
    if "是" not in json_string and "否" in json_string:
        return False
def ai(A, #options
       B, #data
       C, #entry
       aiclient, sfclient):
    options=[]
    letters = "ABCDEFGHIJKLMN"
    for i in range(len(A)):
        options.append("给定选项" + letters[i]+"，内容为"+A[i])

    completion = aiclient.chat.completions.create(
        model="glm-4.5-flash",
        messages=[
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": " 背景知识：已知预应力高强混凝土管桩（PHC）型号定义为PHC-AAA(BB)CC-DDD-E1,E2,E3,E4,其中AAA代表管桩外径，BB代表管桩壁厚，CC表示型号，DDD表示混凝土强度等级，E1/E2/E3/E4表示分段桩长。例如，PHC-500(125)-AB-C80-9,7 表示外径500mm,壁厚125mm,型号AB，混凝土强度C80, 分段桩长分别为9米和7米，总桩长16米"},
        {"role": "user", "content": "问题描述：  给定一段工作内容描述，内容为" + B['mc'] +B['tz'] + "，" + "，".join(options) + "。请从上述选项中选择（总）桩长最恰当的一个选项，并返回代号,如果无法得出总桩长，则默认桩长10米。例如，如果A选项最恰当，请返回A; "},
        ],
        extra_body={"thinking": {"type": "disabled"}},
    )
    answer2 = completion.choices[0].message.content
    print(answer2)
    if len(answer2) < 4:
        if 'A' in answer2:
            return A[0]
        if 'B' in answer2:
            return A[1]
        if 'C' in answer2:
            return A[2]
        if 'D' in answer2:
            return A[3]
        if 'E' in answer2:
            return A[4]
        if 'F' in answer2:
            return A[5]
        if 'G' in answer2:
            return A[6]
        if 'H' in answer2:
            return A[7]
        if 'I' in answer2:
            return A[8]
        if 'J' in answer2:
            return A[9]

    completion = sfclient.chat.completions.create(
        model=simplemodel(),
        messages=[
            {"role": "system", "content": "You are a helpful assistant.请将最终答案以JSON格式输出"},
            {"role": "user", "content": " 给你一段文字如下， " + answer2 + ",其中给出了一个类似于A或者B或者C的表达式作为答案，请将该最终答案输出"},                       
        ],
        extra_body={"thinking": {"type": "disabled"}},
    )
    json_string = completion.choices[0].message.content
    print(json_string)
    answers = json_string.split("\n")
    answers = [x for x in answers if ':' in x]
    answers = [x for x in answers if 'A' in x or 'B' in x or 'C' in x or 'D' in x]
    print(answers)
    if len(answers) == 0:
        return C
    answer2 = answers[0].split(":")[1].replace(" ", "")
    if 'A' in answer2:
        return A[0]
    if 'B' in answer2:
        return A[1]
    if 'C' in answer2:
        return A[2]
    if 'D' in answer2:
        return A[3]
    if 'E' in answer2:
        return A[4]
    if 'F' in answer2:
        return A[5]
    if 'G' in answer2:
        return A[6]
    if 'H' in answer2:
        return A[7]
    if 'I' in answer2:
        return A[8]
    if 'J' in answer2:
        return A[9]
def select(options, data, entry, aiclient, sfclient):
    if len([x for x in options if '桩长在' in x]) == len(options):
        return ai(options, data, entry, aiclient, sfclient)
    else:
        return entry
def postprocess0103(selected, data, aiclient, sfclient):
    if jiezhuang(data, aiclient, sfclient):
        selected = associate_jiezhuang(selected)
    if songzhuang(data, aiclient, sfclient):
        selected = associate(selected)
    correct=[]
    for entry in selected:
        options = []
        for item in rule:
            l = rule[item]
            if entry in l:
                options = l
        if len(options) > 0:
            correct.append(select(options, data, entry, aiclient, sfclient))
        else:
            correct.append(entry)
    return aifilter(correct, data, aiclient, sfclient)