import os

os.environ.setdefault("HF_HOME", "/tmp/.cache/huggingface")
os.environ.setdefault("HF_MODULES_CACHE", "/tmp/hf_modules")
os.environ.setdefault("MPLCONFIGDIR", "/tmp/matplotlib")
os.environ.setdefault("GRADIO_ANALYTICS_ENABLED", "False")
os.environ.setdefault("HF_HUB_DISABLE_SYMLINKS_WARNING", "1")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

import spaces

import time
from pathlib import Path
from typing import Optional, Tuple

import gradio as gr
from gradio.themes.utils import colors, sizes
from gradio_imageslider import ImageSlider
from huggingface_hub import snapshot_download
from PIL import Image, ImageOps

import torch
from diffusers.pipelines import FluxPipeline

from src.flux.condition import Condition
from src.flux.generate import generate, seed_everything
from tools.color_fix import adain_color_fix


_IMAGE_SLIDER_GET_CONFIG = ImageSlider.get_config


def _imageslider_get_config_with_buttons(self, cls=None):
    # gradio_imageslider 0.0.20's Gradio 6 front-end expects this prop.
    config = _IMAGE_SLIDER_GET_CONFIG(self, cls)
    config.setdefault("buttons", ["download", "fullscreen"])
    return config


ImageSlider.get_config = _imageslider_get_config_with_buttons


FLUX_MODEL_ID = os.environ.get("FLUX_MODEL_ID", "black-forest-labs/FLUX.1-dev")
ASASR_MODEL_ID = os.environ.get("ASASR_MODEL_ID", "wafer-bob/ASASR")
SR_LORA_NAME = "sr_lora/pytorch_lora_weights_v2.safetensors"
DPO_LORA_NAME = "dpo_lora/adapter_model.safetensors"
TARGET_RESOLUTION = 512
LR_RESOLUTION = TARGET_RESOLUTION // 4
NUM_INFERENCE_STEPS = 28
GUIDANCE_SCALE = 3.5
EXAMPLES_DIR = Path("examples")
EXAMPLE_NAMES = (
    "portrait.png",
    "landscape.png",
    "text.png",
    "architecture.png",
    "wildlife.png",
    "texture.png",
)
EXAMPLE_FILES = [
    str(EXAMPLES_DIR / name)
    for name in EXAMPLE_NAMES
    if (EXAMPLES_DIR / name).is_file()
]
PIPELINE: Optional[FluxPipeline] = None
FLUX_LOCAL_DIR: Optional[str] = None
ASASR_LOCAL_DIR: Optional[str] = None
PIPELINE_LOAD_SECONDS: Optional[float] = None
LAST_INFERENCE_SECONDS: Optional[float] = None
STARTUP_NOTE = "Assets will be resolved on the first request."


def _token() -> Optional[str]:
    return os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")


def _require_token() -> str:
    token = _token()
    if not token:
        raise gr.Error(
            "HF_TOKEN is not set. Add it as a Space secret with access to "
            "black-forest-labs/FLUX.1-dev."
        )
    return token


def _prepare_assets() -> Tuple[str, str]:
    global FLUX_LOCAL_DIR, ASASR_LOCAL_DIR, STARTUP_NOTE
    if FLUX_LOCAL_DIR and ASASR_LOCAL_DIR:
        return FLUX_LOCAL_DIR, ASASR_LOCAL_DIR

    token = _require_token()
    start = time.perf_counter()
    print("[ASASR] Resolving FLUX.1-dev and ASASR LoRA assets...")
    FLUX_LOCAL_DIR = snapshot_download(
        repo_id=FLUX_MODEL_ID,
        token=token,
        ignore_patterns=[
            "*.bin",
            "*.onnx",
            "*.msgpack",
            "examples/*",
            "ae.safetensors",
            "dev_grid.jpg",
            "flux1-dev.safetensors",
        ],
        max_workers=8,
    )
    ASASR_LOCAL_DIR = snapshot_download(
        repo_id=ASASR_MODEL_ID,
        token=token,
        allow_patterns=[
            SR_LORA_NAME,
            DPO_LORA_NAME,
            "dpo_lora/adapter_config.json",
        ],
        max_workers=4,
    )
    elapsed = time.perf_counter() - start
    STARTUP_NOTE = f"Model assets resolved in {elapsed:.1f}s."
    print(f"[ASASR] Asset resolution complete in {elapsed:.1f}s.")
    return FLUX_LOCAL_DIR, ASASR_LOCAL_DIR


