# -*- coding: utf-8 -*-
import requests
import json
import os, sys
import redis
from concurrent.futures import ThreadPoolExecutor, as_completed

base_path = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(base_path)
import tos_util

redis_client = redis.Redis(host='172.16.0.24', port=6379, db=8, password='aiyingli@@123', socket_timeout=10,
                           decode_responses=True)
doubao_api_token = 'Bearer fcc424e5-58af-494d-9683-5787413a26c9'
dir = '/Users/yaowentong/Desktop/image/'
tos_dir = ''
size_map = {
    '1:1': '2048x2048',
    '4:3': '2304x1728',
    '3:4': '1728x2304',
    '16:9': '2560x1440',
    '9:16': '1440x2560',
    '3:2': '2496x1664',
    '2:3': '1664x2496',
    '21:9': '3024x1296',
}


def gen_image(prompt, size,task_id):
    url = "https://ark.cn-beijing.volces.com/api/v3/images/generations"
    payload = json.dumps({
        "model": "doubao-seedream-4-0-250828",
        "prompt": prompt,
        "size": size_map[size],
        "sequential_image_generation": "disabled",
        "sequential_image_generation_options": {
            "max_images": 1
        },
        "stream": False,
        "response_format": "url",
        "watermark": False
    })
    headers = {
        'Authorization': doubao_api_token,
        'Content-Type': 'application/json'
    }
    try:
        response = requests.request("POST", url, headers=headers, data=payload, timeout=10000)
        return response.json().get('data')

    except Exception as e:
        json_url = f'dso/ai_gen/single/{task_id}/info.json'
        tos_util.put_string_to_tos(f'{json_url}',
                                   ["异常：你输入的文字不符合制作规则，请修改后重试"])

        print(f"下载失败：{str(e)}")


def download_image(url, save_path):
    try:
        # 发送 GET 请求获取图片内容
        response = requests.get(url, stream=True)
        response.raise_for_status()  # 检查请求是否成功

        # 提取保存路径中的目录部分
        save_dir = os.path.dirname(save_path)

        # 若目录不存在则创建（包括多级目录）
        if save_dir and not os.path.exists(save_dir):
            os.makedirs(save_dir, exist_ok=True)  # exist_ok=True 避免目录已存在时报错
        # 写入文件
        with open(save_path, 'wb') as file:
            for chunk in response.iter_content(chunk_size=1024):
                if chunk:  # 过滤空块
                    file.write(chunk)

        print(f"图片已成功下载至：{save_path}")
    except Exception as e:
        print(f"下载失败：{str(e)}")


def process_gen_image(prompt, ratio, task_id):
    # url_list = []
    # for i in range(4):
    #     url_list_result = gen_image(prompt, ratio)
    #     url_list.append(url_list_result[0])
    # url_list_json = []
    # for index, item in enumerate(url_list):
    #     save_path = f'{dir}{task_id}/{index}.jpeg'
    #     download_image(item['url'], save_path)
    #     tos_util.put_video(f'dso/ai_gen/single/{task_id}/{index}.jpeg', save_path)
    #     url_list_json.append(f'https://static2.douchacha.com/dso/ai_gen/single/{task_id}/{index}.jpeg')
    #
    # json_url = f'dso/ai_gen/single/{task_id}/info.json'
    # tos_util.put_string_to_tos(f'{json_url}', url_list_json)


    # 使用线程池并发执行4次请求
    url_list = []
    with ThreadPoolExecutor(max_workers=4) as executor:
        # 提交4个任务
        futures = [executor.submit(gen_image, prompt, ratio,task_id) for _ in range(4)]
        # 获取执行结果（按完成顺序）
        for future in as_completed(futures):
            url_list_result = future.result()
            url_list.append(url_list_result[0])

    # 后续逻辑保持不变
    url_list_json = []
    for index, item in enumerate(url_list):
        save_path = f'{dir}{task_id}/{index}.jpeg'
        download_image(item['url'], save_path)
        tos_util.put_video(f'dso/ai_gen/single/{task_id}/{index}.jpeg', save_path)
        url_list_json.append(f'https://static2.douchacha.com/dso/ai_gen/single/{task_id}/{index}.jpeg')

    json_url = f'dso/ai_gen/single/{task_id}/info.json'
    tos_util.put_string_to_tos(f'{json_url}', url_list_json)
    # os.remove(f'{dir}{task_id}/')


# # 保存路径（可自定义，这里保存为当前目录下的 image.jpg）
# save_path = f'{dir}image.jpeg'
#
# # 调用下载函数
#

if __name__ == '__main__':
    task = redis_client.rpop('dso_image_gen_task')
    try:
        task_json = json.loads(task)
        prompt = task_json['prompt']
        ratio = task_json['ratio']
        task_id = task_json['task_id']
        process_gen_image(prompt, ratio, task_id)
    except Exception as err:
        # redis_client.rpush('dso_image_gen', task)
        print(f'{err}')


        # prompt = '情趣内衣'
        # ratio = '16:9'
        # task_id = 'aaa'
        # process_gen_image(prompt, ratio, task_id)





