import requests
import time
import json
import os
import uuid
from urllib.parse import urlparse

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 RequestImgAndSave(image_url):
    response = requests.get(image_url, stream=True)
    response.raise_for_status()

    # 获取原始图片格式
    ext = get_file_extension(image_url)
    filename = f"{uuid.uuid4()}{ext}"
    save_path = os.path.join(COMFYUI_INPUT_DIR, filename)

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

# 1. 提交任务

#RequestImgAndSave("https://example.com/image.jpg")