def _startup_prefetch() -> None:
    global STARTUP_NOTE
    if os.environ.get("ASASR_PREFETCH", "1") != "1":
        STARTUP_NOTE = "Startup prefetch is disabled."
        print("[ASASR] Startup prefetch disabled.")
        return
    if not _token():
        STARTUP_NOTE = "HF_TOKEN is missing; assets will be resolved after the secret is set."
        print("[ASASR] HF_TOKEN missing; startup prefetch deferred.")
        return
    try:
        _prepare_assets()
    except Exception as exc:
        STARTUP_NOTE = f"Startup prefetch deferred: {type(exc).__name__}: {exc}"
        print(f"[ASASR] Startup prefetch deferred: {type(exc).__name__}: {exc}")


def _get_pipeline() -> FluxPipeline:
    global PIPELINE, PIPELINE_LOAD_SECONDS
    if PIPELINE is not None:
        return PIPELINE

    flux_dir, asasr_dir = _prepare_assets()
    sr_path = Path(asasr_dir) / SR_LORA_NAME
    dpo_path = Path(asasr_dir) / DPO_LORA_NAME

    start = time.perf_counter()
    print("[ASASR] Loading FLUX.1-dev pipeline and dual LoRAs onto cuda...")
    pipe = FluxPipeline.from_pretrained(
        flux_dir,
        torch_dtype=torch.bfloat16,
        local_files_only=True,
    ).to("cuda")
    pipe.load_lora_weights(
        sr_path.parent.as_posix(),
        weight_name=sr_path.name,
        adapter_name="sr",
    )
    pipe.load_lora_weights(
        dpo_path.parent.as_posix(),
        weight_name=dpo_path.name,
        adapter_name="dpo",
    )
    pipe.set_adapters(["sr", "dpo"], adapter_weights=[1.0, 1.0])
    pipe.set_progress_bar_config(disable=True)
    PIPELINE = pipe
    PIPELINE_LOAD_SECONDS = time.perf_counter() - start
    print(f"[ASASR] Pipeline ready in {PIPELINE_LOAD_SECONDS:.1f}s.")
    return PIPELINE


def _center_square(image: Image.Image) -> Image.Image:
    image = ImageOps.exif_transpose(image).convert("RGB")
    width, height = image.size
    side = min(width, height)
    left = (width - side) // 2
    top = (height - side) // 2
    return image.crop((left, top, left + side, top + side))


def _prepare_lr_image(image: Image.Image) -> Image.Image:
    return _center_square(image).resize(
        (LR_RESOLUTION, LR_RESOLUTION),
        Image.Resampling.LANCZOS,
    )


def _prepare_condition_image(lr_image: Image.Image) -> Image.Image:
    return lr_image.resize(
        (TARGET_RESOLUTION, TARGET_RESOLUTION),
        Image.Resampling.BICUBIC,
    )


def _prepare_slider_lr_image(lr_image: Image.Image) -> Image.Image:
    return lr_image.resize(
        (TARGET_RESOLUTION, TARGET_RESOLUTION),
        Image.Resampling.NEAREST,
    )


def _initial_slider_value() -> Tuple[Image.Image, Image.Image]:
    preview = Image.new("RGB", (TARGET_RESOLUTION, TARGET_RESOLUTION), "#f8fafc")
    return preview, preview.copy()


def _gpu_duration(*args, **kwargs) -> int:
    value = os.environ.get("ASASR_GPU_DURATION", "45")
    try:
        duration = int(value)
    except ValueError:
        duration = 240
    return max(30, min(duration, 300))


@spaces.GPU(duration=1)
def _zerogpu_probe() -> str:
    return "ready"


@spaces.GPU(duration=_gpu_duration)
def super_resolve(
    input_image: Image.Image,
    progress: gr.Progress = gr.Progress(track_tqdm=True),
):
    global LAST_INFERENCE_SECONDS
    if input_image is None:
        raise gr.Error("Upload or choose a 128 x 128 low-resolution image first.")

    pipe = _get_pipeline()
    lr_image = _prepare_lr_image(input_image)
    condition_image = _prepare_condition_image(lr_image)
    slider_lr_image = _prepare_slider_lr_image(lr_image)
    condition = Condition("sr", condition_image)

    seed_everything(42)
    start = time.perf_counter()
    result = generate(
        pipe,
        prompt="",
        conditions=[condition],
        default_lora=True,
        height=TARGET_RESOLUTION,
        width=TARGET_RESOLUTION,
        num_inference_steps=NUM_INFERENCE_STEPS,
        guidance_scale=GUIDANCE_SCALE,
    ).images[0]
    result = adain_color_fix(result, condition_image).convert("RGB")
    LAST_INFERENCE_SECONDS = time.perf_counter() - start
    load_text = (
        f"Model load: {PIPELINE_LOAD_SECONDS:.1f}s. "
        if PIPELINE_LOAD_SECONDS is not None
        else ""
    )
    status = (
        f"Input size: {lr_image.width} x {lr_image.height}. "
        f"Output size: {result.width} x {result.height}. "
        f"{load_text}Inference: {LAST_INFERENCE_SECONDS:.1f}s."
    )
    print(f"[ASASR] {status}")
    return (slider_lr_image, result), status


