import os, io, math, time, base64
from typing import Tuple
import gradio as gr
from PIL import Image, ImageFilter, ImageOps
import numpy as np
import torch

# ---- Модель по умолчанию (можно переопределить в Settings → Variables) ----
MODEL_ID = os.getenv("MODEL_ID", "Qwen/Qwen-Image-Edit")
HF_TOKEN = os.getenv("HF_TOKEN", "").strip() or None  # если модель gated

# ---- Подгружаем подходящий пайплайн inpainting ----
# Многие редактирующие модели следуют API diffusers InpaintPipeline
# Попробуем сначала специализированный пайплайн, затем универсальный.
from diffusers import AutoPipelineForInpainting, StableDiffusionInpaintPipeline

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
dtype = torch.float16 if device.type == "cuda" else torch.float32

def _load_pipeline():
    auth = {"token": HF_TOKEN} if HF_TOKEN else {}
    try:
        pipe = AutoPipelineForInpainting.from_pretrained(
            MODEL_ID, torch_dtype=dtype, **auth
        )
    except Exception:
        # fallback на классический SD inpaint, если у модели нет auto-конфига
        pipe = StableDiffusionInpaintPipeline.from_pretrained(
            MODEL_ID, torch_dtype=dtype, **auth
        )
    pipe = pipe.to(device)
    if hasattr(pipe, "enable_attention_slicing"):
        pipe.enable_attention_slicing()
    return pipe

PIPE = None
LOAD_ERR = None
try:
    PIPE = _load_pipeline()
except Exception as e:
    LOAD_ERR = str(e)

# ----------------- Утилиты -----------------
def _pad_canvas(img: Image.Image, left: int, right: int, top: int, bottom: int,
                fill=(127,127,127)) -> Tuple[Image.Image, Tuple[int,int,int,int]]:
    """Создаёт расширенный холст и возвращает (canvas, bbox вставки исходника)."""
    w, h = img.size
    canvas = Image.new("RGB", (w + left + right, h + top + bottom), fill)
    canvas.paste(img, (left, top))
    return canvas, (left, top, left + w, top + h)

def _feather_mask(size: Tuple[int,int], bbox: Tuple[int,int,int,int], feather_px: int = 32) -> Image.Image:
    """Маска: 255 = дорисовать, 0 = оставить исходник. По краю делаем плавный градиент."""
    W, H = size
    L, T, R, B = bbox
    mask = Image.new("L", (W, H), 255)
    base = Image.new("L", (W, H), 0)
    base.paste(255, (L, T, R, B))
    # invert: центральная область = 0 (не трогать), снаружи = 255
    inv = ImageOps.invert(base)
    if feather_px > 0:
        inv = inv.filter(ImageFilter.GaussianBlur(radius=feather_px))
    return inv

def expand_with_model(image: Image.Image, left: int, right: int, top: int, bottom: int,
                      prompt: str, neg: str, steps: int, guidance: float, seed: int | None):
    if LOAD_ERR:
        raise gr.Error(f"Не удалось загрузить модель '{MODEL_ID}'.\n{LOAD_ERR}")

    for name, val in [("left",left),("right",right),("top",top),("bottom",bottom)]:
        if val is None or val < 0 or val > 2048:
            raise gr.Error(f"{name} должен быть в диапазоне 0..2048")

    # 1) строим холст и маску
    canvas, bbox = _pad_canvas(image.convert("RGB"), left, right, top, bottom)
    mask = _feather_mask(canvas.size, bbox, feather_px=48)

    # 2) аккуратный промпт для бесшовного расширения
    clean_prompt = (prompt or "").strip()
    if not clean_prompt:
        clean_prompt = (
            "Seamlessly continue the existing background so it looks like a natural, "
            "wider version of the same image. Match colors, textures, lighting and perspective. "
            "No frames, no collage, no new subjects."
        )
    negative_prompt = (neg or "frames, borders, phone mockup, collage tiles, UI elements, captions, watermark").strip()

    generator = None
    if isinstance(seed, int) and seed >= 0:
        generator = torch.Generator(device=device).manual_seed(seed)

    # 3) инференс
    out = PIPE(
        prompt=clean_prompt,
        negative_prompt=negative_prompt,
        image=canvas,
        mask_image=mask,
        num_inference_steps=int(steps),
        guidance_scale=float(guidance),
        generator=generator
    ).images[0]

    return out, clean_prompt, negative_prompt

# ----------------- Gradio UI -----------------
with gr.Blocks(title="Qwen Image — Seamless Expand") as demo:
    gr.Markdown("## Qwen Image — Seamless Background Expand\n"
                "Загрузи картинку, задай сколько пикселей дорисовать, при желании уточни промпт.")

    if LOAD_ERR:
        gr.Markdown(
            f"**⚠️ Модель не загрузилась.** Проверь `MODEL_ID` в Settings → Variables. Текущая: `{MODEL_ID}`  "
            f"\nСообщение: `{LOAD_ERR}`"
        )

    with gr.Row():
        with gr.Column(scale=1):
            img_in = gr.Image(type="pil", label="Картинка")
            prompt = gr.Textbox(label="Prompt (EN)", placeholder="Describe seamless continuation…")
            neg = gr.Textbox(label="Negative prompt", value="frames, borders, phone mockup, collage tiles, UI elements, captions, watermark")
            with gr.Row():
                left = gr.Number(value=256, label="Left (px)", precision=0)
                right = gr.Number(value=256, label="Right (px)", precision=0)
            with gr.Row():
                top = gr.Number(value=256, label="Top (px)", precision=0)
                bottom = gr.Number(value=256, label="Bottom (px)", precision=0)
            with gr.Row():
                steps = gr.Slider(10, 60, value=30, step=1, label="Steps")
                guidance = gr.Slider(0.5, 12.0, value=5.5, step=0.1, label="Guidance scale")
                seed = gr.Number(value=-1, label="Seed (-1 = random)", precision=0)
            btn = gr.Button("Expand", variant="primary")
        with gr.Column(scale=1):
            img_out = gr.Image(label="Результат")
            used_prompt = gr.Textbox(label="Использованный prompt", interactive=False)
            used_neg = gr.Textbox(label="Использованный negative", interactive=False)

    btn.click(
        expand_with_model,
        inputs=[img_in, left, right, top, bottom, prompt, neg, steps, guidance, seed],
        outputs=[img_out, used_prompt, used_neg]
    )

demo.launch()