import gradio as gr
import os
import random
import torch
from diffusers import StableDiffusionPipeline

HF_TOKEN = os.environ.get("HF_TOKEN", "")
MODEL_ID = "paceailab/WaxFashionStableDiffusion"

pipe = None

def load_pipeline():
    global pipe

    if pipe is not None:
        return pipe

    if not torch.cuda.is_available():
        raise gr.Error(
            "GPU is not available. Please change your Hugging Face Space hardware to GPU, for example T4 Small."
        )

    pipe = StableDiffusionPipeline.from_pretrained(
        MODEL_ID,
        torch_dtype=torch.float16,
        token=HF_TOKEN if HF_TOKEN else None,
        safety_checker=None
    )

    pipe = pipe.to("cuda")
    pipe.enable_attention_slicing()
    try:
        pipe.enable_xformers_memory_efficient_attention()
    except Exception:
        pass

    return pipe


def build_prompt(color, shape, theme):
    palette_map = {
        "Indigo Dusk": "indigo blue and amber colorway",
        "Terracotta": "terracotta red and black colorway",
        "Saharan Ochre": "golden ochre and cream colorway",
        "Forest Kente": "forest green and amber colorway",
        "Royal Plum": "royal purple and gold colorway",
        "Ebony Bone": "black and bone white colorway"
    }

    motif_map = {
        "Circles": "circular concentric motifs",
        "Diamonds": "diamond geometric shapes",
        "Squares": "nested square patterns",
        "Stars": "star motifs",
        "Comb": "comb stripe patterns",
        "Adinkra": "Adinkra symbol patterns"
    }

    return (
        f"African wax print textile, {theme.lower()} style, "
        f"{palette_map.get(color, color + ' colorway')}, "
        f"{motif_map.get(shape, shape)}, high quality Ankara fabric pattern"
    )


def generate_image(color, shape, theme, num_images, progress=gr.Progress()):
    prompt = build_prompt(color, shape, theme)
    generator_pipe = load_pipeline()

    results = []

    try:
        n = max(1, min(3, int(str(num_images).strip())))
    except Exception:
        n = 1

    for i in range(n):
        progress(i / n, desc=f"Generating pattern {i + 1} of {n}")

        try:
            full_prompt = f"{prompt}, variation {random.randint(0, 99999)}"
            with torch.inference_mode():
                image = generator_pipe(
                    full_prompt,
                    num_inference_steps=30,
                    guidance_scale=7.5
                ).images[0]
            results.append(image)

        except torch.cuda.OutOfMemoryError:
            torch.cuda.empty_cache()
            try:
                with torch.inference_mode():
                    image = generator_pipe(
                        full_prompt,
                        num_inference_steps=30,
                        guidance_scale=7.5
                    ).images[0]
                results.append(image)
            except Exception as e:
                raise gr.Error(f"GPU out of memory. Try generating 1 image at a time. ({e})")

        except Exception as e:
            raise gr.Error(f"Generation failed: {e}")

        finally:
            torch.cuda.empty_cache()

    progress(1.0)

    while len(results) < 3:
        results.append(None)

    return results[0], results[1], results[2], prompt


