import spaces
import numpy as np
import gradio as gr
from PIL import Image, ImageEnhance, ImageFilter, ImageDraw, ImageFont


def get_font(size: int):
    candidates = [
        "DejaVuSans-Bold.ttf",
        "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",
        "Arial.ttf",
    ]
    for path in candidates:
        try:
            return ImageFont.truetype(path, size)
        except Exception:
            continue
    try:
        return ImageFont.load_default(size=size)
    except Exception:
        return ImageFont.load_default()


def crop_to_ratio(img: Image.Image, ratio_key: str, max_size: int = 1024) -> Image.Image:
    ratios = {
        "1:1":  (1, 1),
        "4:3":  (4, 3),
        "3:4":  (3, 4),
        "16:9": (16, 9),
        "9:16": (9, 16),
    }
    tw, th = ratios.get(ratio_key, (1, 1))
    target_ratio = tw / th

    img = img.convert("RGB")
    w, h = img.size
    current_ratio = w / h

    if current_ratio > target_ratio:
        new_w = int(h * target_ratio)
        left = (w - new_w) // 2
        img = img.crop((left, 0, left + new_w, h))
    else:
        new_h = int(w / target_ratio)
        top = (h - new_h) // 2
        img = img.crop((0, top, w, top + new_h))

    w, h = img.size
    scale = max_size / max(w, h)
    if scale < 1:
        img = img.resize((int(w * scale), int(h * scale)), Image.LANCZOS)
    return img


def add_grain(arr: np.ndarray, amount: float = 15) -> np.ndarray:
    noise = np.random.normal(0, amount, arr.shape)
    out = arr.astype(np.float32) + noise
    return np.clip(out, 0, 255).astype(np.uint8)


def color_grade(arr: np.ndarray,
                shadow_tint, highlight_tint,
                shadow_strength=0.3,
                highlight_strength=0.3) -> np.ndarray:
    arr = arr.astype(np.float32)
    luminance = arr.mean(axis=2, keepdims=True) / 255.0
    shadow_mask    = (1 - luminance) * shadow_strength
    highlight_mask = luminance * highlight_strength
    for i in range(3):
        arr[:, :, i] += shadow_tint[i]    * shadow_mask[:, :, 0]
        arr[:, :, i] += highlight_tint[i] * highlight_mask[:, :, 0]
    return np.clip(arr, 0, 255).astype(np.uint8)


def channel_shift(arr: np.ndarray, shift_x: int = 3) -> np.ndarray:
    r, g, b = arr[:, :, 0], arr[:, :, 1], arr[:, :, 2]
    return np.dstack([np.roll(r, shift_x, axis=1),
                      g,
                      np.roll(b, -shift_x, axis=1)])


def add_scanlines(arr: np.ndarray, gap: int = 4, darken: int = 25) -> np.ndarray:
    arr = arr.copy()
    arr[::gap, :, :] = np.clip(
        arr[::gap, :, :].astype(int) - darken, 0, 255
    ).astype(np.uint8)
    return arr


def apply_vignette(img: Image.Image, strength: float = 0.5) -> Image.Image:
    w, h = img.size
    y, x  = np.ogrid[:h, :w]
    cx, cy = w / 2, h / 2
    max_dist = np.sqrt(cx**2 + cy**2)
    dist = np.sqrt((x - cx)**2 + (y - cy)**2) / max_dist
    mask = np.clip(1 - strength * dist**2, 0, 1)[:, :, None]
    arr  = np.array(img).astype(np.float32) * mask
    return Image.fromarray(np.clip(arr, 0, 255).astype(np.uint8))


def add_bloom(img: Image.Image,
              threshold=200, blur_radius=10, intensity=0.5) -> Image.Image:
    arr       = np.array(img).astype(np.float32)
    luminance = arr.mean(axis=2)
    mask      = np.clip((luminance - threshold) / max(255 - threshold, 1), 0, 1)
    bright    = (arr * mask[:, :, None]).astype(np.uint8)
    bright_img = Image.fromarray(bright).filter(ImageFilter.GaussianBlur(blur_radius))
    bright_arr = np.array(bright_img).astype(np.float32)
    out = 255 - (255 - arr) * (255 - bright_arr * intensity) / 255
    return Image.fromarray(np.clip(out, 0, 255).astype(np.uint8))


