import gradio as gr
import torch
from diffusers import FluxFillPipeline
from diffusers.utils import load_image
from PIL import Image, ImageDraw
import numpy as np
import spaces
import requests

# Model setup
pipe = FluxFillPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-Fill-dev",
    torch_dtype=torch.bfloat16
).to("cuda")

# Translation function
@spaces.GPU
def translate_albanian_to_english(text):
    if not text.strip():
        return ""
    for attempt in range(2):
        try:
            response = requests.post(
                "https://hal1993-mdftranslation1234567890abcdef1234567890-fc073a6.hf.space/v1/translate",
                json={"from_language": "sq", "to_language": "en", "input_text": text},
                headers={"accept": "application/json", "Content-Type": "application/json"},
                timeout=5
            )
            response.raise_for_status()
            translated = response.json().get("translate", "")
            return translated
        except Exception as e:
            if attempt == 1:
                raise gr.Error(f"Përkthimi dështoi: {str(e)}")
    raise gr.Error("Përkthimi dështoi. Ju lutem provoni përsëri.")

# Aspect ratio function
def update_aspect_ratio(ratio):
    if ratio == "1:1":
        return 640, 640
    elif ratio == "9:16":
        width = 512
        height = int(round(512 * 16 / 9 / 8)) * 8  # Round to nearest multiple of 8
        return width, height
    elif ratio == "16:9":
        width = int(round(512 * 16 / 9 / 8)) * 8  # Round to nearest multiple of 8
        height = 512
        return width, height
    return 640, 640  # Default to 1:1

# Core processing functions
def can_expand(source_width, source_height, target_width, target_height, alignment):
    if alignment in ("Left", "Right") and source_width >= target_width:
        return False
    if alignment in ("Top", "Bottom") and source_height >= target_height:
        return False
    return True

def prepare_image_and_mask(image, width, height, overlap_percentage, resize_option, alignment):
    if image is None:
        raise gr.Error("Ju lutem ngarkoni një imazh.")

    target_size = (width, height)

    # Resize image based on scale factor
    scale_factor = min(target_size[0] / image.width, target_size[1] / image.height)
    new_width = int(image.width * scale_factor)
    new_height = int(image.height * scale_factor)
    source = image.resize((new_width, new_height), Image.LANCZOS)

    # Map resize_option to percentage
    resize_map = {
        "E Plotë": 100,
        "75%": 75,
        "50%": 50,
        "33%": 33,
        "25%": 25
    }
    resize_percentage = resize_map.get(resize_option, 75)  # Default to 75% if invalid

    # Apply resize percentage
    resize_factor = resize_percentage / 100
    new_width = int(source.width * resize_factor)
    new_height = int(source.height * resize_factor)
    new_width = max(new_width, 64)  # Ensure minimum size
    new_height = max(new_height, 64)
    source = source.resize((new_width, new_height), Image.LANCZOS)

    # Calculate overlap in pixels
    overlap_x = int(new_width * (overlap_percentage / 100))
    overlap_y = int(new_height * (overlap_percentage / 100))
    overlap_x = max(overlap_x, 1)
    overlap_y = max(overlap_y, 1)

    # Calculate margins based on alignment
    if alignment == "Middle":
        margin_x = (target_size[0] - new_width) // 2
        margin_y = (target_size[1] - new_height) // 2
    elif alignment == "Left":
        margin_x = 0
        margin_y = (target_size[1] - new_height) // 2
    elif alignment == "Right":
        margin_x = target_size[0] - new_width
        margin_y = (target_size[1] - new_height) // 2
    elif alignment == "Top":
        margin_x = (target_size[0] - new_width) // 2
        margin_y = 0
    elif alignment == "Bottom":
        margin_x = (target_size[0] - new_width) // 2
        margin_y = target_size[1] - new_height

    margin_x = max(0, min(margin_x, target_size[0] - new_width))
    margin_y = max(0, min(margin_y, target_size[1] - new_height))

    # Create background and paste source image
    background = Image.new('RGB', target_size, (255, 255, 255))
    background.paste(source, (margin_x, margin_y))

    # Create mask
    mask = Image.new('L', target_size, 255)
    mask_draw = ImageDraw.Draw(mask)
    white_gaps_patch = 2
    left_overlap = margin_x + overlap_x
    right_overlap = margin_x + new_width - overlap_x
    top_overlap = margin_y + overlap_y
    bottom_overlap = margin_y + new_height - overlap_y

    if alignment == "Left":
        left_overlap = margin_x
    elif alignment == "Right":
        right_overlap = margin_x + new_width
    elif alignment == "Top":
        top_overlap = margin_y
    elif alignment == "Bottom":
        bottom_overlap = margin_y + new_height

    mask_draw.rectangle([
        (left_overlap, top_overlap),
        (right_overlap, bottom_overlap)
    ], fill=0)

    return background, mask

