import json
from config import simplemodel
from template import xuanxiang
def huitianfilter(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-flash",
        messages=[
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": " 背景知识：回填方一般不使用单轮车或者双轮车运输"},
            {"role": "user", "content": " 背景知识：室内回填、房心回填一般不使用压路机、挖掘机、推土机、单轮车、双轮车"},
            {"role": "user", "content": " 背景知识：回填方时，挖掘机挖土与自卸汽车运土需要配套使用，也就是说，如果选择两者之一，另一个也不许选择。如果不选择两者之一，则另一个也不能选择"},
            {"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 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-flash",
        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 ai(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-flash",
        messages=[
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": " 背景知识：平整场地可以使用人工的方法，也可以使用机械的方法，机械的方法更经济"},
        {"role": "user", "content": "问题描述： " +   "，".join(options) + "。请从上述选项中选择一更经济的选项，并返回代号。例如，如果A选项最恰当，请返回A"},
        ],
        extra_body={"thinking": {"type": "disabled"}},
    )
    json_string = completion.choices[0].message.content
    print(json_string)
    if len(json_string) < 4:
        if 'A' in json_string:
            return A[0]
        if 'B' in json_string:
            return A[1]
        if 'C' in json_string:
            return A[2]
        if 'D' in json_string:
            return A[3]
        if 'E' in json_string:
            return A[4]
        if 'F' in json_string:
            return A[5]
        if 'G' in json_string:
            return A[6]
        if 'H' in json_string:
            return A[7]
        if 'I' in json_string:
            return A[8]
        if 'J' in json_string:
            return A[9]

    completion = sfclient.chat.completions.create(
        model=simplemodel(),
        messages=[
            {"role": "system", "content": "You are a helpful assistant.请将最终答案以JSON格式输出"},
            {"role": "user", "content": " 给你一段文字如下， " + json_string + ",其中给出了一个类似于A或者B或者C的表达式作为答案，请将该最终答案输出.例如，文字提到答案为A，请输出A；文字提到答案是B，请输出B"},                       
        ],
        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]
    print(answers)
    if len(answers) == 0:
        return A[0]
    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 tihuan(selected):
    hit = False
    for entry in selected:
        if '自动平地机' in entry:
            hit = True
    if hit:
        left = [x for x in selected if '自动平地机' not in x]
        left.append('第一章  土、石方工程 1.2  机械土、石方 1.2.9  平整场地、碾压 平整场地(厚300mm以内) 推土机(kW以内) 75')
        return left
    else:
        return selected
def postprocess0101(selected, data, aiclient, sfclient):
    if data['bianma'].startswith("010101001"):##平整场地
        if len(selected) > 1:
            return tihuan([ai(selected, data, aiclient, sfclient)])
        else:
            return tihuan(selected)
    elif data['bianma'].startswith("010103"):##回填
        return huitianfilter(selected, data, aiclient, sfclient)
    else:
        return aifilter(selected, data, aiclient, sfclient)
