import spaces
import gradio as gr
import torch
from PIL import Image
from torchvision import transforms
from transformers import AutoModelForImageSegmentation

torch.set_float32_matmul_precision("high")

print("Loading BiRefNet...")
birefnet = AutoModelForImageSegmentation.from_pretrained(
    "ZhengPeng7/BiRefNet", trust_remote_code=True
)
birefnet.to("cuda")
birefnet.eval()
birefnet.half()
print("Model ready.")

IMAGE_SIZE = (1024, 1024)

transform_image = transforms.Compose([
    transforms.Resize(IMAGE_SIZE),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])


@spaces.GPU(duration=30)
def cutout_model(image: Image.Image):
    """Removes the background from the model photo, returns RGBA cutout."""
    original = image.convert("RGB")
    input_tensor = transform_image(original).unsqueeze(0).to("cuda").half()

    with torch.no_grad():
        preds = birefnet(input_tensor)[-1].sigmoid().cpu()

    mask = preds[0].squeeze()
    mask_pil = transforms.ToPILImage()(mask).resize(original.size)

    cutout = original.copy()
    cutout.putalpha(mask_pil)
    return cutout


def compose_thumbnail(thumbnail, model_photo, scale, x_pos, y_pos):
    if thumbnail is None:
        raise gr.Error("Please upload a thumbnail/background image.")
    if model_photo is None:
        raise gr.Error("Please upload a model photo.")

    try:
        thumbnail = thumbnail.convert("RGBA")
        cutout = cutout_model(model_photo)

        # Scale the cutout relative to the thumbnail's height
        thumb_w, thumb_h = thumbnail.size
        target_h = int(thumb_h * (scale / 100.0))
        aspect = cutout.width / cutout.height
        target_w = int(target_h * aspect)
        cutout_resized = cutout.resize((target_w, target_h), Image.LANCZOS)

        # Position: x_pos/y_pos are percentages of thumbnail width/height,
        # representing where the CENTER of the cutout should land.
        center_x = int(thumb_w * (x_pos / 100.0))
        center_y = int(thumb_h * (y_pos / 100.0))
        paste_x = center_x - target_w // 2
        paste_y = center_y - target_h // 2

        result = thumbnail.copy()
        result.paste(cutout_resized, (paste_x, paste_y), cutout_resized)
        return result.convert("RGB")

    except Exception as e:
        import traceback
        traceback.print_exc()
        raise gr.Error(f"Compositing failed: {e}")


css = """
#header {
    text-align: center;
    padding: 24px 0 8px;
}
#header h1 {
    font-size: 32px;
    font-weight: 700;
    background: linear-gradient(135deg, #6366f1, #06b6d4);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    margin-bottom: 4px;
}
#header p {
    color: #888;
    font-size: 14px;
}
#run-btn {
    background: linear-gradient(135deg, #6366f1, #06b6d4) !important;
    color: white !important;
    font-weight: 600 !important;
    border: none !important;
}
"""

with gr.Blocks(title="Peace Network Thumbnail Composer", css=css) as demo:
    gr.HTML(
        """
        <div id="header">
            <h1>🎬 Peace Network Thumbnail Composer</h1>
            <p>Thumbnail background + model photo daaliye — model automatically fit ho jayega.</p>
        </div>
        """
    )

    with gr.Row():
        thumb_input = gr.Image(type="pil", label="1. Thumbnail / background")
        model_input = gr.Image(type="pil", label="2. Model photo")

    with gr.Row():
        scale_slider = gr.Slider(
            minimum=10, maximum=100, value=60, step=1,
            label="Model size (% of thumbnail height)"
        )
        x_slider = gr.Slider(
            minimum=0, maximum=100, value=50, step=1,
            label="Horizontal position (%)"
        )
        y_slider = gr.Slider(
            minimum=0, maximum=100, value=70, step=1,
            label="Vertical position (%)"
        )

    btn = gr.Button("✨ Compose Thumbnail", variant="primary", elem_id="run-btn")
    output = gr.Image(type="pil", label="Result")

    btn.click(
        compose_thumbnail,
        inputs=[thumb_input, model_input, scale_slider, x_slider, y_slider],
        outputs=output,
    )

demo.queue().launch()