import io
import random
import re

import numpy as np
import torch
import spaces
import gradio as gr
from diffusers import DiffusionPipeline, AutoPipelineForText2Image, AutoencoderKL
from compel import Compel, ReturnedEmbeddingsType

import r2_uploader

# Per-Space namespace embedded in every uploaded object key. Deliberately opaque
# (not the readable Space name) but stable so the owner can tell assets apart.
# Legend: s01=ImageStudio, s02=LTX2.3-Studio, s03=wan2-2-fp8da-aoti-preview-2.
R2_NAMESPACE = "s01"

# =============================================================================
# Model registry
# =============================================================================
# Two pipelines are exposed through a single UI via the selector below:
#   * Z-Image-Turbo  -> fast, guidance-free distilled model (8 DiT forwards)
#   * NoobAI-XL v1.1 -> SDXL anime model with Compel prompt weighting
MODEL_ZIMAGE = "Z-Image-Turbo"
MODEL_NOOBXL = "NoobAI-XL v1.1"
MODEL_CHOICES = [MODEL_ZIMAGE, MODEL_NOOBXL]

MAX_SEED = np.iinfo(np.int32).max

NOOBXL_NEGATIVE = (
    "lowres, {bad}, error, fewer, extra, missing, worst quality, jpeg artifacts, "
    "bad quality, watermark, unfinished, displeasing, chromatic aberration, signature, "
    "extra digits, artistic error, username, scan, [abstract]"
)

# -----------------------------------------------------------------------------
# Load both pipelines once at startup. On ZeroGPU the `.to("cuda")` calls are
# captured by the runtime and the work happens inside the @spaces.GPU function.
# -----------------------------------------------------------------------------
print("Loading Z-Image-Turbo pipeline...")
zimage_pipe = DiffusionPipeline.from_pretrained(
    "Tongyi-MAI/Z-Image-Turbo",
    torch_dtype=torch.bfloat16,
    low_cpu_mem_usage=False,
)
zimage_pipe.to("cuda")

print("Loading NoobAI-XL v1.1 pipeline...")
noobxl_vae = AutoencoderKL.from_pretrained(
    "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16
)
noobxl_pipe = AutoPipelineForText2Image.from_pretrained(
    "Menyu/noobaiXLNAIXL_epsilonPred11Version",
    vae=noobxl_vae,
    torch_dtype=torch.float16,
    use_safetensors=True,
    add_watermarker=False,
)
noobxl_pipe.to("cuda")

print("Pipelines loaded!")


# =============================================================================
# Compel-based prompt weighting helpers (ported from the NoobXL11 reference
# space so long prompts and (weight:1.2) syntax work for the SDXL model).
# =============================================================================
def parse_prompt_attention(text):
    re_attention = re.compile(r"""
      \\\(|
      \\\)|
      \\\[|
      \\]|
      \\\\|
      \\|
      \(|
      \[|
      :([+-]?[.\d]+)\)|
      \)|
      ]|
      [^\\()\[\]:]+|
      :
      """, re.X)

    res = []
    round_brackets = []
    square_brackets = []

    round_bracket_multiplier = 1.1
    square_bracket_multiplier = 1 / 1.1

    def multiply_range(start_position, multiplier):
        for p in range(start_position, len(res)):
            res[p][1] *= multiplier

    for m in re_attention.finditer(text):
        text = m.group(0)
        weight = m.group(1)

        if text.startswith('\\'):
            res.append([text[1:], 1.0])
        elif text == '(':
            round_brackets.append(len(res))
        elif text == '[':
            square_brackets.append(len(res))
        elif weight is not None and len(round_brackets) > 0:
            multiply_range(round_brackets.pop(), float(weight))
        elif text == ')' and len(round_brackets) > 0:
            multiply_range(round_brackets.pop(), round_bracket_multiplier)
        elif text == ']' and len(square_brackets) > 0:
            multiply_range(square_brackets.pop(), square_bracket_multiplier)
        else:
            parts = re.split(re.compile(r"\s*\bBREAK\b\s*", re.S), text)
            for i, part in enumerate(parts):
                if i > 0:
                    res.append(["BREAK", -1])
                res.append([part, 1.0])

    for pos in round_brackets:
        multiply_range(pos, round_bracket_multiplier)

    for pos in square_brackets:
        multiply_range(pos, square_bracket_multiplier)

    if len(res) == 0:
        res = [["", 1.0]]

    i = 0
    while i + 1 < len(res):
        if res[i][1] == res[i + 1][1]:
            res[i][0] += res[i + 1][0]
            res.pop(i + 1)
        else:
            i += 1

    return res


