import json
import os
from docx import Document
from docx.shared import Pt, RGBColor, Inches
from docx.enum.text import WD_ALIGN_PARAGRAPH
from docx.oxml.shared import qn, OxmlElement
from docx.oxml.ns import nsdecls

def json_to_docx(json_str, output_path):
    try:
        json_data = json.loads(json_str)
        print(f"\n开始转换JSON到DOCX文档，输出路径: {output_path}")
        doc = Document()
        total_elements = len(json_data)
        print(f"文档包含 {total_elements} 个元素(段落、表格和图片)")

        for i, element in enumerate(json_data, 1):
            try:
                print(f"\n处理元素 {i}/{total_elements}: ", end="")
                if element["type"] == "text":
                    print(f"段落 (长度: {len(element['content']['runs'])}个runs)")
                    add_paragraph_from_json(doc, element["content"])
                elif element["type"] == "table":
                    try:
                        rows = element["content"]["rows"]
                        cols = element["content"]["cols"]
                        merges = len(element["content"].get("merged_cells", []))
                        images = len(element["content"].get("images", []))
                        print(f"表格 ({rows}行×{cols}列, 包含 {merges} 个合并单元格, {images} 张图片)")
                        add_table_from_json(doc, element["content"])
                    except KeyError as e:
                        print(f"表格数据格式错误，缺少必要字段: {str(e)}")
                        continue
                elif element["type"] == "image":
                    print(f"图片 (路径: {element['content']['path']}, 尺寸: {element['content']['width_inches']}×{element['content']['height_inches']}英寸)")
                    add_image_from_json(doc, element["content"])
                else:
                    print(f"未知元素类型: {element['type']}")
            except Exception as e:
                print(f"处理元素 {i} 时出错: {str(e)}")
                continue

        print("\n正在保存文档...")
        doc.save(output_path)
        print(f"文档已成功保存到 {output_path}")
    except Exception as e:
        print(f"转换过程中发生错误: {str(e)}")
        raise

def add_image_from_json(doc, image_data, parent=None):
    """将JSON中的图片数据添加到docx文档或表格单元格"""
    try:
        # 检查图片文件是否存在
        if not os.path.exists(image_data["path"]):
            print(f"    警告: 图片文件不存在 {image_data['path']}")
            return

        # 添加段落用于放置图片
        if parent is None:
            paragraph = doc.add_paragraph()
        else:
            # 如果是表格单元格中的图片，使用单元格的段落
            if len(parent.paragraphs) == 0:
                paragraph = parent.add_paragraph()
            else:
                paragraph = parent.paragraphs[0]

        paragraph.alignment = WD_ALIGN_PARAGRAPH.CENTER

        # 添加图片到段落
        run = paragraph.add_run()

        # 使用Word中的原始尺寸
        width = Inches(image_data["width_inches"])
        height = Inches(image_data["height_inches"])
        run.add_picture(image_data["path"], width=width, height=height)

        print("    图片添加成功")
    except Exception as e:
        print(f"    添加图片失败: {str(e)}")

def add_paragraph_from_json(doc, para_json, parent=None):
    """添加段落，支持在表格单元格内添加"""
    if parent is None:
        paragraph = doc.add_paragraph()
    else:
        paragraph = parent.add_paragraph()

    print(f"  添加段落 (对齐: {para_json['alignment']})")

    # 设置段落对齐方式
    alignment_map = {
        "left": WD_ALIGN_PARAGRAPH.LEFT,
        "center": WD_ALIGN_PARAGRAPH.CENTER,
        "right": WD_ALIGN_PARAGRAPH.RIGHT,
        "justify": WD_ALIGN_PARAGRAPH.JUSTIFY
    }
    paragraph.alignment = alignment_map.get(para_json["alignment"], WD_ALIGN_PARAGRAPH.LEFT)

    # 添加文本运行(runs)
    for run_idx, run_json in enumerate(para_json["runs"], 1):
        run = paragraph.add_run(run_json["text"])
        if run_json["has_page_break"]:
            import docx
            run.add_break(docx.enum.text.WD_BREAK.PAGE)
        font = run.font

        print(f"    添加run {run_idx}: '{run_json['text']}' "
              f"(字体: {run_json['font']['name']}, 大小: {run_json['font']['size']}, "
              f"加粗: {run_json['font']['bold']}, 斜体: {run_json['font']['italic']})")

        # 设置字体样式
        if run_json["font"]["name"]:
            font.name = run_json["font"]["name"]
            run.element.rPr.rFonts.set(qn('w:eastAsia'), run_json["font"]["name"])

        if run_json["font"]["size"]:
            font.size = Pt(run_json["font"]["size"])

        font.bold = run_json["font"]["bold"]
        font.italic = run_json["font"]["italic"]
        font.underline = run_json["font"]["underline"]

        # 设置字体颜色
        if run_json["font"]["color"]:
            color = run_json["font"]["color"]
            font.color.rgb = RGBColor(color["r"], color["g"], color["b"])
            print(f"      设置颜色: RGB({color['r']}, {color['g']}, {color['b']})")

