import os
import numpy as np
import cv2
from tool.lighter import adjust_highlights_shadows
from concurrent.futures import ThreadPoolExecutor
from PIL import Image
import piexif

# parser = argparse.ArgumentParser(description='Depth Anything V2')
# parser.add_argument('--input-paths', type=str, nargs='+', required=True,
#                         help='输入文件夹列表（多个路径，用空格分隔）')
# parser.add_argument('--output-path', type=str, default='./output', 
#                         help="按输出路径的结构输出处理好的图片")
# parser.add_argument('--model-path', type=str, default='./model/depth_anything_v2_vitl.pth',
#                         help='模型路径')
# args = parser.parse_args()


def process_single_image(img_path, depth_anything):
    """处理单个图片
    Args:
        img_path: 图片路径
    Returns:
        tuple: (处理后的图片numpy数组, 原始图片路径)
    """
    # 使用OpenCV读取图片
    img = cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), cv2.IMREAD_COLOR)
    if img is None:
        raise ValueError(f"无法读取图片: {img_path}")
    
    # 处理图片 - 这里假设adjust_highlights_shadows能处理OpenCV格式的图片
    processed_img = adjust_highlights_shadows(img, 180, 253, depth_anything)
    return processed_img, img_path

def save_image_with_exif(input_path, output_path, processed_img):
    """保存图片并保留EXIF信息
    Args:
        input_path: 原始图片路径
        output_path: 输出图片路径
        processed_img: 处理后的numpy数组图片
    """
    # 确保输出目录存在
    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    
    # 读取原始图片的EXIF信息
    exif_dict = None
    try:
        with exif_lock:  # 使用锁防止多线程同时读取EXIF
            with Image.open(input_path) as img:
                if 'exif' in img.info:
                    exif_dict = piexif.load(img.info['exif'])
    except Exception as e:
        print(f"Warning: 无法读取 {input_path} 的EXIF信息: {str(e)}")
    
    # 将OpenCV格式转换为PIL格式以便保存EXIF
    processed_img_rgb = cv2.cvtColor(processed_img, cv2.COLOR_BGR2RGB)
    pil_img = Image.fromarray(processed_img_rgb)
    
    # 保存图片
    try:
        if exif_dict:
            pil_img.save(output_path, exif=piexif.dump(exif_dict))
        else:
            pil_img.save(output_path)
    except Exception as e:
        print(f"Error: 无法保存图片 {output_path}: {str(e)}")
        raise

def process_single_file(input_path, output_root, input_base, depth_anything):
    """处理单个文件"""
    try:
        # 计算相对路径以保持文件夹结构
        rel_path = os.path.relpath(os.path.dirname(input_path), start=input_base)
        output_dir = os.path.join(output_root, rel_path)
        output_path = os.path.join(output_dir, os.path.basename(input_path))
        
        # 处理图片
        processed_img, _ = process_single_image(input_path, depth_anything)
        
        # 保存图片
        save_image_with_exif(input_path, output_path, processed_img)
        print(f"Processed and saved: {output_path}")
    except Exception as e:
        print(f"Error processing {input_path}: {str(e)}")

def process_images(input_paths, output_root, depth_anything, workers=4):
    """处理所有图片并保持原文件夹结构（多线程版本）
    Args:
        input_paths: 输入路径列表
        output_root: 输出根目录
        workers: 线程池大小
    """
    # 收集所有需要处理的文件
    all_files = []
    for input_path in input_paths:
        for root, dirs, files in os.walk(input_path):
            for file in files:
                if file.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.tiff')):
                    input_file_path = os.path.join(root, file)
                    all_files.append((input_file_path, input_path))
    
    # 使用线程池处理文件
    with ThreadPoolExecutor(max_workers=workers) as executor:
        futures = []
        for file_path, input_base in all_files:
            futures.append(
                executor.submit(
                    process_single_file,
                    file_path,
                    output_root,
                    input_base,
                    depth_anything
                )
            )
        
        # 等待所有任务完成
        for future in futures:
            try:
                future.result()
            except Exception as e:
                print(f"Error in processing: {str(e)}")
# if __name__ == '__main__':
    # model_path = args.model_path
    # input_paths = args.input_paths
    # output_path = args.output_path

    # # 线程锁，防止多线程同时访问EXIF数据时出现问题
    # exif_lock = threading.Lock()

    # input_paths = ["/home/dtyx/桌面/yhh/ReportGenerator/测试数据/山东国华无棣风电场叶片外部数据/A1（31301）",]
    # model_path = "./model/depth_anything_v2_vitl.pth"
    # output_path = "./output"
    # #启动模型
    # depth_anything = model_start(model_path)
    # # 创建输出目录
    # os.makedirs(output_path, exist_ok=True)
    # # 处理所有图片
    # process_images(input_paths, output_path, depth_anything)