import os
import uuid
import requests
from flask import Flask, request, jsonify
from urllib.parse import urlparse
from volcenginesdkarkruntime import Ark
import json
import base64
import time
from PIL import Image
from datetime import datetime

APP_ROOT = os.path.dirname(os.path.abspath(__file__))
STATIC_FOLDER = os.path.join(APP_ROOT, 'static')

app = Flask(__name__)

# 请确保您已将 API Key 存储在环境变量 ARK_API_KEY 中
# 初始化Ark客户端，从环境变量中读取您的API Key
client = Ark(
    # 此为默认路径，您可根据业务所在地域进行配置
    base_url="https://ark.cn-beijing.volces.com/api/v3",
    # 从环境变量中获取您的 API Key。此为默认方式，您可根据需要进行修改
    # api_key=os.environ.get("ARK_API_KEY"),
    api_key='14fc0280-fc65-462d-ac2d-50178c0212e3'
)

def parse_facial_data(data_str):
    print(f'原始字符串 {data_str}')
    # 去除多余的标记和换行符
    json_str_clean = data_str.strip('```json\n').strip('```').strip()

    # 转换为Python字典
    data = json.loads(json_str_clean)
    s = data['面部年龄']
    s = s.replace('岁', '')
    parts = s.split('-')
    start_age = int(parts[0])
    end_age = int(parts[1])
    data['start_age'] = start_age
    data['end_age'] = end_age
    return data

def GetPicDesc(img_url):
    response = client.chat.completions.create(
        # 指定您创建的方舟推理接入点 ID，此处已帮您修改为您的推理接入点 ID
        model="doubao-1.5-vision-pro-250328",
        messages=[
            {
                "role": "user",
                "content": [
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": img_url
                        },
                    },
                    {"type": "text", "text": "告诉我图片中的服装长度, 只要答案, 选项有(腰部以上、腰、跨、大腿上部1/3、大腿中部、膝盖、小腿、小腿上部1/3、脚踝、拖地、难以辨认)"},
                ],
            }
        ],
        
    )

    text = response.choices[0].message.content
    result = {}

    try:
        result = parse_facial_data(text)
    except:
        print(f'Error {text}')

    return result


def change(human_name, cloth_name, c_width, c_height):
    with open('/home/szlc/code/ComfyUI/change_cloth/change_new.json', 'r', encoding='utf-8') as file:
        prompt_text = file.read()

    prompt = json.loads(prompt_text)

    prompt["99"]["inputs"]["width"] = int((c_width/c_height) * 1024)

    #input cloth img
    prompt["22"]["inputs"]["image"] = cloth_name

    #input human img
    prompt["61"]["inputs"]["image"] = human_name

    #out put name
    out_img_name = str(uuid.uuid4())[:8]
    prompt["102"]["inputs"]["filename_prefix"] = out_img_name

    p = {"prompt": prompt}

    data = json.dumps(p).encode('utf-8')

    response = requests.post("http://localhost:8188/prompt", data=data)
    prompt_id = response.json()["prompt_id"]


    # 2. 轮询队列，直到任务完成
    while True:
        queue = requests.get("http://localhost:8188/queue").json()
        # print(queue)
        if not queue["queue_running"] and not queue["queue_pending"]:
            break  # 队列为空，任务已完成
        time.sleep(0.5)  # 避免频繁请求

    # 3. 从历史记录中获取结果
    history = requests.get("http://localhost:8188/history").json()
    print("History:", history)
    outputs = history[prompt_id]["outputs"]
    for out in outputs:
        if 'images' in outputs[out]:
            for out_img_name in outputs[out]['images']:
                if 'filename' in out_img_name:
                    return out_img_name['filename']
    return None

def save_base64_image(base64_str, prefix):
    """
    将base64字符串保存为图片文件
    :param base64_str: 带有类型标识的base64字符串
    :param prefix: 文件名前缀
    :return: 保存的文件路径
    """
    try:
        # 分离base64头部和实际数据
        header, data = base64_str.split(',', 1)
        
        # 从头部获取文件扩展名
        file_ext = ''
        if 'image/png' in header:
            file_ext = '.png'
        elif 'image/jpeg' in header:
            file_ext = '.jpg'
        elif 'image/jpg' in header:
            file_ext = '.jpg'
        elif 'image/gif' in header:
            file_ext = '.gif'
        else:
            file_ext = '.png'  # 默认使用png
        
        # 生成唯一文件名
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        unique_id = str(uuid.uuid4())[:8]
        filename = f"{prefix}_{timestamp}_{unique_id}{file_ext}"
        filepath = os.path.join('/home/szlc/code/ComfyUI/change_cloth/static', filename)
        
        # 解码并保存图片
        with open(filepath, 'wb') as f:
            f.write(base64.b64decode(data))

        input_filepath = os.path.join('/home/szlc/code/ComfyUI/input', filename)
        # 解码并保存图片
        with open(input_filepath, 'wb') as f:
            f.write(base64.b64decode(data))
            
        return filename
    
    except Exception as e:
        print(f"Error saving image: {str(e)}")
        return None
    
def get_image_dimensions(image_path):
    try:
        with Image.open(image_path) as img:
            width, height = img.size
            return width, height
    except Exception as e:
        print(f"Error: {e}")
        return None
    
