from typing import Dict
import os
from docx import Document
from tools.document_tools import add_picture_to_table, add_table_to_document
def add_jf_picture_table(
        typical_picture_dict: Dict[str, str],
        defect_picture_dict: Dict[str, str],
        TYPICAL_MUBAN_DIR: str,
        DEFECT_MUBAN_DIR: str,
        output_dir: str,
):
    """添加金风版本的图片展示表格"""
    # 初始化文档对象
    doc = Document()
    target_filename = os.path.join(output_dir, "output.docx")
    
    # 处理典型图
    typical_data = []
    defect_keys = []  # 记录有缺陷的典型图key
    
    # 准备典型图数据
    for key, value in typical_picture_dict.items():
        if "损伤" in key or "损伤" in value:
            # 情况二：有缺陷
            defect_count = int(value.split("损伤类型")[1].split("处")[0])
            typical_data.append([key, "损伤{}处，详见下表".format(defect_count), value])
            defect_keys.append((key, defect_count))
        else:
            # 情况一：正常图片
            typical_data.append([key, value, ""])  # 第三行留空
    
    # 添加典型图表格
    if typical_data:
        # 将数据转换为适合表格的格式（3行×n列）
        # 需要将数据从行优先转为列优先
        cols = len(typical_data)
        table_data = [[], [], []]
        for col_data in typical_data:
            for i in range(3):
                table_data[i].append(col_data[i] if i < len(col_data) else "")
        
        # 添加典型图表格
        add_table_to_document(
            target_filename=target_filename,
            source_filename=TYPICAL_MUBAN_DIR,
            rows=3,
            cols=cols,
            table_num=0,
            data=table_data
        )
        
        # 添加典型图中的图片
        doc = Document(target_filename)
        for col_idx, (key, value) in enumerate(typical_picture_dict.items()):
            if "损伤" not in key and "损伤" not in value:  # 只有正常图片才添加
                add_picture_to_table(
                    target_doc=doc,
                    target_filename=target_filename,
                    row=1,  # 第二行是图片
                    col=col_idx,
                    image_path=value
                )
        doc.save(target_filename)
    
    # 处理缺陷图
    if defect_keys:
        # 遍历所有有缺陷的典型图
        for typical_key, defect_count in defect_keys:
            # 计算需要多少个缺陷图表格（每表5列）
            table_count = (defect_count + 4) // 5  # 向上取整
            
            for table_idx in range(table_count):
                # 计算当前表格的列数
                cols_in_table = min(5, defect_count - table_idx * 5)
                
                # 准备缺陷图数据
                defect_data = []
                for i in range(cols_in_table):
                    defect_idx = table_idx * 5 + i
                    defect_key = list(defect_picture_dict.keys())[defect_idx]
                    defect_value = list(defect_picture_dict.values())[defect_idx]
                    defect_data.append([defect_key, defect_value])
                
                # 将数据转换为适合表格的格式（2行×n列）
                table_data = [[], []]
                for col_data in defect_data:
                    for i in range(2):
                        table_data[i].append(col_data[i] if i < len(col_data) else "")
                
                # 添加缺陷图表格
                add_table_to_document(
                    target_filename=target_filename,
                    source_filename=DEFECT_MUBAN_DIR,
                    rows=2,
                    cols=cols_in_table,
                    table_num=0,
                    data=table_data
                )
                
                # 添加缺陷图中的图片
                doc = Document(target_filename)
                for col_idx in range(cols_in_table):
                    defect_idx = table_idx * 5 + col_idx
                    defect_value = list(defect_picture_dict.values())[defect_idx]
                    add_picture_to_table(
                        target_doc=doc,
                        target_filename=target_filename,
                        row=1,  # 第二行是图片
                        col=col_idx,
                        image_path=defect_value
                    )
                doc.save(target_filename)
    
    return target_filename


add_jf_picture_table({""})