from __future__ import annotations

import os
import secrets

import gradio as gr

from model_runtime import predict_mask
from multiview_runtime import generate_studio_photos, save_studio_outputs
from studio_pipeline import isolate_product, prepare_input


def create_studio_set(uploaded_image):
    try:
        image = prepare_input(uploaded_image)
        mask = predict_mask(image)
        product = isolate_product(image, mask)
        results = generate_studio_photos(product, secrets.randbelow(2**31 - 1))
        files, archive = save_studio_outputs(results)
        return [(photo, title) for title, photo in results], files, archive, "Ready"
    except Exception as exc:
        raise gr.Error(f"Generation failed: {exc}") from exc


CSS = """
.gradio-container { max-width: 1120px !important; margin: 0 auto !important; }
#hero { text-align: center; padding: 34px 8px 14px; }
#hero h1 { font-size: clamp(2.5rem, 6vw, 4.8rem); letter-spacing: -.06em; margin-bottom: 4px; }
#hero p { color: #64748b; font-size: 1.15rem; }
.card { border: 1px solid #e2e8f0 !important; border-radius: 22px !important; box-shadow: 0 14px 40px rgba(15,23,42,.06); }
#generate { min-height: 58px; font-weight: 800; }
#status { text-align: center; color: #64748b; font-size: .9rem; }
"""

with gr.Blocks(css=CSS, title="Studio10") as demo:
    gr.Markdown("# Studio10\nFrom a casual photo to studio shots.", elem_id="hero")
    source = gr.Image(
        label="Upload a product photo",
        type="pil",
        image_mode=None,
        sources=["upload", "clipboard"],
        height=500,
        elem_classes="card",
    )
    generate = gr.Button("Generate 10 studio shots", variant="primary", elem_id="generate")
    status = gr.Markdown("", elem_id="status")
    gallery = gr.Gallery(
        show_label=False,
        columns=2,
        rows=5,
        height=1350,
        object_fit="contain",
        elem_classes="card",
    )
    with gr.Row():
        output_files = gr.File(label="Images", file_count="multiple")
        zip_file = gr.File(label="ZIP")

    generate.click(
        fn=create_studio_set,
        inputs=source,
        outputs=[gallery, output_files, zip_file, status],
        api_name="generate_10_studio_shots",
        concurrency_limit=1,
        scroll_to_output=True,
    )

demo.queue(default_concurrency_limit=1, max_size=4)

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=int(os.getenv("PORT", "7860")),
        show_error=True,
    )