_startup_prefetch()


PAPER_URL = "https://arxiv.org/abs/2605.23264"
GITHUB_URL = "https://github.com/wafer-bob/ASASR"
MODEL_URL = "https://huggingface.co/wafer-bob/ASASR"

DESCRIPTION_MD = (
    "ASASR turns a low-resolution image into a faithful **512 × 512** reconstruction "
    "using a **FLUX.1-dev** backbone with **dual-LoRA inference** — a base SR LoRA plus an "
    "AS-DPO alignment LoRA. Upload a roughly **128 × 128** image or choose an example, "
    "then run the 28-step x4 sampler and scrub the slider to compare input and output."
)

HERO_HTML = f"""
<div class="asasr-hero">
  <div class="asasr-badges">
    <span class="asasr-badge asasr-badge--icml">ICML 2026</span>
    <span class="asasr-badge asasr-badge--soft">FLUX.1-dev · Dual-LoRA · x4 SR</span>
  </div>
  <h1 class="asasr-title">
    Coloring the Noise: Adversarial Sobolev Alignment<br>
    <span class="asasr-title--sub">for Faithful Image Super-Resolution</span>
  </h1>
  <p class="asasr-authors">
    Hongbo Wang · Huaibo Huang · Pin Wang · Jinhua Hao · Chao Zhou · Ran He
  </p>
  <div class="asasr-links">
    <a class="asasr-link asasr-link--primary" href="{PAPER_URL}" target="_blank" rel="noopener">
      <svg viewBox="0 0 24 24" width="16" height="16" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><path d="M14 2H6a2 2 0 0 0-2 2v16a2 2 0 0 0 2 2h12a2 2 0 0 0 2-2V8z"/><path d="M14 2v6h6M16 13H8M16 17H8M10 9H8"/></svg>
      Paper · arXiv
    </a>
    <a class="asasr-link" href="{GITHUB_URL}" target="_blank" rel="noopener">
      <svg viewBox="0 0 24 24" width="16" height="16" fill="currentColor"><path d="M12 .5A11.5 11.5 0 0 0 .5 12 11.5 11.5 0 0 0 8.4 23c.6.1.8-.3.8-.6v-2c-3.2.7-3.9-1.5-3.9-1.5-.5-1.3-1.3-1.7-1.3-1.7-1.1-.7.1-.7.1-.7 1.2.1 1.8 1.2 1.8 1.2 1 1.8 2.8 1.3 3.5 1 .1-.8.4-1.3.7-1.6-2.6-.3-5.3-1.3-5.3-5.7 0-1.3.4-2.3 1.2-3.1-.1-.3-.5-1.5.1-3.2 0 0 1-.3 3.3 1.2a11.4 11.4 0 0 1 6 0C17 4.6 18 4.9 18 4.9c.6 1.7.2 2.9.1 3.2.8.8 1.2 1.8 1.2 3.1 0 4.4-2.7 5.4-5.3 5.7.4.4.8 1.1.8 2.2v3.3c0 .3.2.7.8.6A11.5 11.5 0 0 0 23.5 12 11.5 11.5 0 0 0 12 .5z"/></svg>
      Code · GitHub
    </a>
    <a class="asasr-link" href="{MODEL_URL}" target="_blank" rel="noopener">
      <svg viewBox="0 0 24 24" width="16" height="16" fill="currentColor"><path d="M12 2 2 7v10l10 5 10-5V7zm0 2.2 7.5 3.8L12 11.8 4.5 8zm-8 5.3 7 3.5v7.6l-7-3.5zm16 0v7.6l-7 3.5v-7.6z"/></svg>
      Model · Hugging Face
    </a>
  </div>
</div>
"""

FOOTER_HTML = """
<div class="asasr-footer">
  <div class="asasr-license">
    <strong>License · CC-BY-NC-4.0.</strong>
    ASASR weights and this demo are for non-commercial research use only, and
    additionally inherit the non-commercial terms of FLUX.1-dev.
  </div>
  <details class="asasr-citation">
    <summary>Citation</summary>
    <pre>@inproceedings{wang2026asasr,
  title     = {Coloring the Noise: Adversarial Sobolev Alignment for Faithful Image Super-Resolution},
  author    = {Wang, Hongbo and Huang, Huaibo and Wang, Pin and Hao, Jinhua and Zhou, Chao and He, Ran},
  booktitle = {International Conference on Machine Learning (ICML)},
  year      = {2026}
}</pre>
  </details>
</div>
"""

