import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces
import random
from math import sqrt
from pathlib import Path

import gradio as gr
import torch
from PIL import Image, ImageOps

from larpscaler import LarpScaler


MODEL_ID = "VladimirM388/larpscaler-v2-bf16"
MAX_OUTPUT_PIXELS = 4_194_304
QUALITY_MODE = "Quality"
FAST_MODE = "Fast"
ASSETS = Path(__file__).parent / "assets" / "examples"

PRESETS = {
    QUALITY_MODE: {"steps": 1, "noise_level": 1.0, "guidance_scale": 4.5},
    FAST_MODE: {"steps": 1, "noise_level": 0.35, "guidance_scale": 1.0},
}

EXAMPLES = [
    [str(ASSETS / "mountain.jpg")],
    [str(ASSETS / "architecture.jpg")],
    [str(ASSETS / "portrait.jpg")],
]

CSS = """
:root {
  --page: #0c0e13;
  --panel: #151821;
  --surface: #10131a;
  --text: #f1f4fa;
  --muted: #8d95a6;
  --border: #292e3a;
  --accent: #a89bff;
}

body,
.gradio-container {
  background: radial-gradient(900px 480px at 50% -180px, rgba(120, 105, 234, 0.16), transparent 72%), var(--page) !important;
  color: var(--text) !important;
  font-family: Inter, ui-sans-serif, system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif !important;
}

main,
.gradio-container,
.contain { max-width: none !important; }

#larp-shell {
  max-width: 1260px;
  margin: 0 auto;
  padding: 28px 20px 36px;
}

#workspace {
  align-items: stretch !important;
  gap: 18px !important;
}

#source-panel,
#result-panel {
  min-height: 0;
  padding: 16px !important;
  border: 1px solid var(--border) !important;
  border-radius: 16px !important;
  background: var(--panel) !important;
  box-shadow: 0 18px 44px rgba(0, 0, 0, 0.22);
}

.panel-title {
  margin: 1px 2px 14px;
  color: #aeb6c7;
  font-size: 0.68rem;
  font-weight: 700;
  letter-spacing: 0.12em;
  text-transform: uppercase;
}

#source-image,
#result-image {
  overflow: hidden;
  border: 1px solid var(--border);
  border-radius: 12px !important;
  background: var(--surface);
}

#source-image button,
#result-image button { border-radius: 8px !important; }

.gradio-container [data-testid="block-info"] {
  color: var(--muted) !important;
  font-size: 0.69rem !important;
  font-weight: 700 !important;
  letter-spacing: 0.07em;
  text-transform: uppercase;
}

#control-dock {
  align-items: end !important;
  gap: 14px !important;
  margin-top: 18px !important;
  padding: 14px 16px 16px !important;
  border: 1px solid var(--border) !important;
  border-radius: 16px !important;
  background: rgba(21, 24, 33, 0.96) !important;
}

#control-dock .form {
  background: transparent !important;
}

#scale-control,
#mode-control,
#sample-strip { min-height: 52px; }

#scale-control,
#mode-control {
  padding: 0 !important;
  border: 0 !important;
  background: transparent !important;
}

#scale-control .wrap,
#mode-control .wrap {
  display: flex !important;
  gap: 6px !important;
  padding: 0 !important;
  background: transparent !important;
}

#scale-control label,
#mode-control label {
  display: flex !important;
  flex: 1 1 0;
  min-width: 0;
  min-height: 40px;
  align-items: center;
  justify-content: center;
  margin: 0 !important;
  padding: 0 8px !important;
  border: 1px solid var(--border) !important;
  border-radius: 9px !important;
  background: var(--surface) !important;
  color: #aeb6c7 !important;
  font-size: 0.78rem !important;
  font-weight: 650 !important;
}

#scale-control label span,
#mode-control label span {
  overflow: visible !important;
  white-space: nowrap !important;
}

#scale-control input,
#mode-control input {
  flex: 0 0 auto;
  margin-right: 6px !important;
}

#scale-control label.selected,
#mode-control label.selected {
  border-color: var(--accent) !important;
  background: rgba(168, 155, 255, 0.15) !important;
  color: #f1efff !important;
}

#scale-control input,
#mode-control input { accent-color: var(--accent) !important; }

#sample-strip .examples { margin: 0 !important; }
#sample-strip .gallery { gap: 6px !important; }
#sample-strip .thumbnail-item {
  overflow: hidden;
  border: 1px solid var(--border) !important;
  border-radius: 8px !important;
  background: var(--surface) !important;
}

.control-label {
  height: 17px;
  margin: 0 0 5px 1px;
  color: var(--muted);
  font-size: 0.69rem;
  font-weight: 700;
  letter-spacing: 0.07em;
  text-transform: uppercase;
}

button#upscale-button,
button#clear-button {
  min-height: 40px;
  border-radius: 10px !important;
  font-size: 0.84rem !important;
  font-weight: 750 !important;
}

button#upscale-button {
  border: 0 !important;
  background: var(--accent) !important;
  box-shadow: 0 8px 22px rgba(168, 155, 255, 0.2);
  color: #12131a !important;
}

button#clear-button {
  border: 1px solid var(--border) !important;
  background: var(--surface) !important;
  color: #c9cfdb !important;
}

button#upscale-button:hover { background: #beb5ff !important; }
button#clear-button:hover { border-color: #474e60 !important; background: #191d27 !important; }

#status-card {
  min-height: 18px;
  margin: 10px 2px 0;
  color: var(--muted);
  text-align: center;
}

#status-card p { margin: 0; font-size: 0.76rem; }
footer[aria-label="Gradio footer navigation"] { display: none !important; }

@media (max-width: 760px) {
  #larp-shell { padding: 14px 12px 24px; }
  #workspace, #control-dock { gap: 12px !important; }
}
"""


