import gc
import os
import random

import gradio as gr
import numpy as np
import torch
from PIL import Image, ImageDraw, ImageFilter
from diffusers import AutoPipelineForImage2Image, LCMScheduler


DEVICE = "cpu"
SIZE = 512

BASE_MODEL = os.environ.get("BASE_MODEL", "Lykon/dreamshaper-7")
LCM_LORA = os.environ.get("LCM_LORA", "latent-consistency/lcm-lora-sdv1-5")
IP_ADAPTER_REPO = os.environ.get("IP_ADAPTER_REPO", "h94/IP-Adapter")
IP_ADAPTER_WEIGHT = os.environ.get("IP_ADAPTER_WEIGHT", "ip-adapter_sd15.bin")


print("Loading base model...")
pipe = AutoPipelineForImage2Image.from_pretrained(
    BASE_MODEL,
    torch_dtype=torch.float32,
)

print("Loading LCM scheduler...")
pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)

print("Loading LCM-LoRA...")
pipe.load_lora_weights(LCM_LORA)

try:
    pipe.fuse_lora()
except Exception as exc:
    print(f"Could not fuse LoRA, continuing without fuse: {exc}")

print("Loading IP-Adapter...")
pipe.load_ip_adapter(
    IP_ADAPTER_REPO,
    subfolder="models",
    weight_name=IP_ADAPTER_WEIGHT,
    low_cpu_mem_usage=True,
)

pipe = pipe.to(DEVICE)

# IMPORTANT:
# Do not enable attention slicing with IP-Adapter.
# It can cause:
# AttributeError: 'tuple' object has no attribute 'shape'

try:
    pipe.vae.enable_slicing()
except Exception as exc:
    print(f"Could not enable VAE slicing: {exc}")


SIDE_PROMPTS = {
    "left": (
        "orthographic left side elevation of the same storefront building, "
        "flat lighting, square texture tile, plausible side wall details, "
        "brick or concrete, side windows, grime, utility pipes, urban wear, "
        "no people, no cars, no perspective distortion"
    ),
    "right": (
        "orthographic right side elevation of the same storefront building, "
        "flat lighting, square texture tile, plausible side wall details, "
        "brick or concrete, side windows, grime, utility pipes, urban wear, "
        "no people, no cars, no perspective distortion"
    ),
    "back": (
        "orthographic rear elevation of the same storefront building, "
        "flat lighting, square texture tile, plausible rear wall details, "
        "service door, vents, pipes, brick or concrete, grime, urban wear, "
        "no people, no cars, no perspective distortion"
    ),
    "top": (
        "orthographic roof view of the same storefront building, "
        "flat top-down square texture tile, tar roof or concrete roof, "
        "vents, HVAC units, roof seams, stains, grime, urban wear, "
        "no perspective distortion"
    ),
}


def center_crop_resize(img: Image.Image, size: int = SIZE) -> Image.Image:
    img = img.convert("RGB")
    width, height = img.size
    side = min(width, height)

    left = (width - side) // 2
    top = (height - side) // 2

    img = img.crop((left, top, left + side, top + side))
    img = img.resize((size, size), Image.LANCZOS)

    return img


def average_color(img: Image.Image) -> tuple:
    arr = np.array(img.convert("RGB"))
    avg = np.mean(arr.reshape(-1, 3), axis=0).astype(np.uint8)
    return tuple(avg.tolist())


