import runpod
import auth
import app_imagegen
import json
import uuid
import boto3  # Ensure you have boto3 installed for Cloudflare R2 interaction
from botocore.exceptions import NoCredentialsError
from botocore.client import Config
import os  # Import os to access environment variables
from datetime import datetime
import re

config = json.load(open("config.json"))
model_manager = app_imagegen.ModelManager(config, max_loaded_models=2, lazy_load=True)

def handler(job):
    job_input = job["input"]
    result = app_imagegen.image_gen_api_entry(
        job_input["prompt"],
        job_input.get("negative_prompt",""),
        job_input.get("styles",{}),
        job_input.get("options",{}),
        job_input.get("token",""),
        model_manager
    )
    user_id = job_input.get("user_id","")
    domain = job_input.get("domain","")
    # Save result as WebP using a GUID as filename
    filename = f"{uuid.uuid4()}.webp"
    with open(filename, "wb") as f:
        f.write(result)  # Assuming result is in bytes

    # Upload to Cloudflare R2
    r2_client = boto3.client('s3',
        endpoint_url=f"https://{config['r2']['account_id']}.r2.cloudflarestorage.com",
        aws_access_key_id=os.getenv('AWS_ACCESS_KEY_ID'),  # Read from environment variable
        aws_secret_access_key=os.getenv('AWS_SECRET_ACCESS_KEY'),  # Read from environment variable
        config=Config(signature_version='s3v4'),
        region_name='auto'
    )
    # Prepare metadata
    publish_time = datetime.now().isoformat()
    sanitized_prompt = re.sub(r'[^\x20-\x7E]', ' ', job_input["prompt"])  # Replace non-ASCII characters with space
    sanitized_prompt = sanitized_prompt[:512]  # Limit length to 256 characters
    metadata = {
        'domain': domain,
        'model': job_input.get("styles",{}).get("vibe","") ,
        'prompt': ''.join([i if ord(i) < 128 else ' ' for i in sanitized_prompt]),
        'author': user_id,
        'publish_time': str(publish_time)
    }
    new_url_path = ""
    try:
        r2_key = f"paid/{user_id}/{filename}"
        r2_client.put_object(
            Bucket=config['r2']['bucket_name'],
            Key=r2_key,
            Body=open(filename, 'rb'),
            ContentType='image/webp',  # Adjust this if the image type might vary
            Metadata=metadata
        )
        new_url_path =  f"/{r2_key}"
    except NoCredentialsError:
        print("Credentials not available")

    return {"image_url":new_url_path}


runpod.serverless.start({"handler": handler})