def prompt_attention_to_invoke_prompt(attention):
    tokens = []
    for text, weight in attention:
        weight = round(weight, 2)
        if weight == 1.0:
            tokens.append(text)
        elif weight < 1.0:
            if weight < 0.8:
                tokens.append(f"({text}){weight}")
            else:
                tokens.append(f"({text})-" + "-" * int((1.0 - weight) * 10))
        else:
            if weight < 1.3:
                tokens.append(f"({text})" + "+" * int((weight - 1.0) * 10))
            else:
                tokens.append(f"({text}){weight}")
    return "".join(tokens)


def merge_embeds(prompt_chanks, compel):
    num_chanks = len(prompt_chanks)
    if num_chanks != 0:
        power_prompt = 1 / (num_chanks * (num_chanks + 1) // 2)
        prompt_embs = compel(prompt_chanks)
        t_list = list(torch.split(prompt_embs, 1, dim=0))
        for i in range(num_chanks):
            t_list[-(i + 1)] = t_list[-(i + 1)] * ((i + 1) * power_prompt)
        prompt_emb = torch.stack(t_list, dim=0).sum(dim=0)
    else:
        prompt_emb = compel('')
    return prompt_emb


def detokenize(chunk, actual_prompt):
    chunk[-1] = chunk[-1].replace('</w>', '')
    chanked_prompt = ''.join(chunk).strip()
    while '</w>' in chanked_prompt:
        if actual_prompt[chanked_prompt.find('</w>')] == ' ':
            chanked_prompt = chanked_prompt.replace('</w>', ' ', 1)
        else:
            chanked_prompt = chanked_prompt.replace('</w>', '', 1)
    actual_prompt = actual_prompt.replace(chanked_prompt, '')
    return chanked_prompt.strip(), actual_prompt.strip()


def tokenize_line(line, tokenizer):
    actual_prompt = line.lower().strip()
    actual_tokens = tokenizer.tokenize(actual_prompt)
    max_tokens = tokenizer.model_max_length - 2
    comma_token = tokenizer.tokenize(',')[0]

    chunks = []
    chunk = []
    for item in actual_tokens:
        chunk.append(item)
        if len(chunk) == max_tokens:
            if chunk[-1] != comma_token:
                for i in range(max_tokens - 1, -1, -1):
                    if chunk[i] == comma_token:
                        actual_chunk, actual_prompt = detokenize(chunk[:i + 1], actual_prompt)
                        chunks.append(actual_chunk)
                        chunk = chunk[i + 1:]
                        break
                else:
                    actual_chunk, actual_prompt = detokenize(chunk, actual_prompt)
                    chunks.append(actual_chunk)
                    chunk = []
            else:
                actual_chunk, actual_prompt = detokenize(chunk, actual_prompt)
                chunks.append(actual_chunk)
                chunk = []
    if chunk:
        actual_chunk, _ = detokenize(chunk, actual_prompt)
        chunks.append(actual_chunk)

    return chunks


def get_embed_new(prompt, pipeline, compel, only_convert_string=False, compel_process_sd=False):
    if compel_process_sd:
        return merge_embeds(tokenize_line(prompt, pipeline.tokenizer), compel)
    else:
        prompt = prompt.replace("((", "(").replace("))", ")").replace("\\", "\\\\\\")

    attention = parse_prompt_attention(prompt)
    global_attention_chanks = []

    for att in attention:
        for chank in att[0].split(','):
            temp_prompt_chanks = tokenize_line(chank, pipeline.tokenizer)
            for small_chank in temp_prompt_chanks:
                temp_dict = {
                    "weight": round(att[1], 2),
                    "lenght": len(pipeline.tokenizer.tokenize(f'{small_chank},')),
                    "prompt": f'{small_chank},'
                }
                global_attention_chanks.append(temp_dict)

    max_tokens = pipeline.tokenizer.model_max_length - 2
    global_prompt_chanks = []
    current_list = []
    current_length = 0
    for item in global_attention_chanks:
        if current_length + item['lenght'] > max_tokens:
            global_prompt_chanks.append(current_list)
            current_list = [[item['prompt'], item['weight']]]
            current_length = item['lenght']
        else:
            if not current_list:
                current_list.append([item['prompt'], item['weight']])
            else:
                if item['weight'] != current_list[-1][1]:
                    current_list.append([item['prompt'], item['weight']])
                else:
                    current_list[-1][0] += f" {item['prompt']}"
            current_length += item['lenght']
    if current_list:
        global_prompt_chanks.append(current_list)

    if only_convert_string:
        return ' '.join([prompt_attention_to_invoke_prompt(i) for i in global_prompt_chanks])

    return merge_embeds([prompt_attention_to_invoke_prompt(i) for i in global_prompt_chanks], compel)


# =============================================================================
# Generation
# =============================================================================
@spaces.GPU
def generate_image(
    model_name,
    prompt,
    negative_prompt,
    use_negative_prompt,
    height,
    width,
    num_inference_steps,
    guidance_scale,
    seed,
    randomize_seed,
    progress=gr.Progress(track_tqdm=True),
):
    """Generate an image from the given prompt using the selected model."""
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    seed = int(seed)

    if model_name == MODEL_NOOBXL:
        generator = torch.Generator().manual_seed(seed)
        compel = Compel(
            tokenizer=[noobxl_pipe.tokenizer, noobxl_pipe.tokenizer_2],
            text_encoder=[noobxl_pipe.text_encoder, noobxl_pipe.text_encoder_2],
            returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
            requires_pooled=[False, True],
            truncate_long_prompts=False,
        )
        if not use_negative_prompt:
            negative_prompt = ""
        conv_prompt = get_embed_new(prompt, noobxl_pipe, compel, only_convert_string=True)
        conv_negative = get_embed_new(negative_prompt, noobxl_pipe, compel, only_convert_string=True)
        conditioning, pooled = compel([conv_prompt, conv_negative])

        image = noobxl_pipe(
            prompt_embeds=conditioning[0:1],
            pooled_prompt_embeds=pooled[0:1],
            negative_prompt_embeds=conditioning[1:2],
            negative_pooled_prompt_embeds=pooled[1:2],
            width=int(width),
            height=int(height),
            guidance_scale=float(guidance_scale),
            num_inference_steps=int(num_inference_steps),
            generator=generator,
            use_resolution_binning=True,
        ).images[0]
        return image, seed

    # Default: Z-Image-Turbo (guidance-free distilled model)
    generator = torch.Generator("cuda").manual_seed(seed)
    image = zimage_pipe(
        prompt=prompt,
        height=int(height),
        width=int(width),
        num_inference_steps=int(num_inference_steps),
        guidance_scale=0.0,
        generator=generator,
    ).images[0]
    return image, seed


def generate_and_upload(
    model_name,
    prompt,
    negative_prompt,
    use_negative_prompt,
    height,
    width,
    num_inference_steps,
    guidance_scale,
    seed,
    randomize_seed,
    request: gr.Request = None,
    progress=gr.Progress(track_tqdm=True),
):
    """Generate, then upload the result to R2 outside the GPU window.

    Returns ``(image, seed, r2_status)``. The image is always the original
    HF-generated asset; ``r2_status`` reports the uploaded filekey on success or
    the error on failure. The caller's unique id (``uid`` cookie) is recorded in
    the uploaded object's metadata.
    """
    image, used = generate_image(
        model_name, prompt, negative_prompt, use_negative_prompt,
        height, width, num_inference_steps, guidance_scale, seed, randomize_seed,
        progress=progress,
    )

    uid = r2_uploader.uid_from_request(request)
    buf = io.BytesIO()
    image.save(buf, format="PNG")
    params = {
        "model": model_name,
        "prompt": prompt,
        "negative_prompt": (negative_prompt if use_negative_prompt else ""),
        "height": int(height),
        "width": int(width),
        "num_inference_steps": int(num_inference_steps),
        "guidance_scale": float(guidance_scale),
        "seed": int(used),
        "uid": uid,
    }
    result = r2_uploader.upload_asset(
        namespace=R2_NAMESPACE,
        prompt=prompt,
        params=params,
        data=buf.getvalue(),
        ext=".png",
        content_type="image/png",
        uid=uid,
    )
    if result.get("ok"):
        status = {"r2_filekey": result["filekey"], "r2_bucket": result["bucket"]}
    else:
        status = {"r2_error": result.get("error", "unknown error")}
    return image, used, status


# Recommended defaults per model: (steps, guidance, height, width)
MODEL_DEFAULTS = {
    MODEL_ZIMAGE: dict(steps=9, guidance=0.0, height=1024, width=1024),
    MODEL_NOOBXL: dict(steps=28, guidance=5.0, height=1536, width=1024),
}


def apply_model_defaults(model_name):
    """Update sliders and helper visibility when the model changes."""
    d = MODEL_DEFAULTS[model_name]
    is_noob = model_name == MODEL_NOOBXL
    return (
        gr.update(value=d["steps"]),
        gr.update(value=d["guidance"], interactive=is_noob,
                  info="Z-Image-Turbo is guidance-free (forced to 0)" if not is_noob else "Classifier-free guidance"),
        gr.update(value=d["height"]),
        gr.update(value=d["width"]),
        gr.update(visible=is_noob),
        gr.update(visible=is_noob),
    )


# Example prompts (work well across both models)
examples = [
    ["Young Chinese woman in red Hanfu, intricate embroidery. Impeccable makeup, red floral forehead pattern. Elaborate high bun, golden phoenix headdress, red flowers, beads. Holds round folding fan with lady, trees, bird. Neon lightning-bolt lamp, bright yellow glow, above extended left palm. Soft-lit outdoor night background, silhouetted tiered pagoda, blurred colorful distant lights."],
    ["A majestic dragon soaring through clouds at sunset, scales shimmering with iridescent colors, detailed fantasy art style"],
    ["Cozy coffee shop interior, warm lighting, rain on windows, plants on shelves, vintage aesthetic, photorealistic"],
    ["1girl, nahida (genshin impact), white dress, green hair, looking at viewer, masterpiece, best quality, very aesthetic"],
    ["Portrait of a wise old wizard with a long white beard, holding a glowing crystal staff, magical forest background"],
]

# Custom theme with modern aesthetics (Gradio 6)
custom_theme = gr.themes.Soft(
    primary_hue="yellow",
    secondary_hue="amber",
    neutral_hue="slate",
    font=gr.themes.GoogleFont("Inter"),
    text_size="lg",
    spacing_size="md",
    radius_size="lg"
).set(
    button_primary_background_fill="*primary_500",
    button_primary_background_fill_hover="*primary_600",
    block_title_text_weight="600",
)

APP_CSS = """
.header-text h1 {
    font-size: 2.5rem !important;
    font-weight: 700 !important;
    margin-bottom: 0.5rem !important;
    background: linear-gradient(135deg, #fbbf24 0%, #f59e0b 100%);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
}
.header-text p {
    font-size: 1.1rem !important;
    color: #64748b !important;
    margin-top: 0 !important;
}
.footer-text { padding: 1rem 0; }
.footer-text a {
    color: #f59e0b !important;
    text-decoration: none !important;
    font-weight: 500;
}
.footer-text a:hover { text-decoration: underline !important; }
@media (max-width: 768px) {
    .header-text h1 { font-size: 1.8rem !important; }
    .header-text p { font-size: 1rem !important; }
}
button, .gr-button { transition: all 0.2s ease !important; }
button:hover, .gr-button:hover {
    transform: translateY(-1px);
    box-shadow: 0 4px 12px rgba(0, 0, 0, 0.15) !important;
}
.gradio-container {
    max-width: 1400px !important;
    margin: 0 auto !important;
}
"""

# Build the Gradio interface. In Gradio 6.x, theme/css/footer_links/mcp_server
# are arguments to demo.launch() (see bottom of file), not to Blocks().
with gr.Blocks(fill_height=True) as demo:
    gr.Markdown(
        """
        # 🎨 Image Studio
        **Ultra-fast AI image generation** • Choose a model and create stunning images
        """,
        elem_classes="header-text"
    )

    with gr.Row(equal_height=False):
        # Left column - Input controls
        with gr.Column(scale=1, min_width=320):
            model_name = gr.Dropdown(
                choices=MODEL_CHOICES,
                value=MODEL_ZIMAGE,
                label="🧠 Model",
                info="Z-Image-Turbo: fast & general • NoobAI-XL: anime / illustration",
            )

            prompt = gr.Textbox(
                label="✨ Your Prompt",
                placeholder="Describe the image you want to create...",
                lines=5,
                max_lines=10,
                autofocus=True,
            )

            with gr.Accordion("⚙️ Advanced Settings", open=False):
                use_negative_prompt = gr.Checkbox(
                    label="Use Negative Prompt",
                    value=True,
                    visible=False,  # NoobXL only
                )
                negative_prompt = gr.Textbox(
                    label="🚫 Negative Prompt",
                    value=NOOBXL_NEGATIVE,
                    lines=3,
                    max_lines=6,
                    visible=False,  # NoobXL only
                )

                with gr.Row():
                    height = gr.Slider(
                        minimum=512, maximum=2048, value=1024, step=64,
                        label="Height", info="Image height in pixels",
                    )
                    width = gr.Slider(
                        minimum=512, maximum=2048, value=1024, step=64,
                        label="Width", info="Image width in pixels",
                    )

                num_inference_steps = gr.Slider(
                    minimum=1, maximum=50, value=9, step=1,
                    label="Inference Steps",
                    info="Z-Image: ~9 • NoobAI-XL: ~28",
                )

                guidance_scale = gr.Slider(
                    minimum=0.0, maximum=10.0, value=0.0, step=0.1,
                    label="Guidance Scale",
                    info="Z-Image-Turbo is guidance-free (forced to 0)",
                    interactive=False,
                )

                with gr.Row():
                    randomize_seed = gr.Checkbox(label="🎲 Random Seed", value=True)
                    seed = gr.Number(label="Seed", value=42, precision=0, minimum=0, maximum=MAX_SEED)

            generate_btn = gr.Button(
                "🚀 Generate Image", variant="primary", size="lg", scale=1
            )

            gr.Examples(
                examples=examples,
                inputs=[prompt],
                label="💡 Try these prompts",
                examples_per_page=5,
            )

        # Right column - Output
        with gr.Column(scale=1, min_width=320):
            output_image = gr.Image(
                label="Generated Image",
                type="pil",
                format="png",
                show_label=False,
                height=600,
                buttons=["download", "share"],
            )
            used_seed = gr.Number(
                label="🎲 Seed Used", interactive=False, container=True,
            )
            r2_status = gr.JSON(label="☁️ R2 Upload")

    gr.Markdown(
        """
        ---
        <div style="text-align: center; opacity: 0.7; font-size: 0.9em; margin-top: 1rem;">
        <strong>Models:</strong>
        <a href="https://huggingface.co/Tongyi-MAI/Z-Image-Turbo" target="_blank">Z-Image-Turbo</a> •
        <a href="https://huggingface.co/Menyu/noobaiXLNAIXL_epsilonPred11Version" target="_blank">NoobAI-XL v1.1</a> •
        <strong>Demo by:</strong> <a href="https://x.com/realmrfakename" target="_blank">@mrfakename</a>
        </div>
        """,
        elem_classes="footer-text"
    )

    # Update sliders / helper visibility when the model changes
    model_name.change(
        fn=apply_model_defaults,
        inputs=[model_name],
        outputs=[num_inference_steps, guidance_scale, height, width,
                 use_negative_prompt, negative_prompt],
    )

    gen_inputs = [
        model_name, prompt, negative_prompt, use_negative_prompt,
        height, width, num_inference_steps, guidance_scale, seed, randomize_seed,
    ]

    # api_name pinned so the generator's "/generate_image" endpoint keeps
    # resolving even though the click now runs the upload wrapper.
    generate_btn.click(
        fn=generate_and_upload, inputs=gen_inputs,
        outputs=[output_image, used_seed, r2_status], api_name="generate_image",
    )
    prompt.submit(
        fn=generate_and_upload, inputs=gen_inputs,
        outputs=[output_image, used_seed, r2_status], show_api=False,
    )

if __name__ == "__main__":
    demo.launch(
        theme=custom_theme,
        css=APP_CSS,
        footer_links=["api", "gradio"],
        mcp_server=True,
    )