import os
import boto3
from botocore.client import Config
import uuid
from PIL import Image
import io

def upload_mask(image, prefix="mask"):
    """
    将分割掩码图像上传到DigitalOcean Spaces
    
    Args:
        image: PIL Image对象
        prefix: 文件名前缀
        
    Returns:
        上传文件的URL
    """
    try:
        # 从环境变量获取凭据
        do_key = os.environ.get('DO_SPACES_KEY')
        do_secret = os.environ.get('DO_SPACES_SECRET')
        do_region = os.environ.get('DO_SPACES_REGION')
        do_bucket = os.environ.get('DO_SPACES_BUCKET')
        
        # 校验凭据是否存在
        if not all([do_key, do_secret, do_region, do_bucket]):
            raise ValueError("缺少DigitalOcean Spaces凭据")
        
        # 创建S3客户端
        session = boto3.session.Session()
        client = session.client('s3',
                               region_name=do_region,
                               endpoint_url=f'https://{do_region}.digitaloceanspaces.com',
                               aws_access_key_id=do_key,
                               aws_secret_access_key=do_secret)
        
        # 生成唯一文件名
        filename = f"{prefix}_{uuid.uuid4().hex}.png"
        
        # 将图像转换为字节流
        img_byte_arr = io.BytesIO()
        image.save(img_byte_arr, format='PNG')
        img_byte_arr.seek(0)
        
        # 上传到Spaces
        client.upload_fileobj(
            img_byte_arr,
            do_bucket,
            filename,
            ExtraArgs={'ACL': 'public-read', 'ContentType': 'image/png'}
        )
        
        # 返回公共URL
        url = f'https://{do_bucket}.{do_region}.digitaloceanspaces.com/{filename}'
        return url
    
    except Exception as e:
        print(f"上传失败: {str(e)}")
        return f"上传错误: {str(e)}"