def add_table_from_json(doc, table_json):
    """
    将JSON中的表格数据添加到docx文档，支持精确的图片位置和尺寸

    参数:
        doc: Document对象
        table_json: 包含表格数据的JSON对象

    返回:
        None
    """
    try:
        # 验证必要字段
        if "rows" not in table_json or "cols" not in table_json:
            raise ValueError("表格数据缺少rows或cols字段")

        rows = table_json["rows"]
        cols = table_json["cols"]

        if rows <= 0 or cols <= 0:
            raise ValueError(f"无效的行列数: rows={rows}, cols={cols}")

        print(f"  创建表格: {rows}行 × {cols}列")
        table = doc.add_table(rows=rows, cols=cols)
        table.style = 'Table Grid'

        # 设置列宽
        if "col_widths" in table_json and any(table_json["col_widths"]):
            print("  设置列宽...")
            for col_idx, width in enumerate(table_json["col_widths"]):
                if width is not None and col_idx < cols:
                    try:
                        width_inches = float(width) if isinstance(width, str) else width
                        twips_width = int(width_inches * 1440)  # 1英寸=1440 Twips
                        for cell in table.columns[col_idx].cells:
                            tc = cell._tc
                            tcPr = tc.get_or_add_tcPr()
                            tcW = tcPr.first_child_found_in("w:tcW")
                            if tcW is None:
                                tcW = OxmlElement('w:tcW')
                                tcPr.append(tcW)
                            tcW.set(qn('w:w'), str(twips_width))
                            tcW.set(qn('w:type'), 'dxa')
                    except Exception as e:
                        print(f"    设置列宽{col_idx}时出错: {str(e)}")
                        continue

        # 设置行高
        if "row_heights" in table_json and any(table_json["row_heights"]):
            print("  设置行高...")
            for row_idx, height in enumerate(table_json["row_heights"]):
                if height is not None and row_idx < rows:
                    try:
                        height_inches = float(height) if isinstance(height, str) else height
                        twips_height = int(height_inches * 1440)
                        tr = table.rows[row_idx]._tr
                        trPr = tr.get_or_add_trPr()
                        trHeight = OxmlElement('w:trHeight')
                        trHeight.set(qn('w:val'), str(twips_height))
                        trHeight.set(qn('w:hRule'), 'atLeast')
                        trPr.append(trHeight)
                    except Exception as e:
                        print(f"    设置行高{row_idx}时出错: {str(e)}")
                        continue

        # 处理合并单元格
        if "merged_cells" in table_json:
            for merge_info in table_json["merged_cells"]:
                try:
                    start_row = merge_info["start_row"]
                    start_col = merge_info["start_col"]
                    end_row = merge_info["end_row"]
                    end_col = merge_info["end_col"]

                    if (0 <= start_row < rows and 0 <= end_row < rows and
                            0 <= start_col < cols and 0 <= end_col < cols):
                        print(f"    合并单元格: 从({start_row},{start_col})到({end_row},{end_col})")
                        start_cell = table.cell(start_row, start_col)
                        end_cell = table.cell(end_row, end_col)
                        start_cell.merge(end_cell)
                    else:
                        print(f"    无效的合并范围: 从({start_row},{start_col})到({end_row},{end_col})")
                except Exception as e:
                    print(f"    处理合并单元格时出错: {str(e)}")
                    continue

        # 填充表格内容
        if "cells" in table_json:
            for row_data in table_json["cells"]:
                for cell_data in row_data:
                    try:
                        row_idx = cell_data["row"]
                        col_idx = cell_data["col"]

                        if row_idx >= rows or col_idx >= cols:
                            continue

                        # 跳过被合并的非主单元格
                        if cell_data.get("is_merged", False) and not cell_data.get("merge_info", {}).get("is_primary", True):
                            print(f"    跳过被合并的单元格({row_idx},{col_idx})")
                            continue

                        cell = table.cell(row_idx, col_idx)
                        print(f"    处理单元格({row_idx},{col_idx})")

                        # 清空原有内容
                        for p in cell.paragraphs:
                            p._element.getparent().remove(p._element)

                        # 添加文本内容
                        if "content" in cell_data:
                            for para in cell_data["content"]:
                                add_paragraph_from_json(None, para, parent=cell)

                        # 设置单元格格式
                        set_cell_alignment(cell, cell_data)
                        if "border" in cell_data:
                            set_cell_border(cell, cell_data["border"])
                        if "shading" in cell_data:
                            set_cell_shading(cell, cell_data["shading"])
                        if "margins" in cell_data:
                            set_cell_margins(cell, cell_data["margins"])

                        # 处理单元格中的图片（精确位置）
                        if "images" in cell_data:
                            print(f"    单元格({row_idx},{col_idx})中包含 {len(cell_data['images'])} 张图片")
                            for img_data in cell_data["images"]:
                                try:
                                    # 确保单元格有段落
                                    if len(cell.paragraphs) == 0:
                                        cell.add_paragraph()

                                    # 获取图片尺寸
                                    width = Inches(img_data.get("width_inches", 2.0))  # 默认2英寸
                                    height = Inches(img_data.get("height_inches", 1.5))  # 默认1.5英寸

                                    # 根据位置信息确定插入点
                                    para_index = img_data.get("location", {}).get("para_index", 0)
                                    run_index = img_data.get("location", {}).get("run_index", 0)

                                    # 确保段落存在
                                    while len(cell.paragraphs) <= para_index:
                                        cell.add_paragraph()

                                    paragraph = cell.paragraphs[para_index]

                                    # 确保run存在
                                    if len(paragraph.runs) <= run_index:
                                        for _ in range(run_index - len(paragraph.runs) + 1):
                                            paragraph.add_run()

                                    run = paragraph.runs[run_index]

                                    # 添加图片
                                    if os.path.exists(img_data["file_path"]):
                                        run.add_picture(img_data["file_path"], width=width, height=height)
                                        print(f"      图片 {img_data['file_name']} 添加成功 (位置: 段落{para_index}, run{run_index})")
                                    else:
                                        print(f"      图片文件不存在: {img_data['file_path']}")
                                except Exception as e:
                                    print(f"      处理图片 {img_data.get('file_name', '未知')} 时出错: {str(e)}")
                                    continue

                    except Exception as e:
                        print(f"    处理单元格({row_idx},{col_idx})时出错: {str(e)}")
                        continue

        # 处理表格级别的图片（不属于任何特定单元格）
        if "images" in table_json:
            print(f"  处理表格中的 {len(table_json['images'])} 张独立图片...")
            for img_data in table_json["images"]:
                try:
                    if "location" in img_data and img_data["location"]["type"] == "table":
                        loc = img_data["location"]
                        row = loc.get("row", 0)
                        col = loc.get("cell", 0)

                        if 0 <= row < rows and 0 <= col < cols:
                            cell = table.cell(row, col)

                            # 确保单元格有段落
                            if len(cell.paragraphs) == 0:
                                cell.add_paragraph()

                            # 获取图片尺寸
                            width = Inches(img_data.get("width_inches", 2.0))
                            height = Inches(img_data.get("height_inches", 1.5))

                            # 添加图片到单元格的第一个段落
                            paragraph = cell.paragraphs[0]
                            run = paragraph.add_run()

                            if os.path.exists(img_data["file_path"]):
                                run.add_picture(img_data["file_path"], width=width, height=height)
                                print(f"    图片 {img_data['file_name']} 添加成功 (单元格{row},{col})")
                            else:
                                print(f"    图片文件不存在: {img_data['file_path']}")
                        else:
                            print(f"    无效的图片位置: 行{row}列{col}")
                except Exception as e:
                    print(f"    处理表格图片时出错: {str(e)}")
                    continue

    except Exception as e:
        print(f"创建表格时发生严重错误: {str(e)}")
        raise

