from fastapi import FastAPI, Response
from fastapi.middleware.cors import CORSMiddleware
from diffusers import StableDiffusionPipeline
import torch
from torch import autocast
from io import BytesIO
import base64
import datetime  # Import the datetime module
import uuid  # Import the uuid module
import logging
import os

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_credentials=True,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"]
)

device = "cuda"
model_id = "rowenac/vintagechinesecomics"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe.to(device)

# Define a folder path for saving images
image_folder = "image_gallery"

# Create the directory if it doesn't exist
os.makedirs(image_folder, exist_ok=True)

@app.get("/")
def generate(prompt: str):
    try:
        with autocast(device):
            image = pipe(prompt, guidance_scale=8.5).images[0]

        # Generate a unique filename for each image using a timestamp and UUID
        timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S%f")  # Format the timestamp
        image_uuid = uuid.uuid4()
        image_filename = f"{image_folder}/image_{timestamp}_{image_uuid}.png"

        # Save the image to the specified folder
        image.save(image_filename)
        logging.info(f"Image saved to {image_filename}")

        buffer = BytesIO()
        image.save(buffer, format="PNG")
        imgstr = base64.b64encode(buffer.getvalue()).decode('utf-8')

        return Response(content=imgstr, media_type="image/png")
    except Exception as e:
        logging.error(f"Error while generating image: {str(e)}")
        return Response(content=str(e), status_code=500)

# New API endpoint to fetch image filenames
@app.get("/api/get_images")
def get_images():
    # List all files in the image folder
    images = os.listdir(image_folder)
    return images

# New API endpoint to serve individual images
@app.get("/api/get_image/{image_name}")
def get_image(image_name: str):
    # Define the path to the image
    image_path = os.path.join(image_folder, image_name)

    # Open and serve the image
    with open(image_path, "rb") as image_file:
        return Response(content=image_file.read(), media_type="image/png")