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

def upload_mask(image, prefix="mask"):
    """
    Upload the segmentation mask image to DigitalOcean Spaces
    
    Args:
        image: PIL Image object
        prefix: filename prefix
        
    Returns:
        URL of the uploaded file
    """
    try:
        # Get credentials from environment variables
        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')
        
        # Validate that credentials exist
        if not all([do_key, do_secret, do_region, do_bucket]):
            raise ValueError("Missing DigitalOcean Spaces credentials")
        
        # Create S3 client
        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)
        
        # Generate a unique filename
        filename = f"{prefix}_{uuid.uuid4().hex}.png"
        
        # Convert the image to a byte stream
        img_byte_arr = io.BytesIO()
        image.save(img_byte_arr, format='PNG')
        img_byte_arr.seek(0)
        
        # Upload to Spaces
        client.upload_fileobj(
            img_byte_arr,
            do_bucket,
            filename,
            ExtraArgs={'ACL': 'public-read', 'ContentType': 'image/png'}
        )
        
        # Return the public URL
        url = f'https://{do_bucket}.{do_region}.digitaloceanspaces.com/{filename}'
        return url
    
    except Exception as e:
        print(f"Upload failed: {str(e)}")
        return f"Upload error: {str(e)}"