def format_cell(cell, cell_data):
    """设置单元格完整格式"""
    # 清空原有内容
    for p in cell.paragraphs:
        p._element.getparent().remove(p._element)

    # 添加内容
    for para in cell_data["content"]:
        add_paragraph_from_json(None, para, parent=cell)

    # 设置对齐方式
    set_cell_alignment(cell, cell_data)

    # 设置边框
    set_cell_border(cell, cell_data["border"])

    # 设置背景色
    if cell_data.get("shading"):
        set_cell_shading(cell, cell_data["shading"])

    # 设置边距
    if cell_data.get("margins"):
        set_cell_margins(cell, cell_data["margins"])

    # 处理单元格中的图片
    if "images" in cell_data:
        print(f"    单元格中包含 {len(cell_data['images'])} 张图片")
        for img_data in cell_data["images"]:
            add_image_from_json(None, img_data, parent=cell)

def set_cell_alignment(cell, cell_data):
    """设置单元格对齐（水平和垂直）"""
    # 水平对齐
    if cell.paragraphs:
        align_map = {
            "left": WD_ALIGN_PARAGRAPH.LEFT,
            "center": WD_ALIGN_PARAGRAPH.CENTER,
            "right": WD_ALIGN_PARAGRAPH.RIGHT,
            "justify": WD_ALIGN_PARAGRAPH.JUSTIFY
        }
        cell.paragraphs[0].alignment = align_map.get(cell_data["alignment"], WD_ALIGN_PARAGRAPH.LEFT)

    # 垂直对齐设置
    tcPr = cell._tc.get_or_add_tcPr()
    vAlign = OxmlElement('w:vAlign')
    align_value = cell_data.get('vertical_align', 'top')
    print(f"    设置垂直对齐: {align_value}")

    # 确保使用有效的对齐值
    valid_alignments = ['top', 'center', 'bottom']
    if align_value not in valid_alignments:
        align_value = 'top'  # 默认值

    vAlign.set(qn('w:val'), align_value)
    tcPr.append(vAlign)