@spaces.GPU
def inpaint(image, width, height, overlap_percentage, num_inference_steps, resize_option, prompt, progress=gr.Progress(track_tqdm=True)):
    # Translate Albanian prompt to English
    final_prompt = translate_albanian_to_english(prompt.strip()) if prompt.strip() else ""

    # Prepare image and mask
    background, mask = prepare_image_and_mask(
        image, width, height, overlap_percentage, resize_option, alignment="Middle"
    )

    # Check if expansion is possible
    if not can_expand(background.width, background.height, width, height, "Middle"):
        alignment = "Middle"

    # Create control image
    cnet_image = background.copy()
    cnet_image.paste(0, (0, 0), mask)

    # Run inpainting
    try:
        result = pipe(
            prompt=final_prompt,
            height=height,
            width=width,
            image=cnet_image,
            mask_image=mask,
            num_inference_steps=num_inference_steps,
            guidance_scale=50,
        ).images[0]
    except Exception as e:
        raise gr.Error(f"Gabim gjatë gjenerimit të imazhit: {str(e)}")

    # Combine result with control image
    result = result.convert("RGBA")
    cnet_image.paste(result, (0, 0), mask)

    return np.array(cnet_image)  # Return as NumPy array for gr.Image

# Gradio interface
def create_demo():
    with gr.Blocks() as demo:
        # CSS for 320px gap, download button scaling, and container width constraint
        gr.HTML("""
        <style>
        body::before {
            content: "";
            display: block;
            height: 320px;
            background-color: var(--body-background-fill);
        }
        button[aria-label="Fullscreen"], button[aria-label="Fullscreen"]:hover {
            display: none !important;
            visibility: hidden !important;
            opacity: 0 !important;
            pointer-events: none !important;
        }
        button[aria-label="Share"], button[aria-label="Share"]:hover {
            display: none !important;
        }
        button[aria-label="Download"] {
            transform: scale(3);
            transform-origin: top right;
            margin: 0 !important;
            padding: 6px !important;
        }
        .constrained-container {
            max-width: 600px; /* Limits container width */
            margin: 0 auto; /* Centers the container */
        }
        </style>
        """)

        gr.Markdown("# Zgjeroni Imazhin")
        gr.Markdown("Zgjeroni imazhin duke plotësuar sfondin bazuar në përshkrimin e dhënë")

        with gr.Row():
            with gr.Column(elem_classes="constrained-container"):
                input_image = gr.Image(sources='upload', type="pil", label="Imazhi i Ngarkuar", height=480, width=480)
                prompt = gr.Textbox(label="Përshkrimi", placeholder="Shkruani përshkrimin këtu (opsionale)")
                aspect_ratio = gr.Radio(choices=["9:16", "1:1", "16:9"], value="1:1", label="Raporti i Aspektit")
                resize_option = gr.Radio(
                    choices=["E Plotë", "75%", "50%", "33%", "25%"],
                    value="75%",
                    label="Madhësia e Imazhit të Hyrjes",
                    info="Zgjidhni sa i madh të jetë imazhi i hyrjes në kanavacën përfundimtare"
                )
                generate_button = gr.Button(value="Gjenero")
                result_image = gr.Image(label="Rezultati", type="numpy", height=480, width=480, elem_classes="constrained-container")
                # Hidden components for processing
                width_slider = gr.Slider(label="Gjerësia e Synuar", minimum=256, maximum=1536, value=640, step=8, visible=False)
                height_slider = gr.Slider(label="Lartësia e Synuar", minimum=256, maximum=1536, value=640, step=8, visible=False)
                overlap_percentage = gr.Slider(label="Përqindja e Mbivendosjes", minimum=1, maximum=50, value=10, step=1, visible=False)
                num_inference_steps = gr.Slider(label="Hapat", minimum=2, maximum=50, value=28, step=1, visible=False)

        # Update hidden sliders based on aspect ratio
        aspect_ratio.change(
            fn=update_aspect_ratio,
            inputs=[aspect_ratio],
            outputs=[width_slider, height_slider]
        )

        # Bind the generate button
        inputs = [
            input_image, width_slider, height_slider, overlap_percentage, num_inference_steps,
            resize_option, prompt
        ]
        generate_button.click(
            fn=inpaint,
            inputs=inputs,
            outputs=[result_image]
        )

    return demo

if __name__ == "__main__":
    print(f"Gradio version: {gr.__version__}")
    app = create_demo()
    app.queue(max_size=12).launch(server_name='0.0.0.0')