def sample_region(front_img: Image.Image, side: str, size: int = SIZE) -> Image.Image:
    img = center_crop_resize(front_img, size)
    arr = np.array(img)
    strip = max(32, size // 7)

    if side == "left":
        patch = arr[:, :strip, :]
    elif side == "right":
        patch = arr[:, size - strip:, :]
    elif side == "top":
        patch = arr[:strip, :, :]
    else:
        # For back view, sample the center third to avoid edge artifacts.
        x1 = size // 3
        x2 = size - size // 3
        patch = arr[:, x1:x2, :]

    return Image.fromarray(patch).convert("RGB")


def add_noise_texture(img: Image.Image, amount: int = 16, seed: int = 0) -> Image.Image:
    rng = np.random.default_rng(seed)
    arr = np.array(img).astype(np.int16)
    noise = rng.normal(0, amount, arr.shape).astype(np.int16)
    arr = np.clip(arr + noise, 0, 255).astype(np.uint8)
    return Image.fromarray(arr, "RGB")


def blend_texture(
    base: Image.Image,
    texture_source: Image.Image,
    opacity: float = 0.35,
) -> Image.Image:
    texture = texture_source.resize(base.size, Image.BICUBIC)
    texture = texture.filter(ImageFilter.GaussianBlur(radius=7))
    return Image.blend(base.convert("RGB"), texture.convert("RGB"), opacity)


def draw_rect(draw: ImageDraw.ImageDraw, xy, outline, width=3, fill=None):
    x1, y1, x2, y2 = xy
    draw.rectangle(
        (int(x1), int(y1), int(x2), int(y2)),
        outline=outline,
        width=width,
        fill=fill,
    )


def draw_line(draw: ImageDraw.ImageDraw, xy, fill, width=2):
    draw.line(tuple(int(v) for v in xy), fill=fill, width=width)


def darken(color: tuple, factor: float = 0.65) -> tuple:
    return tuple(max(0, min(255, int(c * factor))) for c in color)


def lighten(color: tuple, factor: float = 1.25) -> tuple:
    return tuple(max(0, min(255, int(c * factor))) for c in color)


def hybrid_side_guide(
    front_img: Image.Image,
    side: str,
    seed: int,
    size: int = SIZE,
) -> Image.Image:
    """
    Hybrid guide image.

    Purpose:
    - The front image is used by IP-Adapter as visual/style reference.
    - This guide image gives img2img a rough structural canvas.
    - It avoids directly warping the front facade into a side view.
    """

    rng = random.Random(seed)

    patch = sample_region(front_img, side, size)
    base_color = average_color(patch)

    if side == "top":
        base_color = darken(base_color, 0.55)

    guide = Image.new("RGB", (size, size), base_color)
    guide = blend_texture(guide, patch, opacity=0.28)
    guide = add_noise_texture(guide, amount=10, seed=seed)

    draw = ImageDraw.Draw(guide)
    line_color = darken(base_color, 0.55)
    light_line = lighten(base_color, 1.35)

    if side in ["left", "right"]:
        # Wall divisions.
        for x in [size * 0.22, size * 0.48, size * 0.74]:
            draw_line(draw, (x, 0, x, size), fill=line_color, width=2)

        for y in [size * 0.18, size * 0.42, size * 0.68, size * 0.86]:
            draw_line(draw, (0, y, size, y), fill=darken(base_color, 0.75), width=1)

        # Generic side windows.
        window_fill = darken(base_color, 0.45)

        for i in range(2):
            x = size * (0.18 + i * 0.38) + rng.randint(-12, 12)
            y = size * 0.22 + rng.randint(-20, 20)

            draw_rect(
                draw,
                (x, y, x + size * 0.18, y + size * 0.18),
                outline=line_color,
                fill=window_fill,
                width=3,
            )

        # Pipe / utility line.
        pipe_x = size * (0.82 if side == "left" else 0.14)

        draw_line(
            draw,
            (pipe_x, size * 0.08, pipe_x, size * 0.92),
            fill=light_line,
            width=4,
        )

        draw_line(
            draw,
            (pipe_x + 8, size * 0.08, pipe_x + 8, size * 0.92),
            fill=line_color,
            width=2,
        )

        # Bottom grime band.
        grime = Image.new("RGBA", (size, size), (35, 30, 25, 0))
        grime_draw = ImageDraw.Draw(grime)
        grime_draw.rectangle(
            (0, int(size * 0.78), size, size),
            fill=(35, 30, 25, 55),
        )

        guide = Image.alpha_composite(guide.convert("RGBA"), grime).convert("RGB")

    elif side == "back":
        # Rear wall frame.
        draw_rect(draw, (0, 0, size - 1, size - 1), outline=line_color, width=4)

        # Service door.
        door_fill = darken(base_color, 0.5)

        draw_rect(
            draw,
            (size * 0.38, size * 0.45, size * 0.62, size * 0.95),
            outline=line_color,
            fill=door_fill,
            width=4,
        )

        # Vents / utility panels.
        vent_fill = darken(base_color, 0.62)

        draw_rect(
            draw,
            (size * 0.12, size * 0.18, size * 0.26, size * 0.34),
            outline=line_color,
            fill=vent_fill,
            width=3,
        )

        draw_rect(
            draw,
            (size * 0.74, size * 0.18, size * 0.88, size * 0.34),
            outline=line_color,
            fill=vent_fill,
            width=3,
        )

        # Rear wall horizontal material / grime lines.
        for y in [size * 0.22, size * 0.44, size * 0.66, size * 0.84]:
            draw_line(
                draw,
                (0, y, size, y),
                fill=darken(base_color, 0.75),
                width=2,
            )

        # Pipes.
        draw_line(
            draw,
            (size * 0.08, size * 0.15, size * 0.08, size * 0.95),
            fill=light_line,
            width=4,
        )

        draw_line(
            draw,
            (size * 0.92, size * 0.10, size * 0.92, size * 0.72),
            fill=light_line,
            width=3,
        )

    elif side == "top":
        roof_overlay = Image.new("RGBA", (size, size), (35, 35, 35, 70))
        guide = Image.alpha_composite(guide.convert("RGBA"), roof_overlay).convert("RGB")
        draw = ImageDraw.Draw(guide)

        roof_line = (130, 130, 130)

        # Roof seams.
        for y in [size * 0.18, size * 0.34, size * 0.52, size * 0.70, size * 0.86]:
            draw_line(draw, (0, y, size, y), fill=roof_line, width=2)

        # HVAC / roof units.
        draw_rect(
            draw,
            (size * 0.16, size * 0.20, size * 0.36, size * 0.40),
            outline=(210, 210, 210),
            fill=(80, 80, 80),
            width=3,
        )

        draw_rect(
            draw,
            (size * 0.58, size * 0.30, size * 0.80, size * 0.52),
            outline=(210, 210, 210),
            fill=(75, 75, 75),
            width=3,
        )

        # Vents / drains.
        for cx, cy in [(0.25, 0.68), (0.70, 0.74), (0.48, 0.18)]:
            r = size * 0.035

            draw.ellipse(
                (
                    int(size * cx - r),
                    int(size * cy - r),
                    int(size * cx + r),
                    int(size * cy + r),
                ),
                outline=(210, 210, 210),
                width=3,
            )

    return guide.filter(ImageFilter.GaussianBlur(radius=0.35))


@torch.inference_mode()
def generate_side(front_image, side, prompt_suffix, strength, steps, seed, ip_scale):
    if front_image is None:
        raise gr.Error("Upload a front image first.")

    seed = int(seed)

    ref_image = center_crop_resize(front_image)
    guide_image = hybrid_side_guide(front_image, side, seed=seed)

    base_prompt = SIDE_PROMPTS[side]

    if prompt_suffix and prompt_suffix.strip():
        prompt = f"{base_prompt}, {prompt_suffix.strip()}"
    else:
        prompt = base_prompt

    pipe.set_ip_adapter_scale(float(ip_scale))

    generator = torch.Generator(device=DEVICE).manual_seed(seed)

    result = pipe(
        prompt=prompt,
        image=guide_image,
        ip_adapter_image=ref_image,
        strength=float(strength),
        num_inference_steps=int(steps),
        guidance_scale=1.5,
        generator=generator,
    ).images[0]

    gc.collect()

    status = (
        "Done. Used DreamShaper-7 + LCM-LoRA + IP-Adapter on CPU. "
        "The guide image is hybrid: sampled material color/texture plus procedural side/back/roof layout. "
        "The uploaded front image is used as the IP-Adapter reference."
    )

    return result, guide_image, status


with gr.Blocks(title="Building Side View Generator - IPAdapter CPU") as demo:
    gr.Markdown(
        """
# Building Side View Generator - IPAdapter CPU

Upload a front-facing cropped building image and generate **one plausible side view**.

This CPU prototype uses:

- `Lykon/dreamshaper-7`
- `latent-consistency/lcm-lora-sdv1-5`
- `h94/IP-Adapter`

The uploaded front image is used as the **IP-Adapter reference image**.  
The guide image uses a **hybrid approach**: sampled material color/texture from the front image plus procedural side/back/roof layout.
"""
    )

    with gr.Row():
        front = gr.Image(label="Front Image", type="pil")
        output = gr.Image(label="Generated Side View", type="pil")

    with gr.Row():
        guide_preview = gr.Image(
            label="Hybrid Guide Image Used For Img2Img",
            type="pil",
        )
        status = gr.Textbox(label="Status")

    with gr.Row():
        side = gr.Dropdown(
            choices=["left", "right", "back", "top"],
            value="left",
            label="View to Generate",
        )

        prompt_suffix = gr.Textbox(
            label="Optional Prompt Add-on",
            placeholder="e.g. weathered brick, old windows, grime, utility pipes",
        )

    with gr.Row():
        strength = gr.Slider(
            minimum=0.35,
            maximum=0.85,
            value=0.62,
            step=0.05,
            label="Img2Img Strength",
        )

        steps = gr.Slider(
            minimum=2,
            maximum=6,
            value=4,
            step=1,
            label="Inference Steps",
        )

    with gr.Row():
        ip_scale = gr.Slider(
            minimum=0.2,
            maximum=1.0,
            value=0.60,
            step=0.05,
            label="IPAdapter Scale",
        )

        seed = gr.Number(
            value=1234,
            precision=0,
            label="Seed",
        )

    run = gr.Button("Generate Side View")

    run.click(
        fn=generate_side,
        inputs=[front, side, prompt_suffix, strength, steps, seed, ip_scale],
        outputs=[output, guide_preview, status],
        api_name="generate_side",
    )


if __name__ == "__main__":
    demo.launch(show_error=True)