import json
from config import simplemodel
from template import xuanxiang
with open("cuoshi_corpus.txt", "r") as f:
    content = f.read()
pairs = content.split("\n")
with open("name_label", "r") as f:
    content = f.read()
label = json.loads(content)
pair_tuple = []
i = 0
while i < len(pairs):
    pair_tuple.append([pairs[i], pairs[i+1]])
    i = i + 2
chaogao = [label['19-' + str(i+1)][30:] for i in range(18)]
jiangshui = [
  [label['22-11'], label['22-12'], label['22-13']],
  [label['22-14'], label['22-15'], label['22-16']],
  [label['22-17'], label['22-19'], label['22-20']]
]
def handle_jichu(
       A,
       B, #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": "问题描述： 设备基础一般可分成块体设备基础(记作A)和框架设备基础(记作B)。给定一段工作内容: " + B['label'] + " " +  B['mc'] + " " + B['tz'] +   "," + "请判断描述的设备基础属于A还是B，并返回结果。如果无法确定，则返回A。"},
        ],
        extra_body={"thinking": {"type": "disabled"}},
    )
    json_string = completion.choices[0].message.content
    print(json_string)
    answer = 'A'
    if len(json_string) < 5:
        if 'B' in json_string :
            answer = 'B'
    else:
        completion = sfclient.chat.completions.create(
        model=simplemodel(),
        messages=[
            {"role": "system", "content": "You are a helpful assistant.请将结果以json格式输出"},
            {"role": "user", "content": "问题描述： 给定一段文字: " + json_string + "。文字给出了一个类似于A或者B的英文字母作为答案，请返回这个英文字母结果,不要输出分析过程"},
        ],
        extra_body={"thinking": {"type": "disabled"}},
        )
        json_string = completion.choices[0].message.content
        print(json_string)
        answer='A'
        if 'B' in json_string :
            answer = 'B'
    if answer == 'A':
        return [label['21-14']]
    else:
        return [label['21-18']]
def handle_water(
       A,
       B, #data
       aiclient, sfclient):
    hit = False
    hit_group = []
    for entry in A:
        if hit:
            break
        for group in jiangshui:
            if entry in group:
                hit = True
                hit_group = group
                break
    return hit_group
def handle_chaogao(
       B, #data
       aiclient, sfclient):
    options=[]
    letters = "ABCDEFGHIJKLMNOPQRST"
    for i in range(18):
        options.append("给定选项" + letters[i]+"，内容为"+chaogao[i])

    completion = aiclient.chat.completions.create(
        model="glm-4.5-air",
        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"}},
    )
    json_string = completion.choices[0].message.content
    print(json_string)
    answer = 'A'
    if len(json_string) < 8:
        if 'B' in json_string :
            answer = 'B'
        if 'C' in json_string :
            answer = 'C'
        if 'D' in json_string :
            answer = 'D'
        if 'E' in json_string :
            answer = 'E'
        if 'F' in json_string :
            answer = 'F'
        if 'G' in json_string :
            answer = 'G'
        if 'H' in json_string :
            answer = 'H'
        if 'I' in json_string :
            answer = 'I'
        if 'J' in json_string :
            answer = 'J'
        if 'K' in json_string :
            answer = 'K'
        if 'L' in json_string :
            answer = 'L'
        if 'M' in json_string :
            answer = 'M'
        if 'N' in json_string :
            answer = 'N'
        if 'O' in json_string :
            answer = 'O'
        if 'P' in json_string :
            answer = 'P'
        if 'Q' in json_string :
            answer = 'Q'
        if 'R' in json_string :
            answer = 'R'
    else:
        completion = sfclient.chat.completions.create(
        model=simplemodel(),
        messages=[
            {"role": "system", "content": "You are a helpful assistant.请将结果以json格式输出"},
            {"role": "user", "content": "问题描述： 给定一段文字: " + json_string + "。文字给出了一个类似于A或者B的英文字母作为答案，请返回这个英文字母结果,不要输出分析过程"},
        ],
        extra_body={"thinking": {"type": "disabled"}},
        )
        json_string = completion.choices[0].message.content
        print(json_string)
        answer='A'
        if 'B' in json_string :
            answer = 'B'
        if 'C' in json_string :
            answer = 'C'
        if 'D' in json_string :
            answer = 'D'
        if 'E' in json_string :
            answer = 'E'
        if 'F' in json_string :
            answer = 'F'
        if 'G' in json_string :
            answer = 'G'
        if 'H' in json_string :
            answer = 'H'
        if 'I' in json_string :
            answer = 'I'
        if 'J' in json_string :
            answer = 'J'
        if 'K' in json_string :
            answer = 'K'
        if 'L' in json_string :
            answer = 'L'
        if 'M' in json_string :
            answer = 'M'
        if 'N' in json_string :
            answer = 'N'
        if 'O' in json_string :
            answer = 'O'
        if 'P' in json_string :
            answer = 'P'
        if 'Q' in json_string :
            answer = 'Q'
        if 'R' in json_string :
            answer = 'R'
    if answer == 'A':
        return [label['19-1']]
    if answer == 'B':
        return [label['19-2']]
    if answer == 'C':
        return [label['19-3']]
    if answer == 'D':
        return [label['19-4']]
    if answer == 'E':
        return [label['19-5']]
    if answer == 'F':
        return [label['19-6']]
    if answer == 'G':
        return [label['19-7']]
    if answer == 'H':
        return [label['19-8']]
    if answer == 'I':
        return [label['19-9']]
    if answer == 'J':
        return [label['19-10']]
    if answer == 'K':
        return [label['19-11']]
    if answer == 'L':
        return [label['19-12']]
    if answer == 'M':
        return [label['19-13']]
    if answer == 'N':
        return [label['19-14']]
    if answer == 'O':
        return [label['19-15']]
    if answer == 'P':
        return [label['19-16']]
    if answer == 'Q':
        return [label['19-17']]
    if answer == 'R':
        return [label['19-18']] 
    return [label['19-1']]
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['label'] + " " +  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)
    if len(json_string) < 8:
        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])
        return answer
    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 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])
    return answer
def postprocess0117(selected, data, aiclient, sfclient):
    water = 0
    non_water = 0
    for entry in selected:
        if '施工排水' in entry:
            water = water + 1
        else:
            non_water = non_water + 1
    if water > non_water:
        return handle_water(selected, data, aiclient, sfclient)
    jichu = 0
    non_jichu = 0
    for entry in selected:
        if '设备基础' in entry:
            jichu = jichu + 1
        else:
            non_jichu = non_jichu + 1
    if jichu > non_jichu:
        return handle_jichu(selected, data, aiclient, sfclient)

    if '超高增加' in data['mc'] or '超高增加' in data['tz'] or '超高施工增加' in data['mc'] or '超高施工增加' in data['tz']:
        return handle_chaogao(data, aiclient, sfclient)
    correct=[]
    for entry in selected:
        correct.append(entry)
        for item in pair_tuple:
            if entry in item:
                correct = correct + item
    correct = list(set(correct))
    return aifilter(correct, data, aiclient, sfclient)