import gradio as gr
import requests
import io
import random
import os
import time
import numpy as np
from PIL import Image, ImageEnhance
from deep_translator import GoogleTranslator
import json
from random import randint

# API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-schnell"
API_URL = "black-forest-labs/FLUX.1-schnell"

mod_list = {
    "FLUX.1 Schnell": API_URL,
    "FLUX.1 Schnell | Face Realism": "prithivMLmods/Canopus-LoRA-Flux-FaceRealism",
    "FLUX.1 Schnell | Midjourney": "Jovie/Midjourney_Schnell",
    "FLUX.1 Dev": "black-forest-labs/FLUX.1-dev",
    "FLUX.1 Schnell | realism": "hugovntr/flux-schnell-realism",
    "FLUX.1 Schnell | MJ v6": "strangerzonehf/Flux-Midjourney-Mix2-LoRA",
    "FLUX.1 Dev | Flux Realism": "XLabs-AI/flux-RealismLora",
    "FLUX.1 Redux": "black-forest-labs/FLUX.1-Redux-dev",
    "FLUX.1 Dev | UltraRealism 2.0": "prithivMLmods/Canopus-LoRA-Flux-UltraRealism-2.0",
    "Lumina Image 2.0": "Alpha-VLLM/Lumina-Image-2.0",
    "FLUX.1 Dev | Vector Journey": "Shakker-Labs/FLUX.1-dev-LoRA-Vector-Journey",
    "FLUX.1 Dev | Deep Blue": "fffiloni/deep-blue-v2",
    "Realism Engine v1.0": "digiplay/RealismEngine_v1",
    "Absolute Reality v1.8.1": "digiplay/AbsoluteReality_v1.8.1",
    "FLUX.1 Dev | Midjourney Anime": "brushpenbob/flux-midjourney-anime",
    "FLUX.1 Dev | Pencil v2": "brushpenbob/flux-pencil-v2",
    "FLUX.1 Dev | Add Details": "Shakker-Labs/FLUX.1-dev-LoRA-add-details",
    "FLUX.1 Dev | Pastel Anime": "Raelina/Flux-Pastel-Anime",
    "FLUX.1 Dev | CartoonStyle": "Norod78/CartoonStyle-flux-lora"
}
API_TOKEN = os.getenv("HF_READ_TOKEN")
headers = {"Authorization": f"Bearer {API_TOKEN}"}
timeout = 100

def add_noise(image, intensity=25):
    """
    Добавляет цифровой шум к изображению.
    :param image: Изображение (PIL.Image)
    :param intensity: Интенсивность шума (от 0 до 255)
    :return: Изображение с шумом (PIL.Image)
    """
    # Преобразуем изображение в массив NumPy
    img_array = np.array(image)
    
    # Генерируем шум
    noise = np.random.randint(-intensity, intensity, img_array.shape, dtype=np.int32)
    
    # Добавляем шум к изображению
    noisy_array = np.clip(img_array + noise, 0, 255).astype(np.uint8)
    
    # Преобразуем массив обратно в изображение
    noisy_image = Image.fromarray(noisy_array)
    return noisy_image

def resize_and_crop(image, target_width=768, target_height=1024):
    """
    Подгоняет изображение под размер target_width x target_height с сохранением пропорций.
    Если изображение не соответствует соотношению сторон, обрезает его по бокам.
    :param image: Изображение (PIL.Image)
    :param target_width: Целевая ширина
    :param target_height: Целевая высота
    :return: Изображение с измененным размером (PIL.Image)
    """
    # Сохраняем исходные пропорции
    original_width, original_height = image.size
    target_ratio = target_width / target_height
    original_ratio = original_width / original_height

    # Масштабируем изображение с сохранением пропорций
    if original_ratio > target_ratio:
        # Если изображение шире, чем нужно, обрезаем по бокам
        new_height = target_height
        new_width = int(original_width * (target_height / original_height))
        resized_image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
        left = (new_width - target_width) / 2
        top = 0
        right = (new_width + target_width) / 2
        bottom = target_height
    else:
        # Если изображение уже, чем нужно, обрезаем сверху и снизу
        new_width = target_width
        new_height = int(original_height * (target_width / original_width))
        resized_image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
        left = 0
        top = (new_height - target_height) / 2
        right = target_width
        bottom = (new_height + target_height) / 2

    # Обрезаем изображение до целевого размера
    cropped_image = resized_image.crop((left, top, right, bottom))
    return cropped_image