CSS = """
@import url('https://fonts.googleapis.com/css2?family=Playfair+Display:ital,wght@0,700;0,900;1,400;1,700&family=DM+Sans:wght@300;400;500;600&display=swap');
* { box-sizing: border-box; margin: 0; padding: 0; }
body, .gradio-container {
    font-family: 'DM Sans', sans-serif !important;
    background: #F5EDE0 !important;
    min-height: 100vh;
    color: #1A1410 !important;
}
footer, .gr-prose { display: none !important; }
#kente-stripe {
    height: 10px;
    background: repeating-linear-gradient(90deg,
        #C84B31 0px, #C84B31 20px,
        #E8A33D 20px, #E8A33D 40px,
        #1A1410 40px, #1A1410 60px,
        #2D5F3F 60px, #2D5F3F 80px,
        #1E4D6B 80px, #1E4D6B 100px,
        #C84B31 100px, #C84B31 120px,
        #E8A33D 120px, #E8A33D 140px,
        #1A1410 140px, #1A1410 160px);
    position: relative;
    z-index: 10;
}
#app-header {
    padding: 40px 40px 36px;
    text-align: center;
    border-bottom: 1px solid rgba(26,20,16,0.15);
    animation: fadeDown 0.6s ease;
    position: relative;
    z-index: 1;
}
#app-header .eyebrow {
    font-family: 'DM Sans', sans-serif;
    font-size: 10px;
    letter-spacing: 2.5px;
    text-transform: uppercase;
    color: #3A2E25;
    margin-bottom: 20px;
    display: flex;
    align-items: center;
    justify-content: center;
    gap: 8px;
}
#app-header .eyebrow::before {
    content: '';
    display: inline-block;
    width: 8px;
    height: 8px;
    background: #C84B31;
    transform: rotate(45deg);
}
#app-header h1 {
    font-family: 'Playfair Display', serif;
    font-size: 52px;
    font-weight: 900;
    line-height: 1.1;
    color: #1A1410;
    letter-spacing: -1px;
    margin-bottom: 20px;
}
#app-header h1 em {
    font-style: italic;
    font-weight: 700;
    color: #C84B31;
}
#app-header .subtitle {
    font-family: 'DM Sans', sans-serif;
    font-size: 15px;
    color: #3A2E25;
    max-width: 520px;
    margin: 0 auto;
    line-height: 1.7;
}
.section-heading {
    display: flex;
    align-items: center;
    gap: 12px;
    margin: 20px 0 12px;
}
.section-num {
    font-family: 'Playfair Display', serif;
    font-size: 18px;
    font-weight: 900;
    color: #1A1410;
    flex-shrink: 0;
}
.section-lbl {
    font-family: 'DM Sans', sans-serif;
    font-size: 9px;
    letter-spacing: 2.5px;
    text-transform: uppercase;
    color: #3A2E25;
    flex-shrink: 0;
}
.section-line {
    flex: 1;
    height: 1px;
    background: rgba(26,20,16,0.15);
}
/* Colorway Radio */
#colorway-radio > div {
    display: grid !important;
    grid-template-columns: 1fr 1fr !important;
    gap: 10px !important;
}
#colorway-radio label {
    background-color: #EDE3D1 !important;
    border: 2px solid transparent !important;
    border-radius: 4px !important;
    cursor: pointer !important;
    overflow: hidden !important;
    padding: 0 !important;
    margin: 0 !important;
    display: flex !important;
    flex-direction: column !important;
    box-shadow: 2px 2px 0 rgba(26,20,16,0.1) !important;
    transition: all 0.15s ease !important;
    background-repeat: no-repeat !important;
    background-size: 100% 48px !important;
    background-position: top !important;
    padding-top: 48px !important;
}
#colorway-radio label:hover {
    box-shadow: 3px 3px 0 #1A1410 !important;
    transform: translate(-1px,-1px) !important;
}
#colorway-radio label:has(input:checked),
#colorway-radio label.selected {
    border-color: #1A1410 !important;
    box-shadow: 3px 3px 0 #1A1410 !important;
}
#colorway-radio input[type="radio"] {
    display: none !important;
}
#colorway-radio label:nth-child(1) {
    background-image: linear-gradient(90deg, #1E4D6B 25%, #2A6190 50%, #F5EDE0 75%, #E8A33D 100%) !important;
}
#colorway-radio label:nth-child(2) {
    background-image: linear-gradient(90deg, #C84B31 25%, #9C3A22 50%, #F5EDE0 75%, #1A1410 100%) !important;
}
#colorway-radio label:nth-child(3) {
    background-image: linear-gradient(90deg, #E8A33D 25%, #C8860B 50%, #F5EDE0 75%, #FFFFFF 100%) !important;
}
#colorway-radio label:nth-child(4) {
    background-image: linear-gradient(90deg, #2D5F3F 25%, #1F4530 50%, #E8A33D 75%, #F5EDE0 100%) !important;
}
#colorway-radio label:nth-child(5) {
    background-image: linear-gradient(90deg, #5C2A4D 25%, #7A3A66 50%, #E8A33D 75%, #F5EDE0 100%) !important;
}
#colorway-radio label:nth-child(6) {
    background-image: linear-gradient(90deg, #1A1410 25%, #3A2E25 50%, #F5EDE0 75%, #C84B31 100%) !important;
}
#colorway-radio label span {
    font-family: 'DM Sans', sans-serif !important;
    font-size: 12px !important;
    font-weight: 500 !important;
    color: #1A1410 !important;
    padding: 8px 10px 9px !important;
    display: block !important;
}
/* Motif Radio */
#motif-radio > div {
    display: grid !important;
    grid-template-columns: 1fr 1fr 1fr !important;
    gap: 8px !important;
}
#motif-radio label {
    background: #EDE3D1 !important;
    border: 2px solid transparent !important;
    border-radius: 4px !important;
    padding: 14px 8px 10px !important;
    cursor: pointer !important;
    text-align: center !important;
    transition: all 0.15s ease !important;
    display: flex !important;
    flex-direction: column !important;
    align-items: center !important;
    gap: 8px !important;
    box-shadow: 2px 2px 0 rgba(26,20,16,0.1) !important;
    margin: 0 !important;
}
#motif-radio label:hover {
    box-shadow: 3px 3px 0 #1A1410 !important;
    transform: translate(-1px,-1px) !important;
}
#motif-radio label:has(input:checked),
#motif-radio label.selected {
    border-color: #1E4D6B !important;
    background: #C84B31 !important;
    box-shadow: 3px 3px 0 #1E4D6B !important;
}
#motif-radio input[type="radio"] {
    display: none !important;
}
#motif-radio label::before {
    display: block !important;
    width: 48px !important;
    height: 48px !important;
    border-radius: 50% !important;
    background: rgba(245,237,224,0.7) !important;
    font-size: 22px !important;
    color: #3A2E25 !important;
    text-align: center !important;
    line-height: 48px !important;
    margin: 0 auto !important;
}
#motif-radio label:has(input:checked)::before,
#motif-radio label.selected::before {
    background: rgba(255,255,255,0.25) !important;
    color: #F5EDE0 !important;
}
#motif-radio label:nth-child(1)::before { content: '◎' !important; }
#motif-radio label:nth-child(2)::before { content: '◇' !important; }
#motif-radio label:nth-child(3)::before { content: '▣' !important; }
#motif-radio label:nth-child(4)::before { content: '✦' !important; }
#motif-radio label:nth-child(5)::before { content: '≡' !important; }
#motif-radio label:nth-child(6)::before { content: '✚' !important; }
#motif-radio label span {
    font-family: 'DM Sans', sans-serif !important;
    font-size: 8px !important;
    letter-spacing: 1.5px !important;
    text-transform: uppercase !important;
    color: #3A2E25 !important;
    font-weight: 600 !important;
    display: block !important;
}
#motif-radio label:has(input:checked) span,
#motif-radio label.selected span {
    color: #F5EDE0 !important;
}
/* Style Radio */
#style-radio > div {
    display: flex !important;
    flex-wrap: wrap !important;
    gap: 8px !important;
}
#style-radio label {
    padding: 7px 14px !important;
    border-radius: 3px !important;
    border: 1.5px solid rgba(26,20,16,0.25) !important;
    cursor: pointer !important;
    font-size: 13px !important;
    color: #3A2E25 !important;
    background: #EDE3D1 !important;
    margin: 0 !important;
    transition: all 0.15s ease !important;
    font-family: 'DM Sans', sans-serif !important;
}
#style-radio label:hover {
    border-color: #1A1410 !important;
    box-shadow: 2px 2px 0 #1A1410 !important;
    transform: translate(-1px,-1px) !important;
}
#style-radio label:has(input:checked),
#style-radio label.selected {
    border-color: #C84B31 !important;
    color: #C84B31 !important;
    font-weight: 500 !important;
    box-shadow: 2px 2px 0 #C84B31 !important;
}
#style-radio input[type="radio"] {
    display: none !important;
}
/* Num Radio */
#num-radio > div {
    display: flex !important;
    gap: 8px !important;
}
#num-radio label {
    width: 40px !important;
    height: 40px !important;
    border-radius: 4px !important;
    border: 1.5px solid rgba(26,20,16,0.25) !important;
    cursor: pointer !important;
    font-size: 16px !important;
    color: #3A2E25 !important;
    background: #EDE3D1 !important;
    margin: 0 !important;
    padding: 0 !important;
    display: flex !important;
    align-items: center !important;
    justify-content: center !important;
    transition: all 0.15s ease !important;
    font-family: 'DM Sans', sans-serif !important;
}
#num-radio label:hover {
    border-color: #1A1410 !important;
    box-shadow: 2px 2px 0 #1A1410 !important;
    transform: translate(-1px,-1px) !important;
}
#num-radio label:has(input:checked),
#num-radio label.selected {
    border-color: #C84B31 !important;
    color: #C84B31 !important;
    font-weight: 600 !important;
    box-shadow: 2px 2px 0 #C84B31 !important;
}
#num-radio input[type="radio"] {
    display: none !important;
}
/* Prompt + Button */
#prompt-display textarea {
    background: transparent !important;
    border: none !important;
    border-top: 1px solid rgba(26,20,16,0.15) !important;
    border-radius: 0 !important;
    font-family: 'DM Sans', sans-serif !important;
    font-size: 11px !important;
    color: #C84B31 !important;
    padding: 10px 0 0 !important;
    resize: none !important;
}
#prompt-display label {
    display: none !important;
}
#prompt-display .wrap {
    border: none !important;
    background: transparent !important;
    box-shadow: none !important;
}
#gen-btn button {
    background: #1A1410 !important;
    color: #F5EDE0 !important;
    border: none !important;
    border-radius: 4px !important;
    font-family: 'DM Sans', sans-serif !important;
    font-size: 11px !important;
    font-weight: 600 !important;
    letter-spacing: 2px !important;
    text-transform: uppercase !important;
    padding: 14px 24px !important;
    width: 100% !important;
    transition: all 0.15s ease !important;
    box-shadow: 3px 3px 0 #C84B31 !important;
}
#gen-btn button:hover {
    background: #C84B31 !important;
    box-shadow: 3px 3px 0 #1A1410 !important;
    transform: translate(-1px,-1px) !important;
}
#gen-btn button:active {
    transform: translate(0,0) !important;
    box-shadow: 1px 1px 0 #1A1410 !important;
}
/* Output */
.output-card {
    background: #EDE3D1;
    border: 1.5px solid rgba(26,20,16,0.2);
    border-radius: 4px;
    padding: 16px;
    margin-bottom: 12px;
}
#output-row-1, #output-row-2 {
    gap: 8px !important;
    margin-bottom: 8px !important;
}
#out1 img, #out2 img, #out3 img {
    border-radius: 4px !important;
    border: 1.5px solid #1A1410 !important;
    max-height: 280px !important;
    object-fit: contain !important;
    width: 100% !important;
}
#out3 {
    max-width: 50% !important;
}
#out1 .empty, #out2 .empty, #out3 .empty {
    display: none !important;
}
#out1 .upload-container, #out2 .upload-container, #out3 .upload-container {
    display: none !important;
}
#out1 > .svelte-1ipelgc > svg,
#out2 > .svelte-1ipelgc > svg,
#out3 > .svelte-1ipelgc > svg {
    display: none !important;
}
#out1 label, #out2 label, #out3 label {
    display: none !important;
}
/* Hide expand and share, keep only download */
#out1 button[aria-label="Expand"],
#out2 button[aria-label="Expand"],
#out3 button[aria-label="Expand"],
#out1 button[aria-label="Share"],
#out2 button[aria-label="Share"],
#out3 button[aria-label="Share"] {
    display: none !important;
}
/* Download button styling */
#out1 button[aria-label="Download"],
#out2 button[aria-label="Download"],
#out3 button[aria-label="Download"] {
    background: #1A1410 !important;
    color: #F5EDE0 !important;
    border-radius: 3px !important;
    opacity: 1 !important;
    padding: 6px 10px !important;
}
#out1 button[aria-label="Download"]:hover,
#out2 button[aria-label="Download"]:hover,
#out3 button[aria-label="Download"]:hover {
    background: #C84B31 !important;
}
.ref-grid {
    display: flex;
    gap: 12px;
}
.ref-item {
    flex: 1;
}
.ref-item img {
    width: 100%;
    border-radius: 4px;
    border: 1.5px solid rgba(26,20,16,0.2);
}
.ref-caption {
    font-family: 'DM Sans', sans-serif;
    font-size: 9px;
    text-transform: uppercase;
    letter-spacing: 1px;
    color: #3A2E25;
    margin-top: 5px;
}
#app-footer {
    padding: 16px 40px;
    display: flex;
    justify-content: space-between;
    font-family: 'DM Sans', sans-serif;
    font-size: 9px;
    letter-spacing: 1.5px;
    text-transform: uppercase;
    color: #3A2E25;
    border-top: 1px solid rgba(26,20,16,0.15);
    position: relative;
    z-index: 1;
}
@keyframes fadeDown {
    from { opacity: 0; transform: translateY(-10px); }
    to   { opacity: 1; transform: translateY(0); }
}
"""