def image_to_base64(file_path, mime_type=None):
    if mime_type is None:
        extension = file_path.split('.')[-1].lower()
        mime_types = {
            'jpg': 'image/jpeg',
            'jpeg': 'image/jpeg',
            'png': 'image/png',
            'gif': 'image/gif',
            'webp': 'image/webp',
            'bmp': 'image/bmp'
        }
        mime_type = mime_types.get(extension, 'application/octet-stream')
    
    # 读取文件内容并编码为Base64
    with open(file_path, 'rb') as image_file:
        encoded_string = base64.b64encode(image_file.read()).decode('utf-8')
    
    # 组合成Data URI格式
    return f"data:{mime_type};base64,{encoded_string}"

def process_change_cloth(human_filename, cloth_filename, output_format):
    w,h = get_image_dimensions(f'/home/szlc/code/ComfyUI/input/{human_filename}') 

    out_put_name = change(human_filename, cloth_filename, w, h)
    if out_put_name == None:
        return jsonify({"ret":-1, 'msg': 'Failed to change cloth'}), 500
    
    image = Image.open(f'/home/szlc/code/ComfyUI/output/{out_put_name}')
    jpg_name = out_put_name.replace(".png", ".jpg")
    jpg_path_name = f'/home/szlc/code/ComfyUI/change_cloth/static/{jpg_name}'
    image.save(jpg_path_name, quality=95)

    if 'base64' in output_format:
        return jsonify({
            "ret":0,
            "msg":"success",
            "data":image_to_base64(jpg_path_name)
        })
    else:
        return jsonify({
            "ret":0,
            "msg":"success",
            "url":f"http://112.126.94.241:18018/static/{jpg_name}"
        })

@app.route('/change_cloth_base64', methods=['POST'])
def change_cloth_base64():
    # 获取参数
    data = request.get_json()
    if not data:
        return jsonify({"ret":-1,'msg': 'No JSON data provided'}), 400
    
    human_img = data.get('human_img')
    cloth_img = data.get('cloth_img')
    output_format = data.get('output_format')
    
    if not human_img or not cloth_img:
        return jsonify({"ret":-1, 'msg': 'Both human_img and cloth_img are required'}), 400
    
    try:
        
        if 'base64' in output_format:
            return jsonify({
                    "ret":0,
                    "msg":"success",
                    "data":image_to_base64("/home/szlc/code/ComfyUI/input/girl_full.jpg")
                })
        else:
            return jsonify({
                "ret":0,
                "msg":"success",
                "url":f"http://43.143.205.217/imgs/processed_f60d0a9862834e1c92bc1c64ecf8f3ba.jpg"
            })
        
    except Exception as e:
        return jsonify({"ret":-1, 'error': str(e)}), 500
    
def get_file_extension(url_or_filename):
    """从 URL 或文件名中提取扩展名（如 .jpg、.png）"""
    # 处理 URL 情况（如 https://example.com/image.jpg?width=200）
    if url_or_filename.startswith(('http://', 'https://')):
        parsed = urlparse(url_or_filename)
        path = parsed.path
    else:
        path = url_or_filename

    # 提取扩展名（转换为小写，去掉问号后的参数）
    ext = os.path.splitext(path)[1].lower().split('?')[0]
    return ext if ext else '.png'  # 默认 PNG（如果无法提取）

def save_image_from_url(image_url):
    """从 URL 下载图片并保留原始格式"""
    try:
        response = requests.get(image_url, stream=True)
        response.raise_for_status()

        # 获取原始图片格式
        ext = get_file_extension(image_url)
        filename = f"{uuid.uuid4()}{ext}"
        filepath = os.path.join('/home/szlc/code/ComfyUI/change_cloth/static', filename)

        with open(filepath, "wb") as f:
            for chunk in response.iter_content(1024):
                f.write(chunk)

        input_filepath = os.path.join('/home/szlc/code/ComfyUI/input', filename)
        with open(input_filepath, "wb") as f:
            for chunk in response.iter_content(1024):
                f.write(chunk)

        return filename
    except Exception as e:
        print(f"Error saving image from URL: {e}")
        return None
    
@app.route('/change_cloth', methods=['POST'])
def change_cloth():
    """从 URL 下载图片"""
    data = request.json
    human_url = data.get("human_url")

    if not human_url:
        return jsonify({"error": "Missing 'human_url' parameter"}), 400

    
    cloth_url = data.get("cloth_url")

    if not cloth_url:
        return jsonify({"error": "Missing 'cloth_url' parameter"}), 400
    
    output_format = data.get('output_format')
    if not output_format:
        return jsonify({"error": "Missing 'output_format' parameter"}), 500

    if 'base64' in output_format:
        return jsonify({
            "ret":0,
            "msg":"success",
            "data":image_to_base64("/home/szlc/code/ComfyUI/input/girl_full.jpg")
        })
    else:
        return jsonify({
            "ret":0,
            "msg":"success",
            "url":f"http://43.143.205.217/imgs/processed_f60d0a9862834e1c92bc1c64ecf8f3ba.jpg"
        })

if __name__ == '__main__':
    app.run(host="0.0.0.0", port=8888, debug=False)