def set_cell_shading(cell, shading):
    """设置单元格背景色"""
    tcPr = cell._tc.get_or_add_tcPr()
    shd = OxmlElement('w:shd')
    shd.set(qn('w:fill'), shading["color"])
    if shading.get("theme"):
        shd.set(qn('w:themeColor'), shading["theme"])
    tcPr.append(shd)

def set_cell_margins(cell, margins):
    """设置单元格边距"""
    tcPr = cell._tc.get_or_add_tcPr()
    tcMar = OxmlElement('w:tcMar')

    for side, margin in margins.items():
        side_el = OxmlElement(f'w:{side}')
        side_el.set(qn('w:w'), margin["w"])
        side_el.set(qn('w:type'), margin["type"])
        tcMar.append(side_el)

    tcPr.append(tcMar)

def set_cell_border(cell, border_data):
    """
    设置单元格边框
    :param cell: 单元格对象
    :param border_data: 边框数据
    """
    tc = cell._tc
    tcPr = tc.get_or_add_tcPr()

    # 检查是否存在边框元素，不存在则创建
    tcBorders = tcPr.first_child_found_in("w:tcBorders")
    if tcBorders is None:
        tcBorders = OxmlElement('w:tcBorders')
        tcPr.append(tcBorders)

    # 设置各边边框
    for side in ['top', 'left', 'bottom', 'right']:
        if side in border_data:
            border = border_data[side]
            border_el = OxmlElement(f'w:{side}')
            border_el.set(qn('w:val'), border.get('style', 'single'))
            border_el.set(qn('w:sz'), str(border.get('size', 4)))
            border_el.set(qn('w:color'), border.get('color', '000000'))
            tcBorders.append(border_el)

# 使用示例
if __name__ == "__main__":
    # 假设我们已经有了之前生成的JSON数据
    input_json = "output.json"
    output_path = "restored.docx"

    print(f"从 {input_json} 读取JSON数据...")
    with open(input_json, "r", encoding="utf-8") as f:
        json_data = json.load(f)

    # 将JSON转换回DOCX
    json_to_docx(json.dumps(json_data), output_path)