import os
import argparse
import subprocess
import cv2
import numpy as np
import glob
import sys

def read_image_safe(path):
    """
    [解决中文路径问题] 读取图片
    使用 numpy 读取字节流再解码，绕过 cv2.imread 的路径编码 bug
    """
    try:
        # np.fromfile 读取文件内容到内存
        img_array = np.fromfile(path, dtype=np.uint8)
        # cv2.imdecode 解码内存数据
        img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
        return img
    except Exception as e:
        print(f"Error reading image {path}: {e}")
        return None

def write_image_safe(path, img):
    """
    [解决中文路径问题] 保存图片
    使用 cv2.imencode 编码再用 numpy 保存，绕过 cv2.imwrite 的路径编码 bug
    """
    try:
        # 获取文件扩展名 (例如 .png)
        ext = os.path.splitext(path)[1]
        if not ext:
            ext = ".png"
        
        # cv2.imencode 编码图片
        success, img_array = cv2.imencode(ext, img)
        if success:
            # tofile 保存到文件
            img_array.tofile(path)
            return True
        else:
            print(f"Error encoding image for {path}")
            return False
    except Exception as e:
        print(f"Error writing image {path}: {e}")
        return False

def run_encoder(encoder, img_path, out_dir, grayscale):
    """
    调用 run.py 生成指定 encoder 的深度图
    """
    cmd = [
        sys.executable, 'run.py',
        '--encoder', encoder,
        '--img-path', img_path,
        '--outdir', out_dir,
        '--pred-only'
    ]
    
    if grayscale:
        cmd.append('--grayscale')
    
    print(f"Executing: {' '.join(cmd)}")
    try:
        subprocess.check_call(cmd)
    except subprocess.CalledProcessError as e:
        print(f"Error occurred while running {encoder} (Gray={grayscale}): {e}")
        return False
    return True

def add_text(img, text):
    """
    在图片左上角添加文字标注
    """
    font = cv2.FONT_HERSHEY_SIMPLEX
    font_scale = 1.0
    font_color = (255, 255, 255)
    thickness = 2
    position = (30, 50)
    
    cv2.putText(img, text, position, font, font_scale, (0, 0, 0), thickness + 3, cv2.LINE_AA)
    cv2.putText(img, text, position, font, font_scale, font_color, thickness, cv2.LINE_AA)
    return img

def create_comparison(image_files, folder_name, encoders, mode_suffix, out_folder_name):
    """
    生成对比图的核心逻辑
    """
    compare_dir = os.path.join(folder_name, out_folder_name)
    os.makedirs(compare_dir, exist_ok=True)
    
    print(f"\n--- Generating Comparisons ({out_folder_name}) -> {compare_dir} ---")
    
    for img_file in image_files:
        file_basename = os.path.basename(img_file)
        file_name_no_ext = os.path.splitext(file_basename)[0]
        
        # 1. 读取原图
        raw_img = read_image_safe(img_file)
        if raw_img is None:
            continue
        
        # 获取原图尺寸 (高度, 宽度)
        h, w = raw_img.shape[:2]
            
        # 标注原图
        raw_img_labeled = raw_img.copy()
        raw_img_labeled = add_text(raw_img_labeled, "Original")
        
        images_to_concat = [raw_img_labeled]
        valid_group = True
        
        # 2. 依次读取 4 个 encoder 的结果
        for encoder in encoders:
            result_dir = os.path.join(folder_name, f'V2_{encoder}{mode_suffix}')
            result_path = os.path.join(result_dir, file_name_no_ext + '.png')
            
            if not os.path.exists(result_path):
                print(f"Missing result for {file_basename} in {encoder}, skipping.")
                valid_group = False
                break
            
            depth_img = read_image_safe(result_path)
            
            if depth_img is None:
                print(f"Failed to read {result_path}, skipping.")
                valid_group = False
                break

            # 确保通道数正确 (转为3通道以便拼接)
            if len(depth_img.shape) == 2:
                depth_img = cv2.cvtColor(depth_img, cv2.COLOR_GRAY2BGR)
            
            # --- 关键修改：强制 Resize 到原图尺寸 ---
            # 无论原始输出尺寸如何，都强制缩放为 (w, h)
            if depth_img.shape[:2] != (h, w):
                depth_img = cv2.resize(depth_img, (w, h), interpolation=cv2.INTER_LINEAR)
            
            # 标注文字
            label = f"Depth_Anything_V2_{encoder}"
            depth_img = add_text(depth_img, label)
            
            images_to_concat.append(depth_img)
            
        if valid_group:
            # 水平拼接
            try:
                combined_result = cv2.hconcat(images_to_concat)
                # 保存
                save_path = os.path.join(compare_dir, f"{file_name_no_ext}_compare.png")
                if write_image_safe(save_path, combined_result):
                    print(f"Saved: {save_path}")
                else:
                    print(f"Failed to save: {save_path}")
            except Exception as e:
                print(f"Error combining images for {file_name_no_ext}: {e}")

def main():
    parser = argparse.ArgumentParser(description='Batch process depth estimation (Gray & Color) and compare.')
    parser.add_argument('--img-path', type=str, required=True, help='Path to input images folder')
    args = parser.parse_args()

    # 1. 确定根目录名称: 输入文件夹名 + "_results"
    folder_name = os.path.basename(os.path.normpath(args.img_path)) + "_results"
    
    encoders = ['vits', 'vitb', 'vitl', 'vitg']
    
    # 2. 批量生成深度图 (Run Inference)
    for encoder in encoders:
        # A. 灰度图
        gray_out_dir = os.path.join(folder_name, f'V2_{encoder}_gray_result')
        if not os.path.exists(gray_out_dir):
            run_encoder(encoder, args.img_path, gray_out_dir, grayscale=True)
        
        # B. 彩色图
        color_out_dir = os.path.join(folder_name, f'V2_{encoder}_colorful_result')
        if not os.path.exists(color_out_dir):
            run_encoder(encoder, args.img_path, color_out_dir, grayscale=False)

    # 3. 准备原图列表
    exts = ['*.jpg', '*.jpeg', '*.png', '*.bmp', '*.tif', '*.tiff']
    image_files = []
    for ext in exts:
        image_files.extend(glob.glob(os.path.join(args.img_path, ext)))
        image_files.extend(glob.glob(os.path.join(args.img_path, '**', ext), recursive=True))
    image_files = sorted(list(set(image_files)))
    
    if not image_files:
        print("No images found in input path.")
        return

    # 4. 生成灰度对比图
    create_comparison(
        image_files=image_files, 
        folder_name=folder_name, 
        encoders=encoders, 
        mode_suffix='_gray_result', 
        out_folder_name='compare_all_gray'
    )

    # 5. 生成彩色对比图
    create_comparison(
        image_files=image_files, 
        folder_name=folder_name, 
        encoders=encoders, 
        mode_suffix='_colorful_result', 
        out_folder_name='compare_all_colorful'
    )

    print(f"\nAll done! Results are in: {folder_name}")

if __name__ == "__main__":
    main()