KENTE = '<div id="kente-stripe"></div>'

HEADER = (
    '<div id="app-header">'
    '<div class="eyebrow">Generative African Textiles</div>'
    '<h1>Design African wax textile prints with <em>AI.</em></h1>'
    '<p class="subtitle">Pick your color, shape, and style —<br>your pattern generates in minutes.</p>'
    '</div>'
)

REF_HTML = (
    '<div class="output-card">'
    '<div class="section-heading" style="margin-bottom:12px">'
    '<span class="section-num"></span><span class="section-lbl">Reference Patterns</span>'
    '<div class="section-line"></div></div>'
    '<div class="ref-grid">'
    '<div class="ref-item"><img src="https://i.imgur.com/m7a9sZk.png"/><div class="ref-caption">Blue and White Wavy</div></div>'
    '<div class="ref-item"><img src="https://i.imgur.com/3GJarE9.png"/><div class="ref-caption">Olive Tribal Ovals</div></div>'
    '<div class="ref-item"><img src="https://i.imgur.com/xa6QOnC.png"/><div class="ref-caption">Red and Gold Swirling</div></div>'
    '</div></div>'
)

FOOTER = '<div id="app-footer"><span>Pace AI Lab</span><span>African Wax Patterns</span></div>'

SCROLL_JS = """
<script>
(function() {
    const waitForBtn = setInterval(() => {
        const btn = document.querySelector('#gen-btn button');
        if (!btn) return;
        clearInterval(waitForBtn);
        btn.addEventListener('click', () => {
            setTimeout(() => {
                const out = document.querySelector('#output-section');
                if (out) out.scrollIntoView({ behavior: 'smooth', block: 'start' });
            }, 800);
        });
    }, 300);
})();
</script>
"""

