from markitdown import MarkItDown
from openpyxl import load_workbook
from io import BytesIO
from PIL import Image
import base64
from zipfile import ZipFile

def extract_xlsx_all_images_base64(xlsx_path: str, max_width=1200):
    img_list = []
    wb = load_workbook(xlsx_path)
    # 1.先读取单元格内嵌图片
    for sheet_name in wb.sheetnames:
        ws = wb[sheet_name]
        for img in ws._images:
            img_bytes = img.ref
            pil_img = Image.open(BytesIO(img_bytes))
            if pil_img.width > max_width:
                ratio = max_width / pil_img.width
                new_size = (max_width, int(pil_img.height * ratio))
                pil_img = pil_img.resize(new_size, Image.Resampling.LANCZOS)
            buf = BytesIO()
            pil_img.save(buf, format="PNG", optimize=True)
            b64_str = base64.b64encode(buf.getvalue()).decode("utf-8")
            uri = f"data:image/png;base64,{b64_str}"
            img_list.append({"sheet": sheet_name, "data_uri": uri})
    wb.close()

    # 2.解析xlsx压缩包内drawing.xml，提取浮动形状图片
    with ZipFile(xlsx_path, "r") as zf:
        file_list = zf.namelist()
        # 遍历所有drawing，匹配media图片
        drawing_files = [f for f in file_list if f.startswith("xl/drawings/drawing") and f.endswith(".xml")]
        for drawing_path in drawing_files:
            # 解析drawing.xml获取图片rId
            import xml.etree.ElementTree as ET
            ns_a = {"a": "http://schemas.openxmlformats.org/drawingml/2006/main"}
            ns_xdr = {"xdr": "http://schemas.openxmlformats.org/drawingml/2006/spreadsheetDrawing"}
            xml_data = zf.read(drawing_path)
            root = ET.fromstring(xml_data)
            for blip in root.findall(".//a:blip", ns_a):
                r_id = blip.get("{http://schemas.openxmlformats.org/officeDocument/2006/relationships}embed")
                if not r_id:
                    continue
                # 读取drawing的关系文件
                rel_path = drawing_path.replace(".xml", ".rels")
                if rel_path not in file_list:
                    continue
                rel_xml = zf.read(rel_path)
                rel_root = ET.fromstring(rel_xml)
                ns_rel = {"r": "http://schemas.openxmlformats.org/package/2006/relationships"}
                for rel in rel_root.findall(".//r:Relationship", ns_rel):
                    if rel.get("Id") == r_id:
                        target = rel.get("Target")
                        media_path = f"xl/drawings/{target}"
                        if media_path in file_list:
                            media_bytes = zf.read(media_path)
                            pil_img = Image.open(BytesIO(media_bytes))
                            if pil_img.width > max_width:
                                ratio = max_width / pil_img.width
                                new_size = (max_width, int(pil_img.height * ratio))
                                pil_img = pil_img.resize(new_size, Image.Resampling.LANCZOS)
                            buf = BytesIO()
                            pil_img.save(buf, format="PNG", optimize=True)
                            b64_str = base64.b64encode(buf.getvalue()).decode("utf-8")
                            uri = f"data:image/png;base64,{b64_str}"
                            # 缺陷：无法精确归属到sheet名称，统一标记为浮动图片
                            img_list.append({"sheet": "FloatingShape", "data_uri": uri})
    return img_list


if __name__ == "__main__":
    xlsx_file = "test_xlsx.xlsx"
    out_md = "xlsx_out.md"
    md = MarkItDown()
    res = md.convert(source=xlsx_file)
    md_text = res.text_content
    images = extract_xlsx_all_images_base64(xlsx_file)
    print(f"一共提取图片：{len(images)}")
    if len(images) > 0:
        md_text += "\n\n## 提取的内嵌图片\n"
        for item in images:
            md_text += f"\n【Sheet：{item['sheet']}】\n![]({item['data_uri']})\n"
    with open(out_md, "w", encoding="utf-8") as f:
        f.write(md_text)
    print(f"✅转换完成，输出文件：{out_md}")