def add_light_leak(img: Image.Image,
                   color=(255, 120, 60), strength=0.3) -> Image.Image:
    w, h = img.size
    xv, yv = np.meshgrid(np.linspace(0, 1, w), np.linspace(0, 1, h))
    gradient = np.clip(1 - np.abs((xv + yv) / 2 - 0.75) * 3, 0, 1)
    arr = np.array(img).astype(np.float32)
    for i in range(3):
        arr[:, :, i] += gradient * color[i] * strength
    return Image.fromarray(np.clip(arr, 0, 255).astype(np.uint8))


def draw_caption(img: Image.Image, text: str,
                 prefix: str = "",
                 text_color=(255, 255, 255, 255),
                 bg_color=(0, 0, 0, 130),
                 position: str = "bottom") -> Image.Image:
    img     = img.convert("RGBA")
    overlay = Image.new("RGBA", img.size, (0, 0, 0, 0))
    draw    = ImageDraw.Draw(overlay)
    w, h    = img.size

    full_text = f"{prefix} {text}".strip() if prefix else text
    font_size = max(16, w // 30)
    font      = get_font(font_size)

    bbox   = draw.textbbox((0, 0), full_text, font=font)
    text_w = bbox[2] - bbox[0]
    text_h = bbox[3] - bbox[1]
    pad    = 12

    if position == "bottom":
        box = (0, h - text_h - pad * 2, w, h)
    else:
        box = (0, 0, w, text_h + pad * 2)

    draw.rectangle(box, fill=bg_color)
    draw.text(((w - text_w) // 2, box[1] + pad),
              full_text, font=font, fill=text_color)

    return Image.alpha_composite(img, overlay).convert("RGB")


def add_polaroid_border(img: Image.Image,
                        border=20, bottom_extra=70) -> Image.Image:
    w, h = img.size
    canvas = Image.new("RGB",
                       (w + border * 2, h + border * 2 + bottom_extra),
                       (250, 250, 245))
    canvas.paste(img, (border, border))
    return canvas


def apply_1980s(img: Image.Image, text: str) -> Image.Image:
    out = ImageEnhance.Color(img).enhance(1.4)
    out = ImageEnhance.Contrast(out).enhance(1.15)
    out = Image.fromarray(color_grade(np.array(out),
          shadow_tint=(20, -10, 60), highlight_tint=(255, 60, 180),
          shadow_strength=0.5, highlight_strength=0.35))
    out = add_bloom(out, threshold=180, blur_radius=12, intensity=0.6)
    out = add_light_leak(out, color=(255, 80, 180), strength=0.3)
    out = Image.fromarray(add_grain(np.array(out), amount=10))
    out = draw_caption(out, text, prefix="◆ 1980s ◆",
                       text_color=(255, 0, 200, 255))
    return out


def apply_1990s(img: Image.Image, text: str) -> Image.Image:
    out = ImageEnhance.Color(img).enhance(0.85)
    out = ImageEnhance.Contrast(out).enhance(0.9)
    out = ImageEnhance.Brightness(out).enhance(1.05)
    arr = color_grade(np.array(out),
          shadow_tint=(10, 5, -10), highlight_tint=(255, 220, 150),
          shadow_strength=0.3, highlight_strength=0.4)
    arr = channel_shift(arr, shift_x=3)
    arr = add_scanlines(arr, gap=4, darken=25)
    arr = add_grain(arr, amount=18)
    out = Image.fromarray(arr).filter(ImageFilter.GaussianBlur(0.6))
    out = draw_caption(out, text, prefix="● REC",
                       text_color=(255, 60, 60, 255), position="top")
    return out


def apply_2000s(img: Image.Image, text: str) -> Image.Image:
    out = ImageEnhance.Brightness(img).enhance(1.25)
    out = ImageEnhance.Contrast(out).enhance(1.1)
    out = ImageEnhance.Sharpness(out).enhance(1.5)
    out = Image.fromarray(color_grade(np.array(out),
          shadow_tint=(-10, 0, 40), highlight_tint=(30, 30, 60),
          shadow_strength=0.3, highlight_strength=0.2))
    out = apply_vignette(out, strength=0.55)
    out = Image.fromarray(add_grain(np.array(out), amount=14))
    out = draw_caption(out, text, prefix="📷",
                       text_color=(255, 140, 0, 255))
    return out


def apply_2010s(img: Image.Image, text: str) -> Image.Image:
    out = ImageEnhance.Color(img).enhance(1.1)
    out = ImageEnhance.Contrast(out).enhance(0.92)
    arr = np.clip(np.array(out).astype(np.float32) * 0.9 + 25, 0, 255).astype(np.uint8)
    arr = color_grade(arr,
          shadow_tint=(30, 10, -10), highlight_tint=(255, 200, 120),
          shadow_strength=0.35, highlight_strength=0.3)
    out = apply_vignette(Image.fromarray(arr), strength=0.3)
    out = Image.fromarray(add_grain(np.array(out), amount=6))
    out = add_polaroid_border(out)
    out = draw_caption(out, text, prefix="📸",
                       text_color=(60, 60, 60, 255),
                       bg_color=(0, 0, 0, 0), position="bottom")
    return out


ERA_FUNCS = {
    "1980-е 🕺": apply_1980s,
    "1990-е 📼": apply_1990s,
    "2000-е 💿": apply_2000s,
    "2010-е 📱": apply_2010s,
}

@spaces.GPU(duration=10)
def generate_all_eras(selfie, scene_text, aspect_ratio):
    if selfie is None:
        raise gr.Error("Пожалуйста, загрузите ваше селфи! 🤳")

    scene_text = (scene_text or "").strip() or "путешествие во времени"
    base = crop_to_ratio(selfie, aspect_ratio)

    results = []
    for era_name, era_func in ERA_FUNCS.items():
        results.append((era_func(base.copy(), scene_text), era_name))
    return results

CUSTOM_CSS = """
#title { text-align: center; font-family: 'Georgia', serif; }
.gradio-container {
    background: linear-gradient(135deg, #1a1a2e 0%, #16213e 50%, #0f3460 100%);
}
#generate-btn {
    background: linear-gradient(90deg, #f7971e, #ffd200) !important;
    font-weight: bold;
    font-size: 18px !important;
}
"""

with gr.Blocks(css=CUSTOM_CSS, theme=gr.themes.Soft(primary_hue="orange")) as demo:
    gr.Markdown(
        "<h1 id='title'>⏳ Eras Time Machine AI</h1>"
        "<p style='text-align:center;'>"
        "Загрузи селфи — получи себя в стиле 1980-х, 1990-х, 2000-х и 2010-х! ✨"
        "</p>"
    )

    with gr.Row():
        with gr.Column(scale=1):
            selfie_input = gr.Image(
                type="pil",
                label="📸 Загрузите ваше селфи",
            )
            scene_input = gr.Textbox(
                label="✍️ Опишите сцену / момент",
                placeholder="Например: сижу на капоте машины на закате",
                lines=2,
            )
            aspect_input = gr.Radio(
                choices=["1:1", "4:3", "3:4", "16:9", "9:16"],
                value="1:1",
                label="📐 Формат изображения",
            )
            generate_btn = gr.Button(
                "Отправиться сквозь время 🚀",
                elem_id="generate-btn",
                variant="primary",
            )

        with gr.Column(scale=1):
            gallery_output = gr.Gallery(
                label="🖼️ Результаты путешествия во времени",
                columns=2,
                rows=2,
                height="auto",
                object_fit="contain",
                preview=True,
            )

    generate_btn.click(
        fn=generate_all_eras,
        inputs=[selfie_input, scene_input, aspect_input],
        outputs=gallery_output,
    )

    gr.Markdown(
        "<p style='text-align:center; opacity:0.6; font-size:12px;'>"
        "CPU-фильтры · ZeroGPU · Hugging Face Spaces · Бесплатно для всех"
        "</p>"
    )

if __name__ == "__main__":
    demo.queue().launch()