def _open_image(path: str | None) -> Image.Image:
    if not path:
        raise gr.Error("Upload an image before starting the upscale.")
    with Image.open(path) as loaded:
        return ImageOps.exif_transpose(loaded).convert("RGB")


def _fit_input_to_output_budget(image: Image.Image, scale: int) -> tuple[Image.Image, bool]:
    output_pixels = image.width * image.height * scale * scale
    if output_pixels <= MAX_OUTPUT_PIXELS:
        return image, False
    ratio = sqrt(MAX_OUTPUT_PIXELS / output_pixels)
    width = max(1, int(image.width * ratio))
    height = max(1, int(image.height * ratio))
    return image.resize((width, height), Image.Resampling.LANCZOS), True


def _run_upscale(
    image_path: str | None,
    scale: str | int,
    mode: str,
    seed: float,
    randomize_seed: bool,
    adapter_path: str | None,
    use_image_adapter: bool,
) -> tuple[Image.Image, int, str]:
    image = _open_image(image_path)
    image, _ = _fit_input_to_output_budget(image, int(scale))
    adapter_image = _open_image(adapter_path) if adapter_path and use_image_adapter else None
    used_seed = random.randint(0, 2**31 - 1) if randomize_seed else int(seed)
    preset = PRESETS[mode]

    result = UPSCALER.upscale(
        image,
        scale=int(scale),
        steps=preset["steps"],
        noise_level=preset["noise_level"],
        guidance_scale=preset["guidance_scale"],
        seed=used_seed,
        adapter_image=adapter_image,
        use_image_adapter=bool(use_image_adapter),
        tile_size=1024,
        tile_overlap=256,
        tile_batch_size="auto",
    )
    status = f"{image.width} × {image.height}  →  {result.width} × {result.height}"
    return result, used_seed, status


UPSCALER = LarpScaler.from_pretrained(
    MODEL_ID,
    device="cuda",
    dtype=torch.bfloat16,
)


@spaces.GPU(duration=60)
def upscale_image(
    image_path: str | None,
    scale: str,
    mode: str,
    seed: float,
    randomize_seed: bool,
    adapter_path: str | None,
    use_image_adapter: bool,
) -> tuple[Image.Image, int, str]:
    """Upscale an image with LARP-Scaler on ZeroGPU.

    Args:
        image_path: Uploaded source image.
        scale: Enlargement factor, one of 2, 4, or 8.
        mode: Fast or quality inference preset.
        seed: Seed used for deterministic refinement.
        randomize_seed: Choose a new seed for this request.
        adapter_path: Optional reference image for image guidance.
        use_image_adapter: Whether to enable LARP-Scaler's guidance adapter.
    """
    return _run_upscale(
        image_path,
        scale,
        mode,
        seed,
        randomize_seed,
        adapter_path,
        use_image_adapter,
    )