def compress_image(image, quality=50):
    """
    Сжимает изображение и понижает его качество в формате JPEG.
    :param image: Изображение (PIL.Image)
    :param quality: Качество JPEG (от 1 до 100)
    :return: Сжатое изображение (PIL.Image)
    """
    # Сохраняем изображение в буфер с пониженным качеством
    buffer = io.BytesIO()
    image.save(buffer, format="JPEG", quality=quality)
    buffer.seek(0)
    
    # Загружаем изображение обратно из буфера
    compressed_image = Image.open(buffer)
    return compressed_image

def query(prompt, is_realistic, num_inference_steps, width, height, mod = None):
    if prompt == "" or prompt == None:
        return None

    key = random.randint(0, 999)
    
    API_TOKEN = random.choice([os.getenv("HF_READ_TOKEN")])
    headers = {"Authorization": f"Bearer {API_TOKEN}"}
    
    payload = {
        "inputs": prompt,
        "seed": random.randint(1, 1000000000),
        "parameters": {"num_inference_steps": num_inference_steps, "width": 1024, "height": 1024}
    }

    model = API_URL
    for mod_name, mod_link in mod_list.items():
        if (mod == mod_name):
            model = mod_link

    model = "https://api-inference.huggingface.co/models/" + model;
    
    response = requests.post(model, headers=headers, json=payload, timeout=timeout)
    if response.status_code != 200:
        print(f"Error: Failed to get image. Response status: {response.status_code}")
        print(f"Response content: {response.text}")
        if response.status_code == 503:
            raise gr.Error(f"{response.status_code} : The model is being loaded")
        raise gr.Error(f"{response.status_code}")
    
    try:
        image_bytes = response.content
        image = Image.open(io.BytesIO(image_bytes))

        image = resize_and_crop(image, width, height)

        if (is_realistic == True):
            image = add_noise(image, intensity = randint(100, 150) / 10)
            image = compress_image(image, randint(70, 80))

            enhancer = ImageEnhance.Contrast(image)
            image = enhancer.enhance(randint(75,80) / 100)
            
            # Изменение насыщенности
            enhancer = ImageEnhance.Color(image)
            image = enhancer.enhance(randint(70,80) / 100)
    
            # Изменение экспозиции (яркости)
            enhancer = ImageEnhance.Brightness(image)
            image = enhancer.enhance(randint(70,100) / 100)

        
        print(f'\033[1mGeneration {key} completed!\033[0m ({prompt})')
        return image
    except Exception as e:
        print(f"Error when trying to open the image: {e}")
        return None

css = """
#app-container {
    max-width: 600px;
    margin-left: auto;
    margin-right: auto;
}
"""

test_prompt = """[amateur perspective of a house],
[on this house staying white russian female (if you can see it) with white t-shirt (if you can see it) on her body, with black jeans (if you can see it) on her legs, with lime sneakers (if you can see it) on her feet and looks aside (if you can see it)]
[night melancholic lighting]"""

with gr.Blocks(theme='gstaff/xkcd', css=css) as app:
    gr.HTML("<center><h1>FLUX.1-Schnell</h1></center>")
    with gr.Column(elem_id="app-container"):
        with gr.Row():
            with gr.Column(elem_id="prompt-container"):
                text_prompt = gr.Textbox(label="Prompt", placeholder="Enter a prompt here", lines=2, elem_id="prompt-text-input", value=test_prompt)
                is_realistic = gr.Checkbox(label="Realistic filter", value=True)
                num_inference_steps = gr.Slider(label="Number of inference steps", minimum=1, maximum=50, step=1, value=4)
                with gr.Row():
                    width = gr.Slider(label="Width", minimum = 128, maximum = 1024, step = 32, value = 480)
                    height = gr.Slider(label="Height", minimum = 128, maximum = 1024, step = 32, value = 640)
                mod_choice = gr.Dropdown(list(mod_list.keys()), label="Mod")
                    
        with gr.Row():
            text_button = gr.Button("Run", variant='primary', elem_id="gen-button")
        with gr.Row():
            image_output = gr.Image(type="pil", label="Image Output", elem_id="gallery")
        
        text_button.click(query, inputs=[text_prompt, is_realistic, num_inference_steps, width, height, mod_choice], outputs=image_output)

app.launch(show_api=True, share=True)