from fastapi import FastAPI, File, UploadFile, Form
from fastapi.responses import StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from typing import List
import io
from PIL import Image, ImageOps
import numpy as np
import compColors
import dominantColors
import recolorReinhardAlgo
import recolorOTAlgo
import recolorTransferAlgo
import recolorLumaConverterAlgo
import recolorPaletteBasedTransfer

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)
@app.post("/dominantColor/")
async def dominant_color(file: UploadFile = File(...), num_colors: int = Form(...)):
    """
    Receive an image file and an integer and return the dominant color(s).
    """
    print('num_colors: ', num_colors)
    file_content = await file.read()
    image_bytes = io.BytesIO(file_content)
    im = Image.open(image_bytes)
    dominantColorsRGB = dominantColors.findDominantColors(image_bytes, num_colors, False)
    dominantColorsHex = [dominantColors.rgb_to_hex(color) for color in dominantColorsRGB]
    return {"dominantColors": dominantColorsHex}

@app.post("/ColorPalettes/")
async def color_palettes(colors: str = Form(...)):
    """
    Receive an array of strings representing colors and return a color palette based on these colors.
    """
    #maybe this isn't necessary. converting the string to an array of strings
    colors = [color.strip() for color in colors.split(',')]
    #generate the first pallete, which is the complementary colors of the given colors
    complementaryColors = []
    for color in colors:
        complementaryColors.append(compColors.complementary_colors(color))

    #generate the second palette using the adjacent colors algorithm:
    adjacentColors = []
    for color in colors:
        _adjcolors = compColors.adjacent_colors(color)
        for _color in _adjcolors:
            if _color not in adjacentColors:
                adjacentColors.append(_color)
    #generate the third palette using the gradient colors algorithm:
    gradientColors = []
    for i in range(len(colors)-1):
        gradientColors.append(compColors.gradient_colors(colors[i], colors[i+1]))

    #Fixing size of palletes to 5 colors:
    complementaryColors = [complementaryColors[i:i + 5] for i in range(0, len(complementaryColors), 5)]
    adjacentColors = [adjacentColors[i:i + 5] for i in range(0, len(adjacentColors), 5)]
    colors = [colors[i:i + 5] for i in range(0, len(colors), 5)]
    
    return {"inputColor": colors, "complementaryColors": complementaryColors, "adjacentColors": adjacentColors, "gradientColors": gradientColors}

@app.post("/recolor/")
async def recolor(file: UploadFile = File(...), colors: str = Form(...),  model: str = Form(...)):
    """
    Receive an image file and an array of strings representing colors of a selected pallete and recolor an image.
    """
    method = model
    invertColors = False
    colors = [color.strip() for color in colors.split(',')]
    file_content = await file.read()
    image_bytes = io.BytesIO(file_content)
    image = Image.open(image_bytes)
    if invertColors:
        image = ImageOps.invert(image)
    image_np = np.array(image)

    if method == "CCA":
        #Characteristic Color Analysis
        recolorReinhardAlgo.recolor(image_np, colors)
    elif method == "OTA":
        #Optimal Transport Algorithm transfer
        recolorOTAlgo.recolor(image_np, colors)
    elif method =="KMEANS":
        #K-means clustering transfer
        recolorTransferAlgo.recolor(image_np, colors)
    elif method == "LUMA":
        #Luma converter transfer
        recolorLumaConverterAlgo.remap_image_colors(image_np, colors)
    elif method == "palette":
        #palette transfer
        recolorPaletteBasedTransfer.recolor(image_np, colors)
    img_file = open("./result.jpg", "rb")
    return StreamingResponse(img_file, media_type="image/jpeg")
@app.get("/test/")
async def test():
    """
    Test endpoint to check if the server is running.
    """
    return {"message": "Server is running"}





if __name__ == "__main__":
    import uvicorn
    print("Server is running")
    uvicorn.run(app, host="0.0.0.0", port=7860)




#how to run:
#source env/bin/activate     
#uvicorn server:app --reload