"""
将文章中标题行后插入两幅图片
"""
import importlib
import os
import sys

from docx import Document
import pythoncom
from docx.opc.exceptions import PackageNotFoundError
from docx.shared import Inches

importlib.reload(sys)


# 读取文件夹下的docx文件名列表
def file_name(file_dir, file_format):
    file_list = []
    for root, dirs, files in os.walk(file_dir):
        for file in files:
            if os.path.splitext(file)[1] == file_format:
                file_list.append(os.path.join(root, file))
    return file_list


# 插入两张图片
def add_picture(doc, picture1, picture2):
    print(len(doc.paragraphs))
    try:
        # 若文章仅有两段，则在标题行之后插入两张图片
        if len(doc.paragraphs) <= 2:
            run = doc.paragraphs[0].add_run()
            run.add_break()
            run.add_picture(picture1, width=Inches(5.25))
            run.add_break()
            run.add_picture(picture2, width=Inches(5.25))
            return True
        # 若文章多于两段，则在标题行后插入空行，在正文第一段之后插入一张图片，第二段之后插入一张图片
        elif len(doc.paragraphs) >= 3:
            doc.paragraphs[1] = doc.paragraphs[1].insert_paragraph_before()
            print(len(doc.paragraphs))
            run1 = doc.paragraphs[2].add_run()
            run1.add_break()
            run1.add_picture(picture1, width=Inches(5.25))
            run2 = doc.paragraphs[3].add_run()
            run2.add_break()
            run2.add_picture(picture2, width=Inches(5.25))
            return True
        else:
            print("文档段落数不满足要求！")
            return False
    except:
        print("图片不符合要求：", picture1, picture2)
        return False


# 随机取列表中的两个值返回为新列表
def pick_random(list):
    import random
    if len(list) < 2:
        pass
    else:
        return random.sample(list, 2)


def main():
    file_list = file_name("D:\\文件处理\\2020.8.12-350篇足球", ".docx")
    print(len(file_list))
    picture_list = file_name("D:\\文件处理\\足球图片", ".jpg")
    print(len(picture_list))

    # file_list = file_name("D:\\文件处理\\2020.8.12-88篇棋牌\\M", ".docx")
    # print(len(file_list))
    # picture_list = file_name("D:\\文件处理\\麻将图片", ".jpg")
    # print(len(picture_list))

    # file_list = file_name("D:\\文件处理\\2020.8.12-88篇棋牌\\N", ".docx")
    # print(len(file_list))
    # picture_list = file_name("D:\\文件处理\\扑克图片", ".jpg")
    # print(len(picture_list))

    if len(picture_list) < 2 * len(file_list):
        print("图片素材不足！必须是文章数量的两倍以上")
        return

    for file in file_list:
        pythoncom.CoInitialize()
        try:
            doc = Document(file)
            picture = pick_random(picture_list)
            # print(picture)
            while not add_picture(doc, picture[0], picture[1]):
                # 若插入图片失败，删除不适用的图片素材，并重新选图
                picture_list.remove(picture[0])
                picture_list.remove(picture[1])
                picture = pick_random(picture_list)
            # 插入图片成功，则保存文件
            doc.save(file)
            # 删除已使用过的图片素材
            picture_list.remove(picture[0])
            picture_list.remove(picture[1])
        except PackageNotFoundError:
            pass
        finally:
            # 释放资源
            pythoncom.CoUninitialize()


if __name__ == '__main__':
    main()