with gr.Blocks(title="African Wax Pattern Generator") as demo:
    gr.HTML(KENTE)
    gr.HTML(HEADER)

    with gr.Row(equal_height=False):
        with gr.Column(scale=0, min_width=320):
            gr.HTML('<div class="section-heading"><span class="section-num"></span><span class="section-lbl">Color</span><div class="section-line"></div></div>')
            colorway_radio = gr.Radio(
                choices=["Indigo Dusk", "Terracotta", "Saharan Ochre", "Forest Kente", "Royal Plum", "Ebony Bone"],
                value="Indigo Dusk",
                show_label=False,
                elem_id="colorway-radio"
            )

            gr.HTML('<div class="section-heading"><span class="section-num"></span><span class="section-lbl">Shape</span><div class="section-line"></div></div>')
            motif_radio = gr.Radio(
                choices=["Circles", "Diamonds", "Squares", "Stars", "Comb", "Adinkra"],
                value="Circles",
                show_label=False,
                elem_id="motif-radio"
            )

            gr.HTML('<div class="section-heading"><span class="section-num"></span><span class="section-lbl">Style</span><div class="section-line"></div></div>')
            style_radio = gr.Radio(
                choices=["Traditional", "Modern", "Abstract", "Geometric", "Floral"],
                value="Traditional",
                show_label=False,
                elem_id="style-radio"
            )

            gr.HTML('<div class="section-heading"><span class="section-num"></span><span class="section-lbl">Variations</span><div class="section-line"></div></div>')
            num_radio = gr.Radio(
                choices=["1", "2", "3"],
                value="1",
                show_label=False,
                elem_id="num-radio"
            )

            prompt_display = gr.Textbox(
                show_label=False,
                interactive=False,
                elem_id="prompt-display",
                lines=2,
                value="African wax print textile, traditional style, indigo blue and amber colorway, circular concentric motifs"
            )

            generate_btn = gr.Button("Generate Pattern", elem_id="gen-btn", variant="primary")

            gr.HTML(
                '<div style="font-family:DM Sans,sans-serif;font-size:9px;color:#3A2E25;text-align:center;margin-top:8px;letter-spacing:1px;text-transform:uppercase">1-5 min per image</div>'
            )

        with gr.Column(scale=1):
            with gr.Group():
                gr.HTML(
                    '<div class="output-card" id="output-section">'
                    '<div class="section-heading" style="margin-bottom:12px">'
                    '<span class="section-num"></span>'
                    '<span class="section-lbl">Output — hover image to download ↓</span>'
                    '<div class="section-line"></div></div>'
                )
                with gr.Row(elem_id="output-row-1"):
                    out1 = gr.Image(
                        show_label=False, type="pil", elem_id="out1",
                        container=False
                    )
                    out2 = gr.Image(
                        show_label=False, type="pil", elem_id="out2",
                        container=False
                    )
                with gr.Row(elem_id="output-row-2"):
                    out3 = gr.Image(
                        show_label=False, type="pil", elem_id="out3",
                        container=False
                    )
                gr.HTML('</div>')
                gr.HTML(REF_HTML)

    gr.HTML(SCROLL_JS)
    gr.HTML(FOOTER)

    def update_prompt(c, s, t):
        return build_prompt(c, s, t)

    for radio in [colorway_radio, motif_radio, style_radio]:
        radio.change(
            fn=update_prompt,
            inputs=[colorway_radio, motif_radio, style_radio],
            outputs=prompt_display
        )

    generate_btn.click(
        fn=generate_image,
        inputs=[colorway_radio, motif_radio, style_radio, num_radio],
        outputs=[out1, out2, out3, prompt_display],
    )


if __name__ == "__main__":
    demo.launch(css=CSS)
else:
    demo.launch(css=CSS)