CSS = """
#asasr-root { max-width: 1180px; margin: 0 auto; }

/* ---------- Hero ---------- */
.asasr-hero { text-align: center; padding: 2.4rem 1rem 1.4rem; }
.asasr-badges { display: flex; gap: .55rem; justify-content: center; flex-wrap: wrap; margin-bottom: 1.1rem; }
.asasr-badge {
  display: inline-flex; align-items: center; gap: .4rem;
  padding: .32rem .8rem; border-radius: 999px;
  font-size: .72rem; font-weight: 600; letter-spacing: .05em; text-transform: uppercase;
}
.asasr-badge--icml { background: linear-gradient(135deg, #6366f1, #8b5cf6); color: #fff; box-shadow: 0 4px 14px rgba(99,102,241,.35); }
.asasr-badge--soft { background: rgba(99,102,241,.10); color: #4338ca; border: 1px solid rgba(99,102,241,.25); }
.asasr-title {
  font-size: clamp(1.55rem, 3.6vw, 2.55rem); line-height: 1.12; font-weight: 800;
  margin: 0 auto; max-width: 940px; letter-spacing: -0.025em; color: var(--body-text-color);
}
.asasr-title--sub { color: var(--body-text-color-subdued); font-weight: 600; }
.asasr-authors { color: var(--body-text-color-subdued); margin: .85rem 0 1.25rem; font-size: .98rem; }
.asasr-links { display: inline-flex; gap: .55rem; flex-wrap: wrap; justify-content: center; }
.asasr-link {
  display: inline-flex; align-items: center; gap: .45rem; padding: .5rem 1.05rem;
  border-radius: 999px; text-decoration: none; font-weight: 600; font-size: .9rem;
  border: 1px solid var(--border-color-primary); color: var(--body-text-color);
  background: var(--background-fill-primary); transition: all .15s ease;
}
.asasr-link:hover { border-color: #6366f1; color: #4338ca; background: rgba(99,102,241,.08); }
.asasr-link--primary { background: #4338ca; color: #fff; border-color: #4338ca; box-shadow: 0 4px 14px rgba(67,56,202,.3); }
.asasr-link--primary:hover { background: #3730a3; color: #fff; border-color: #3730a3; }

/* ---------- Description ---------- */
.asasr-desc { max-width: 760px; margin: 0 auto 1.6rem; text-align: center; color: var(--body-text-color-subdued); font-size: 1.02rem; line-height: 1.55; }
.asasr-desc strong { color: var(--body-text-color); }

/* ---------- Section labels ---------- */
.asasr-section-label {
  display: flex; align-items: center; gap: .6rem; margin: .2rem 0 .9rem;
  font-size: .78rem; font-weight: 700; letter-spacing: .14em; text-transform: uppercase;
  color: var(--body-text-color-subdued);
}
.asasr-section-label::after { content: ""; flex: 1; height: 1px; background: var(--border-color-primary); }

/* ---------- Interface panels ---------- */
.asasr-panel { border: 1px solid var(--border-color-primary); border-radius: 16px; padding: .9rem .9rem 1.1rem; background: var(--background-fill-primary); }
.asasr-panel--input { display: flex; flex-direction: column; gap: .8rem; }

/* Input image frame */
.asasr-input-image { border-radius: 14px !important; border: 1px solid var(--border-color-primary) !important; overflow: hidden; }
.asasr-input-image .image-frame, .asasr-input-image img { border-radius: 14px !important; }

/* The centerpiece slider gets a subtle accent ring */
.asasr-slider-wrap { position: relative; border-radius: 18px; padding: .55rem; background: linear-gradient(180deg, rgba(99,102,241,.10), rgba(139,92,246,.05)); border: 1px solid rgba(99,102,241,.18); }
.asasr-slider-wrap .component-wrapper { border: none !important; }

/* Run button */
.asasr-run { width: 100%; height: 52px !important; font-size: 1rem !important; font-weight: 700 !important; border-radius: 12px !important; letter-spacing: .01em; }

/* Status line */
.asasr-status textarea { font-family: var(--font-mono, ui-monospace, Menlo, monospace); font-size: .82rem !important; color: var(--body-text-color-subdued) !important; }

/* ---------- Examples ---------- */
#asasr-examples,
#asasr-examples .examples,
#asasr-examples .table-wrap,
#asasr-examples .table,
#asasr-examples table,
#asasr-examples tbody,
#asasr-examples tr {
  max-height: none !important;
  overflow: hidden !important;
}
#asasr-examples img {
  width: 50px !important;
  height: 50px !important;
  min-width: 50px !important;
  min-height: 50px !important;
  max-width: 50px !important;
  max-height: 50px !important;
  object-fit: cover !important;
  border-radius: 8px !important;
}
#asasr-examples button,
#asasr-examples .example,
#asasr-examples td {
  width: 58px !important;
  height: 58px !important;
  min-width: 58px !important;
  max-width: 58px !important;
  padding: 4px !important;
  overflow: hidden !important;
}
#asasr-examples,
#asasr-examples * {
  scrollbar-width: none !important;
}
#asasr-examples::-webkit-scrollbar,
#asasr-examples *::-webkit-scrollbar {
  display: none !important;
}

/* ---------- Footer ---------- */
.asasr-footer { margin-top: 1.8rem; padding-top: 1.2rem; border-top: 1px solid var(--border-color-primary); }
.asasr-license { font-size: .82rem; color: var(--body-text-color-subdued); line-height: 1.5; }
.asasr-license strong { color: var(--body-text-color); }
.asasr-citation { margin-top: 1rem; border: 1px solid var(--border-color-primary); border-radius: 12px; background: var(--background-fill-secondary); overflow: hidden; }
.asasr-citation summary { cursor: pointer; padding: .65rem 1rem; font-weight: 600; font-size: .9rem; list-style: none; display: flex; align-items: center; gap: .5rem; }
.asasr-citation summary::before { content: "▸"; color: var(--body-text-color-subdued); transition: transform .15s ease; }
.asasr-citation[open] summary::before { transform: rotate(90deg); }
.asasr-citation summary::-webkit-details-marker { display: none; }
.asasr-citation pre { margin: 0; padding: 1rem 1.1rem; font-family: var(--font-mono, ui-monospace, Menlo, monospace); font-size: .82rem; line-height: 1.5; overflow-x: auto; background: transparent; border-top: 1px solid var(--border-color-primary); }
"""

