# -*- coding: utf-8 -*-
"""
@author: yq
@time: 2024/12/18
@desc: 策略流节点解析
"""
import json
import re
import time
from tqdm import tqdm

import pandas as pd
from PIL import Image
from openpyxl import load_workbook

from commom import f_get_save_path, call_llm, f_file_upload, GeneralException
from enums import ResultCodesEnum
from prompt import f_get_prompt_parse_node, f_get_prompt_parse_flow


def _f_parse_flow(ws, node_list: list):
    image = ws._images[0]
    img = Image.open(image.ref).convert("RGB")
    save_path = f_get_save_path("流程图.png")
    img.save(save_path)
    time.sleep(1)
    file_id = f_file_upload(save_path)
    prompt = f_get_prompt_parse_flow(node_list)
    print(prompt)
    prompt = [
        {
            "type": "text",
            "text": prompt
        },
        {
            "type": "image",
            "file_id": file_id
        }
    ]
    prompt = json.dumps(prompt, ensure_ascii=False)
    llm_answer = call_llm(prompt, "object_string")
    print(llm_answer)
    code = re.findall(r"```python\n(.*)\n```", llm_answer, flags=re.DOTALL)[0]
    save_path = f_get_save_path("flow.py")
    with open(save_path, mode="w", encoding="utf8") as f:
        f.write(code)

    save_path = f_get_save_path("__init__.py")
    with open(save_path, mode="w", encoding="utf8") as f:
        f.write("")

def _f_parse_node(df: pd.DataFrame, node_name):
    rules = ""
    for idx, row in df.iterrows():
        var_name = row["变量"]
        var_name = var_name.replace("\n", " ")
        rule_content = row["逻辑"]
        rule_out = row["输出"]
        rules = f"{rules}规则{idx + 1}: 变量:{var_name} 逻辑:{rule_content} 输出:{rule_out}\n"

    prompt = f_get_prompt_parse_node(node_name, rules)
    print(prompt)
    llm_answer = call_llm(prompt)
    code = re.findall(r"```python\n(.*)\n```", llm_answer, flags=re.DOTALL)[0]
    func_name = re.findall(r"def (.*)\(data", code)[0]
    save_path = f_get_save_path(f"{func_name}.py")
    print(code)
    with open(save_path, mode="w", encoding="utf8") as f:
        f.write(code)
    return func_name, code


def f_parse_strategy(file_path):
    wb = load_workbook(file_path)
    excel = pd.ExcelFile(file_path)
    sheet_names = excel.sheet_names
    if "流程图" not in sheet_names:
        GeneralException(ResultCodesEnum.NOT_FOUND, message=f"sheet【流程图】不存在")
    node_list = []
    for node_name in tqdm(sheet_names):
        if node_name == "流程图":
            continue
        df = excel.parse(sheet_name=node_name)
        func_name, code = _f_parse_node(df, node_name)
        node_list.append((node_name, func_name, code))
    _f_parse_flow(wb["流程图"], node_list)
    wb.close()
    excel.close()


if __name__ == "__main__":
    f_parse_strategy("./cache/策略节点配置demo.xlsx")