import argparse
import os
import base64
import requests
import json
from PIL import Image
from io import BytesIO

def load_local_env():
    """
    Load optional local credentials from config/baidu_api.env.
    This file is ignored by git; see config/baidu_api.env.example.
    """
    root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
    env_path = os.path.join(root, 'config', 'baidu_api.env')
    if not os.path.exists(env_path):
        return
    with open(env_path, 'r', encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith('#') or '=' not in line:
                continue
            key, value = line.split('=', 1)
            os.environ.setdefault(key.strip(), value.strip().strip('"').strip("'"))


load_local_env()

API_KEY = os.environ.get("BAIDU_API_KEY", "")
SECRET_KEY = os.environ.get("BAIDU_SECRET_KEY", "")

# 输入和输出目录
INPUT_DIR = './去雾图像-北航合作'
OUTPUT_DIR = './Result_Baidu'

def get_access_token():
    """
    使用 AK，SK 生成鉴权签名（Access Token）
    """
    if not API_KEY or not SECRET_KEY:
        print("缺少 BAIDU_API_KEY 或 BAIDU_SECRET_KEY，请设置环境变量或 config/baidu_api.env")
        return None

    url = "https://aip.baidubce.com/oauth/2.0/token"
    params = {
        "grant_type": "client_credentials",
        "client_id": API_KEY,
        "client_secret": SECRET_KEY
    }
    try:
        response = requests.post(url, params=params)
        return response.json().get("access_token")
    except Exception as e:
        print(f"获取 Access Token 失败: {e}")
        return None

def process_image(file_path, access_token):
    """
    读取图片，调用百度API，返回处理后的图片数据
    """
    request_url = "https://aip.baidubce.com/rest/2.0/image-process/v1/dehaze"
    request_url = request_url + "?access_token=" + access_token

    try:
        # 使用 PIL 读取图片
        with Image.open(file_path) as img:
            # 将 PIL Image 对象转换为字节流
            img_buffer = BytesIO()
            # 保持原格式保存到内存中
            save_format = img.format if img.format else 'JPEG'
            img.save(img_buffer, format=save_format)
            img_bytes = img_buffer.getvalue()
        
        # base64 编码
        img_base64 = base64.b64encode(img_bytes)
        
        params = {"image": img_base64}
        headers = {'content-type': 'application/x-www-form-urlencoded'}
        
        # 调用 API
        response = requests.post(request_url, data=params, headers=headers)
        
        if response.status_code == 200:
            result = response.json()
            # 检查是否有 image 字段
            if "image" in result:
                return base64.b64decode(result["image"])
            else:
                print(f"API 返回错误: {result}")
                return None
        else:
            print(f"请求失败，状态码: {response.status_code}")
            return None

    except Exception as e:
        print(f"处理图片 {file_path} 时出错: {e}")
        return None

def main():
    global INPUT_DIR, OUTPUT_DIR

    parser = argparse.ArgumentParser(description="Run Baidu dehaze API on a folder of images.")
    parser.add_argument("--input-dir", default=INPUT_DIR)
    parser.add_argument("--output-dir", default=OUTPUT_DIR)
    args = parser.parse_args()
    INPUT_DIR = args.input_dir
    OUTPUT_DIR = args.output_dir

    # 1. 获取 Token
    access_token = get_access_token()
    if not access_token:
        return

    # 2. 确保输出目录存在
    if not os.path.exists(OUTPUT_DIR):
        os.makedirs(OUTPUT_DIR)
        print(f"已创建输出目录: {OUTPUT_DIR}")

    # 3. 遍历 ./Data 目录
    if not os.path.exists(INPUT_DIR):
        print(f"输入目录 {INPUT_DIR} 不存在")
        return

    supported_exts = ('.png', '.bmp', '.jpg', '.jpeg')
    
    print("开始批量处理...")
    
    files = os.listdir(INPUT_DIR)
    for filename in files:
        # 检查文件扩展名
        if filename.lower().endswith(supported_exts):
            input_path = os.path.join(INPUT_DIR, filename)
            # V1
            # output_path = os.path.join(OUTPUT_DIR, filename)
            # V2
            # 1. 分离文件名和扩展名
            name, ext = os.path.splitext(filename)
            # 2. 拼接新文件名
            new_filename = f"{name}_result{ext}"
            # 3. 生成最终输出路径
            output_path = os.path.join(OUTPUT_DIR, new_filename)
            
            
            print(f"正在处理: {filename} ...")
            
            # 处理图片
            processed_data = process_image(input_path, access_token)
            
            # 保存结果
            if processed_data:
                with open(output_path, 'wb') as f:
                    f.write(processed_data)
                print(f"  已保存至: {output_path}")
            else:
                print(f"  处理失败: {filename}")

    print("批量处理完成。")

if __name__ == '__main__':
    main()