THEME = gr.themes.Soft(
    primary_hue=colors.indigo,
    secondary_hue=colors.violet,
    neutral_hue=colors.slate,
    radius_size=sizes.radius_lg,
    text_size=sizes.text_lg,
)


with gr.Blocks(
    title="ASASR · Faithful x4 Image Super-Resolution",
    elem_id="asasr-root",
) as demo:
    gr.HTML(HERO_HTML, container=False, padding=False, apply_default_css=False)
    gr.Markdown(DESCRIPTION_MD, elem_classes=["asasr-desc"])

    with gr.Row(equal_height=False):
        with gr.Column(scale=1, min_width=320):
            gr.HTML('<div class="asasr-section-label">Input</div>')
            with gr.Column(elem_classes=["asasr-panel", "asasr-panel--input"]):
                input_image = gr.Image(
                    label="Low-resolution input (≈128×128)",
                    type="pil",
                    height=360,
                    elem_classes=["asasr-input-image"],
                )
                run_button = gr.Button(
                    "✦ Super-Resolve",
                    variant="primary",
                    elem_classes=["asasr-run"],
                )
        with gr.Column(scale=2, min_width=480):
            gr.HTML('<div class="asasr-section-label">Result · drag to compare</div>')
            with gr.Column(elem_classes=["asasr-slider-wrap"]):
                comparison_slider = ImageSlider(
                    label="Input LR  |  ASASR HR",
                    value=_initial_slider_value(),
                    type="pil",
                    height=560,
                    position=0.5,
                    interactive=False,
                    slider_color="#4338ca",
                )
            status = gr.Textbox(
                label="Runtime",
                value=STARTUP_NOTE,
                interactive=False,
                elem_classes=["asasr-status"],
            )

    gr.HTML('<div class="asasr-section-label">Examples</div>')
    gr.Examples(
        examples=[[path] for path in EXAMPLE_FILES],
        inputs=input_image,
        outputs=[comparison_slider, status],
        fn=super_resolve,
        cache_examples=True,
        cache_mode="lazy",
        examples_per_page=6,
        label="Example inputs",
        elem_id="asasr-examples",
        run_on_click=True,
    )

    gr.HTML(FOOTER_HTML, container=False, padding=False, apply_default_css=False)

    run_button.click(
        fn=super_resolve,
        inputs=input_image,
        outputs=[comparison_slider, status],
        api_name="super_resolve",
    )

if __name__ == "__main__":
    demo.launch(theme=THEME, css=CSS)