#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Excel读取工具 - 支持图片提取
遍历每个sheet，从第二行开始读取数据，并提取图片保存到本地
"""

import os
import sys
from pathlib import Path
import shutil

def install_requirements():
    """安装必要的依赖包"""
    import subprocess
    
    packages = ['openpyxl', 'Pillow']
    
    for package in packages:
        try:
            if package == 'Pillow':
                import PIL
            else:
                __import__(package)
        except ImportError:
            print(f"正在安装 {package}...")
            subprocess.check_call([sys.executable, '-m', 'pip', 'install', package])

def extract_images_from_excel(file_path, output_dir="excel_images"):
    """
    从Excel文件中提取图片并保存到本地
    
    Args:
        file_path (str): Excel文件路径
        output_dir (str): 图片保存目录
    
    Returns:
        dict: 包含图片信息的字典
    """
    try:
        import openpyxl
        from openpyxl.drawing.image import Image as ExcelImage
        from PIL import Image
        import io
        
        print(f"正在从Excel文件中提取图片...")
        print("=" * 50)
        
        # 创建图片保存目录
        if os.path.exists(output_dir):
            shutil.rmtree(output_dir)
        os.makedirs(output_dir, exist_ok=True)
        
        # 打开Excel文件
        workbook = openpyxl.load_workbook(file_path)
        sheet_names = workbook.sheetnames
        
        print(f"发现 {len(sheet_names)} 个sheet: {sheet_names}")
        
        all_images = {}
        image_counter = 1
        
        for sheet_name in sheet_names:
            print(f"\n处理sheet: {sheet_name}")
            print("-" * 30)
            
            sheet = workbook[sheet_name]
            sheet_images = []
            
            # 遍历sheet中的所有图片
            for image in sheet._images:
                try:
                    # 获取图片数据
                    image_data = image._data()
                    
                    # 创建PIL图片对象
                    pil_image = Image.open(io.BytesIO(image_data))
                    
                    # 生成图片文件名
                    image_filename = f"{sheet_name}_image_{image_counter:03d}.png"
                    image_path = os.path.join(output_dir, image_filename)
                    
                    # 保存图片
                    pil_image.save(image_path, 'PNG')
                    
                    # 获取图片在Excel中的位置信息
                    anchor = image.anchor
                    if hasattr(anchor, '_from'):
                        row = anchor._from.row + 1  # Excel行号从1开始
                        col = anchor._from.col + 1  # Excel列号从1开始
                    else:
                        row = 1
                        col = 1
                    
                    image_info = {
                        'filename': image_filename,
                        'path': image_path,
                        'sheet': sheet_name,
                        'row': row,
                        'col': col,
                        'size': pil_image.size,
                        'format': pil_image.format
                    }
                    
                    sheet_images.append(image_info)
                    print(f"  保存图片: {image_filename} (位置: 行{row}, 列{col})")
                    
                    image_counter += 1
                    
                except Exception as e:
                    print(f"  处理图片时出错: {e}")
                    continue
            
            all_images[sheet_name] = sheet_images
            print(f"  Sheet '{sheet_name}' 提取了 {len(sheet_images)} 张图片")
        
        workbook.close()
        return all_images
        
    except Exception as e:
        print(f"提取图片失败: {e}")
        return {}

def read_excel_with_images(file_path, start_row=2, output_dir="excel_images"):
    """
    读取Excel文件并提取图片
    
    Args:
        file_path (str): Excel文件路径
        start_row (int): 开始读取的行号
        output_dir (str): 图片保存目录
    
    Returns:
        tuple: (数据字典, 图片信息字典)
    """
    try:
        import openpyxl
        
        print(f"正在读取Excel文件: {file_path}")
        print("=" * 50)
        
        # 读取数据
        workbook = openpyxl.load_workbook(file_path, data_only=True)
        sheet_names = workbook.sheetnames
        
        all_data = {}
        all_images = {}
        image_counter = 1
        
        # 创建图片保存目录
        if os.path.exists(output_dir):
            shutil.rmtree(output_dir)
        os.makedirs(output_dir, exist_ok=True)
        
        for sheet_name in sheet_names:
            print(f"\n处理sheet: {sheet_name}")
            print("-" * 30)
            
            sheet = workbook[sheet_name]
            
            # 读取数据
            data_rows = []
            row_count = 0
            
            for row in sheet.iter_rows(min_row=start_row, values_only=True):
                row_count += 1
                
                row_data = []
                for cell in row:
                    if cell is None:
                        row_data.append("")
                    else:
                        row_data.append(str(cell))
                
                if not all(cell.strip() == '' for cell in row_data):
                    data_rows.append(row_data)
            
            all_data[sheet_name] = data_rows
            
            # 提取图片
            sheet_images = []
            for image in sheet._images:
                try:
                    import io
                    from PIL import Image
                    
                    # 获取图片数据
                    image_data = image._data()
                    
                    # 创建PIL图片对象
                    pil_image = Image.open(io.BytesIO(image_data))
                    
                    # 生成图片文件名
                    image_filename = f"{sheet_name}_image_{image_counter:03d}.png"
                    image_path = os.path.join(output_dir, image_filename)
                    
                    # 保存图片
                    pil_image.save(image_path, 'PNG')
                    
                    # 获取图片位置
                    anchor = image.anchor
                    if hasattr(anchor, '_from'):
                        row = anchor._from.row + 1
                        col = anchor._from.col + 1
                    else:
                        row = 1
                        col = 1
                    
                    image_info = {
                        'filename': image_filename,
                        'path': image_path,
                        'sheet': sheet_name,
                        'row': row,
                        'col': col,
                        'size': pil_image.size,
                        'format': pil_image.format
                    }
                    
                    sheet_images.append(image_info)
                    print(f"  保存图片: {image_filename} (位置: 行{row}, 列{col})")
                    
                    image_counter += 1
                    
                except Exception as e:
                    print(f"  处理图片时出错: {e}")
                    continue
            
            all_images[sheet_name] = sheet_images
            
            print(f"  数据行数: {len(data_rows)}")
            print(f"  图片数量: {len(sheet_images)}")
        
        workbook.close()
        return all_data, all_images
        
    except Exception as e:
        print(f"读取Excel失败: {e}")
        return {}, {}

def create_data_with_images(data, images):
    """
    将图片信息关联到数据数组中
    
    Args:
        data (dict): 数据字典
        images (dict): 图片信息字典
    
    Returns:
        dict: 包含图片关联的数据字典
    """
    print("\n" + "=" * 50)
    print("关联图片信息到数据")
    print("=" * 50)
    
    enhanced_data = {}
    
    for sheet_name in data.keys():
        print(f"\n处理sheet: {sheet_name}")
        print("-" * 30)
        
        sheet_data = data[sheet_name]
        sheet_images = images.get(sheet_name, [])
        
        # 为每行数据添加关联的图片信息
        enhanced_sheet_data = []
        
        for row_idx, row_data in enumerate(sheet_data):
            # 计算实际行号（从第二行开始）
            actual_row = row_idx + 2
            
            # 查找该行关联的图片
            row_images = []
            for img_info in sheet_images:
                if img_info['row'] == actual_row:
                    row_images.append({
                        'filename': img_info['filename'],
                        'path': img_info['path'],
                        'size': img_info['size'],
                        'col': img_info['col']
                    })
            
            # 创建增强的行数据
            enhanced_row = {
                'row_number': actual_row,
                'data': row_data,
                'images': row_images,
                'image_count': len(row_images)
            }
            
            enhanced_sheet_data.append(enhanced_row)
            
            if row_images:
                print(f"  行 {actual_row}: 关联了 {len(row_images)} 张图片")
                for img in row_images:
                    print(f"    - {img['filename']} (列{img['col']})")
        
        enhanced_data[sheet_name] = enhanced_sheet_data
        
        print(f"  Sheet '{sheet_name}': {len(enhanced_sheet_data)} 行数据，{len(sheet_images)} 张图片")
    
    return enhanced_data

def save_enhanced_data(enhanced_data, output_file="enhanced_excel_data.json"):
    """保存增强的数据到JSON文件"""
    try:
        import json
        
        # 转换为可序列化的格式
        serializable_data = {}
        for sheet_name, sheet_data in enhanced_data.items():
            serializable_data[sheet_name] = []
            for row_info in sheet_data:
                serializable_data[sheet_name].append({
                    'row_number': row_info['row_number'],
                    'data': row_info['data'],
                    'images': row_info['images'],
                    'image_count': row_info['image_count']
                })
        
        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump(serializable_data, f, ensure_ascii=False, indent=2)
        
        print(f"\n增强数据已保存到: {output_file}")
        
    except Exception as e:
        print(f"保存数据失败: {e}")

def create_image_index(images, output_file="image_index.txt"):
    """创建图片索引文件"""
    try:
        with open(output_file, 'w', encoding='utf-8') as f:
            f.write("Excel图片索引\n")
            f.write("=" * 50 + "\n\n")
            
            total_images = 0
            for sheet_name, sheet_images in images.items():
                f.write(f"Sheet: {sheet_name}\n")
                f.write("-" * 30 + "\n")
                
                if sheet_images:
                    for img in sheet_images:
                        f.write(f"文件: {img['filename']}\n")
                        f.write(f"  位置: 行{img['row']}, 列{img['col']}\n")
                        f.write(f"  尺寸: {img['size'][0]}x{img['size'][1]}\n")
                        f.write(f"  格式: {img['format']}\n")
                        f.write(f"  路径: {img['path']}\n\n")
                        total_images += 1
                else:
                    f.write("无图片\n\n")
            
            f.write(f"\n总计: {total_images} 张图片\n")
        
        print(f"图片索引已保存到: {output_file}")
        
    except Exception as e:
        print(f"创建图片索引失败: {e}")

def main():
    """主函数"""
    print("Excel读取工具 - 支持图片提取")
    print("=" * 40)
    
    # 查找Excel文件
    excel_files = []
    for f in os.listdir('.'):
        if f.lower().endswith(('.xlsx', '.xls')) and not f.startswith('~$'):
            excel_files.append(f)
    
    if not excel_files:
        print("当前目录下没有找到Excel文件")
        return
    
    print(f"找到 {len(excel_files)} 个Excel文件:")
    for i, excel_file in enumerate(excel_files, 1):
        print(f"{i}. {excel_file}")
    
    # 选择文件
    if len(excel_files) == 1:
        selected_file = excel_files[0]
        print(f"\n自动选择: {selected_file}")
    else:
        try:
            choice = int(input(f"\n请选择文件 (1-{len(excel_files)}): ")) - 1
            selected_file = excel_files[choice]
        except (ValueError, IndexError):
            selected_file = excel_files[0]
    
    # 安装依赖
    install_requirements()
    
    # 读取Excel数据和图片
    data, images = read_excel_with_images(selected_file, start_row=2)
    
    if not data:
        print("没有读取到任何数据")
        return
    
    # 关联图片信息到数据
    enhanced_data = create_data_with_images(data, images)
    
    # 保存增强数据
    save_enhanced_data(enhanced_data)
    
    # 创建图片索引
    create_image_index(images)
    
    # 显示统计信息
    print("\n" + "=" * 50)
    print("处理完成统计")
    print("=" * 50)
    
    total_rows = 0
    total_images = 0
    
    for sheet_name in enhanced_data.keys():
        sheet_data = enhanced_data[sheet_name]
        sheet_images = images.get(sheet_name, [])
        
        rows = len(sheet_data)
        img_count = len(sheet_images)
        
        total_rows += rows
        total_images += img_count
        
        print(f"Sheet '{sheet_name}': {rows} 行数据, {img_count} 张图片")
    
    print(f"\n总计: {total_rows} 行数据, {total_images} 张图片")
    print(f"图片保存目录: excel_images/")
    print(f"数据文件: enhanced_excel_data.json")
    print(f"图片索引: image_index.txt")

if __name__ == "__main__":
    main()