def reset_editor() -> tuple[None, str, str, float, bool, None, bool, None, str]:
    return (
        None,
        "4",
        QUALITY_MODE,
        1234,
        False,
        None,
        True,
        None,
        "",
    )


def clear_result() -> tuple[None, str]:
    """Clear the generated image and its status."""
    return None, ""


with gr.Blocks(title="LARP-Scaler") as demo:
    with gr.Column(elem_id="larp-shell"):
        with gr.Row(equal_height=True, elem_id="workspace"):
            with gr.Column(scale=1, min_width=360, elem_id="source-panel"):
                gr.HTML('<div class="panel-title">Source</div>')
                source_image = gr.Image(
                    show_label=False,
                    type="filepath",
                    sources=["upload", "clipboard"],
                    image_mode="RGB",
                    height=500,
                    elem_id="source-image",
                )

            with gr.Column(scale=1, min_width=360, elem_id="result-panel"):
                gr.HTML('<div class="panel-title">Result</div>')
                result_image = gr.Image(
                    show_label=False,
                    type="pil",
                    format="png",
                    height=500,
                    buttons=["download", "fullscreen"],
                    elem_id="result-image",
                    placeholder="",
                )

        with gr.Row(equal_height=True, elem_id="control-dock"):
            with gr.Column(scale=2, min_width=210):
                scale = gr.Radio(
                    choices=[("2×", "2"), ("4×", "4"), ("8×", "8")],
                    value="4",
                    label="Scale",
                    elem_id="scale-control",
                )
            with gr.Column(scale=2, min_width=210):
                mode = gr.Radio(
                    choices=[QUALITY_MODE, FAST_MODE],
                    value=QUALITY_MODE,
                    label="Mode",
                    elem_id="mode-control",
                )
            with gr.Column(scale=2, min_width=210, elem_id="sample-strip"):
                gr.Examples(
                    examples=EXAMPLES,
                    inputs=[source_image],
                    label="Samples",
                    cache_examples=False,
                )
            with gr.Column(scale=2, min_width=210):
                gr.HTML('<div class="control-label">Action</div>')
                with gr.Row():
                    run_button = gr.Button(
                        "Upscale",
                        variant="primary",
                        size="md",
                        elem_id="upscale-button",
                        scale=3,
                    )
                    reset_button = gr.Button(
                        "Clear",
                        variant="secondary",
                        size="md",
                        elem_id="clear-button",
                        scale=1,
                    )

        status = gr.Markdown(
            "",
            elem_id="status-card",
        )

        seed = gr.Number(value=1234, precision=0, visible=False)
        randomize_seed = gr.Checkbox(value=False, visible=False)
        use_image_adapter = gr.Checkbox(value=False, visible=False)
        adapter_image = gr.Image(type="filepath", visible=False)

        run_button.click(
            fn=upscale_image,
            inputs=[
                source_image,
                scale,
                mode,
                seed,
                randomize_seed,
                adapter_image,
                use_image_adapter,
            ],
            outputs=[result_image, seed, status],
            api_name="upscale",
            concurrency_limit=1,
            concurrency_id="larp_gpu",
            time_limit=60,
            scroll_to_output=True,
            show_progress="full",
        )
        reset_button.click(
            fn=reset_editor,
            outputs=[
                source_image,
                scale,
                mode,
                seed,
                randomize_seed,
                adapter_image,
                use_image_adapter,
                result_image,
                status,
            ],
            queue=False,
            api_visibility="undocumented",
        )
        source_image.clear(
            fn=clear_result,
            outputs=[result_image, status],
            queue=False,
            api_visibility="undocumented",
        )


demo.queue(default_concurrency_limit=1).launch(
    mcp_server=True,
    theme=gr.themes.Base(),
    css=CSS,
)