import datetime

from fastapi import FastAPI, File, UploadFile
from PIL import Image
import os
import io
from typing import List
from pydantic import BaseModel
import random
import glob
import hashlib

app = FastAPI()

BACKGROUND_DIR = 'backgrounds'
IMAGE_DIR = 'images'
IMAGE_EXT = '.png'


class MosaicRequest(BaseModel):
    rows: int
    columns: int
    buffer_size: int
    keywords: List[str]
    background_file: str

class BackgroundRequest(BaseModel):
    backgrounds: List[str]


def get_images():
    image_files = [f for f in os.listdir(IMAGE_DIR) if f.endswith(IMAGE_EXT)]
    image_files.sort()
    return [Image.open(os.path.join(IMAGE_DIR, f)) for f in image_files]


def create_mosaic(images, rows, cols, buffer_size):
    mosaic_width = cols * images[0].size[0] + (cols - 1) * buffer_size
    mosaic_height = rows * images[0].size[1] + (rows - 1) * buffer_size
    mosaic = Image.new('RGBA', (mosaic_width, mosaic_height), color=(0, 0, 0, 0))
    for i, img in enumerate(images):
        row = i // cols
        col = i % cols
        offset_x = (images[0].size[0] - img.size[0]) // 2
        offset_y = (images[0].size[1] - img.size[1]) // 2
        pos_x = col * (img.size[0] + buffer_size)
        pos_y = row * (img.size[1] + buffer_size)
        pos = (pos_x + offset_x, pos_y + offset_y)
        mosaic.paste(img, pos)
    return mosaic


def place_image_on_background(foreground: Image.Image, background: Image.Image) -> Image.Image:
    # Compute the position to paste the foreground image
    x = (background.width - foreground.width) // 2
    y = (background.height - foreground.height) // 2

    # Create a new image with the background and paste the foreground onto it
    result = Image.new("RGBA", (background.width, background.height), (255, 255, 255, 0))
    result.paste(background, (0, 0))
    result.alpha_composite(foreground, (x, y))

    return result

@app.get("/backgrounds/")
async def list_backgrounds():
    background_files = [f for f in os.listdir(BACKGROUND_DIR) if f.endswith(IMAGE_EXT)]
    background_files.sort()
    background_dict = {"backgrounds": background_files}
    return BackgroundRequest(**background_dict)

@app.post("/upload_file/")
async def upload_file(file: UploadFile = File(...), directory=IMAGE_DIR):
    # check if the file is an image file
    if not file.content_type.startswith('image'):
        raise ValueError('File is not an image')

    # read the file data and create a Pillow Image object
    image = Image.open(file.file)

    # save the PNG image to the "images/" directory
    filename = os.path.join(directory, os.path.splitext(file.filename)[0] + '.png')
    image.save(filename)

    return {"filename": filename}


@app.post("/upload_background")
async def upload_background(file: UploadFile = File(...)):
    filename_d = await upload_file(file, directory=BACKGROUND_DIR)
    return filename_d


@app.post("/generate_mosaic")
async def generate_mosaic(request: MosaicRequest):
    images = []
    for keyword in request.keywords:
        file_list = glob.glob(f'images/{keyword}*.png')
        if not file_list:
            return {"message": f"No images found for keyword: {keyword}"}
        selected_file = random.choice(file_list)
        with open(selected_file, "rb") as f:
            images.append(Image.open(io.BytesIO(f.read())))

    # select background image
    background = Image.open(f"backgrounds/{request.background_file}")

    # resize images to have the same height
    height = min(background.height - (request.rows + 1) * request.buffer_size,
                 min(img.size[1] for img in images)) / request.rows
    images = [img.resize((round(img.size[0]*height/img.size[1]),
                          round(height))) for img in images]

    # create grid
    grid = create_mosaic(images, request.rows, request.columns, request.buffer_size)

    # select background image
    background = Image.open(f"backgrounds/{request.background_file}")

    # place grid on background
    final_image = place_image_on_background(grid, background)

    # convert image to bytes
    img_bytes = final_image.tobytes()

    # create MD5 hash object
    hash_obj = hashlib.md5()

    # update hash object with image bytes
    hash_obj.update(img_bytes)

    # get hex representation of hash
    hash_hex = hash_obj.hexdigest()

    # convert hex to string
    hash_str = str(hash_hex)

    # save final image
    now = datetime.datetime.now()
    filename_datetime = now.strftime("%Y-%m-%d_%H-%M-%S")

    final_image.save(f"mosaics/{hash_str}-{filename_datetime}.png")

    return {"message": "Mosaic image created successfully"}