"""
app.py — SILHOUETTE Studio
==========================
Monochrome neural pose-guided image generation and latent synthesis platform.

Tabs:
  1. Pose Studio (ControlNet OpenPose + LCM-LoRA)
  2. Fast Synthesis (Text-to-Image LCM-LoRA)

Run with:
  python app.py
"""

import logging
import random
import os

# ──────────────────────────────────────────────────────────────────────────────
# Gradio 4.x / gradio_client compatibility patch:
# In Gradio 4.x, components with dictionary properties (such as Sketchpad)
# generate JSON schemas containing additionalProperties as a boolean (True/False).
# gradio_client's schema parser crashes with TypeError or APIInfoParseError on booleans.
# Patch all schema parser functions and get_api_info at startup to be completely resilient.
# ──────────────────────────────────────────────────────────────────────────────
try:
    import gradio_client.utils as _gu

    _orig_get_type = getattr(_gu, "get_type", None)
    if _orig_get_type:
        def _safe_get_type(schema):
            if not isinstance(schema, dict):
                return "unknown"
            return _orig_get_type(schema)
        _gu.get_type = _safe_get_type

    _orig_sub_schema = getattr(_gu, "_json_schema_to_python_type", None)
    if _orig_sub_schema:
        def _safe_sub_schema(schema, defs=None):
            if isinstance(schema, bool) or not isinstance(schema, dict):
                return "Any"
            try:
                return _orig_sub_schema(schema, defs)
            except Exception:
                return "Any"
        _gu._json_schema_to_python_type = _safe_sub_schema

    _orig_main_schema = getattr(_gu, "json_schema_to_python_type", None)
    if _orig_main_schema:
        def _safe_main_schema(schema):
            if isinstance(schema, bool) or not isinstance(schema, dict):
                return "Any"
            try:
                return _orig_main_schema(schema)
            except Exception:
                return "Any"
        _gu.json_schema_to_python_type = _safe_main_schema
except Exception:
    pass

import gradio as gr

# Guard Blocks.get_api_info against schema parse crashes
try:
    _orig_get_api_info = gr.Blocks.get_api_info
    def _safe_get_api_info(self):
        try:
            return _orig_get_api_info(self)
        except Exception:
            return {"named_endpoints": {}, "unnamed_endpoints": {}}
    gr.Blocks.get_api_info = _safe_get_api_info
except Exception:
    pass
from PIL import Image

# ──────────────────────────────────────────────────────────────────────────────
# Hugging Face ZeroGPU Support:
# When running on HF Spaces ZeroGPU, functions performing neural inference
# must be decorated with @spaces.GPU so an A100 GPU is allocated on demand.
# Fallback to no-op when running locally without the spaces package.
# ──────────────────────────────────────────────────────────────────────────────
try:
    import spaces
    gpu_decorator = spaces.GPU
except Exception:
    def gpu_decorator(func=None, **kwargs):
        if func is not None and callable(func):
            return func
        def _wrapper(f):
            return f
        return _wrapper

from pipeline.processor import extract_pose_from_image, process_sketch, prepare_image
from pipeline.generate import generate_with_pose, generate_text_to_image
from utils.canvas import draw_skeleton_overlay
from utils.image_utils import save_output, pil_to_numpy

# ──────────────────────────────────────────────────────────────────────────────
# Logging
# ──────────────────────────────────────────────────────────────────────────────

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)-7s | %(name)s -> %(message)s",
    datefmt="%H:%M:%S",
)
logger = logging.getLogger(__name__)

# ──────────────────────────────────────────────────────────────────────────────
# Constants & Defaults
# ──────────────────────────────────────────────────────────────────────────────

APP_NAME       = "SILHOUETTE"
APP_SUBTITLE   = "Neural Pose Synthesis & Latent Generation Engine"
DEFAULT_PROMPT = "a subject standing in dynamic posture, studio lighting, highly detailed, photorealistic, 8k"
DEFAULT_NEG    = (
    "low quality, blurry, distorted, deformed, bad anatomy, "
    "watermark, signature, text, extra limbs, amputated"
)

# ──────────────────────────────────────────────────────────────────────────────
# Callbacks: Pose Studio
# ──────────────────────────────────────────────────────────────────────────────

@gpu_decorator
def on_extract_pose(reference_image: Image.Image):
    """Extract OpenPose skeleton from an uploaded reference photo."""
    if reference_image is None:
        return None, None, "[Warning] Upload a reference photo first."
    try:
        pose_img = extract_pose_from_image(reference_image)
        return pose_img, pil_to_numpy(pose_img), "[Success] Pose skeleton extracted successfully."
    except Exception as e:
        logger.exception("Pose extraction failed.")
        return None, None, f"[Error] {e}"


@gpu_decorator(duration=60)
def on_generate_pose(
    reference_image: Image.Image,
    sketch,
    extracted_pose_img: Image.Image | None,
    prompt: str,
    negative_prompt: str,
    num_steps: int,
    guidance_scale: float,
    controlnet_scale: float,
    seed_mode: str,
    manual_seed: int,
):
    """
    Main generation callback for Pose Studio.

    Pose Priority:
      1. Extracted pose from reference photo (if extracted)
      2. User's sketch drawing on canvas
      3. Error - nothing to condition on
    """
    seed = manual_seed if seed_mode == "Fixed" else random.randint(0, 2**32 - 1)
    pose_image = None

    if extracted_pose_img is not None:
        pose_image = prepare_image(extracted_pose_img)
    elif sketch is not None:
        pose_image = process_sketch(sketch)

    if pose_image is None:
        return (
            gr.update(value=None),
            gr.update(value=None),
            gr.update(value="[Warning] Please extract a pose from a photo or draw a skeleton on the canvas."),
        )

    if not prompt.strip():
        return (
            gr.update(value=None),
            gr.update(value=pil_to_numpy(pose_image)),
            gr.update(value="[Warning] Please enter a prompt."),
        )

    try:
        output = generate_with_pose(
            prompt=prompt,
            pose_image=pose_image,
            negative_prompt=negative_prompt,
            num_inference_steps=int(num_steps),
            guidance_scale=float(guidance_scale),
            controlnet_conditioning_scale=float(controlnet_scale),
            seed=int(seed),
        )
        save_output(output, prefix="silhouette_pose")
        return (
            gr.update(value=pil_to_numpy(output)),
            gr.update(value=pil_to_numpy(pose_image)),
            gr.update(value=f"[Success] Generation complete. Seed: {seed}"),
        )
    except Exception as e:
        logger.exception("Pose-guided generation failed.")
        return (
            gr.update(value=None),
            gr.update(value=pil_to_numpy(pose_image) if pose_image is not None else None),
            gr.update(value=f"[Error] {e}"),
        )


# ──────────────────────────────────────────────────────────────────────────────
# Callbacks: Fast Synthesis (T2I)
# ──────────────────────────────────────────────────────────────────────────────

@gpu_decorator(duration=30)
def on_generate_t2i(
    prompt: str,
    negative_prompt: str,
    num_steps: int,
    guidance_scale: float,
    seed_mode: str,
    manual_seed: int,
):
    """Fast text-to-image generation callback."""
    seed = manual_seed if seed_mode == "Fixed" else random.randint(0, 2**32 - 1)

    if not prompt.strip():
        return gr.update(value=None), gr.update(value="[Warning] Please enter a prompt.")

    try:
        output = generate_text_to_image(
            prompt=prompt,
            negative_prompt=negative_prompt,
            num_inference_steps=int(num_steps),
            guidance_scale=float(guidance_scale),
            seed=int(seed),
        )
        save_output(output, prefix="silhouette_t2i")
        return gr.update(value=pil_to_numpy(output)), gr.update(value=f"[Success] Generation complete. Seed: {seed}")
    except Exception as e:
        logger.exception("Text-to-image generation failed.")
        return gr.update(value=None), gr.update(value=f"[Error] {e}")


def _seed_interactivity(mode: str):
    return gr.update(visible=(mode == "Fixed"))


# ──────────────────────────────────────────────────────────────────────────────
# Custom Minimalist Black & White Theme + CSS
# ──────────────────────────────────────────────────────────────────────────────

_CUSTOM_CSS = """
/* Global Container & Typography */
body, .gradio-container {
    background-color: #09090b !important;
    color: #f4f4f5 !important;
    font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, Helvetica, Arial, sans-serif !important;
}

/* Header & Titles */
#app-title {
    text-align: center;
    font-size: 2.2rem;
    font-weight: 800;
    letter-spacing: 0.25em;
    color: #ffffff;
    margin-top: 10px;
    margin-bottom: 2px;
    text-transform: uppercase;
}

#app-subtitle {
    text-align: center;
    font-size: 0.85rem;
    font-weight: 400;
    letter-spacing: 0.12em;
    color: #71717a;
    margin-bottom: 28px;
    text-transform: uppercase;
}

/* Tab Navigation */
.tabs {
    border-bottom: 1px solid #27272a !important;
    margin-bottom: 20px !important;
}

.tab-nav button {
    font-size: 0.9rem !important;
    font-weight: 600 !important;
    letter-spacing: 0.05em !important;
    color: #a1a1aa !important;
    border: none !important;
    background: transparent !important;
    padding: 10px 20px !important;
    text-transform: uppercase !important;
    transition: all 0.2s ease !important;
}

.tab-nav button.selected {
    color: #ffffff !important;
    border-bottom: 2px solid #ffffff !important;
}

/* Panels & Cards */
.gr-panel, .gr-box, .gr-form, .gr-block {
    background-color: #121215 !important;
    border: 1px solid #27272a !important;
    border-radius: 6px !important;
}

/* Inputs & Textareas */
input[type="text"], textarea, .gr-input, .gr-dropdown {
    background-color: #18181b !important;
    border: 1px solid #27272a !important;
    color: #ffffff !important;
    border-radius: 4px !important;
}

input[type="text"]:focus, textarea:focus {
    border-color: #ffffff !important;
    box-shadow: 0 0 0 1px #ffffff !important;
}

/* Labels */
label span, .gr-label {
    color: #a1a1aa !important;
    font-size: 0.78rem !important;
    font-weight: 600 !important;
    letter-spacing: 0.06em !important;
    text-transform: uppercase !important;
}

/* Buttons */
.btn-primary, button.primary-btn, button.primary {
    background: #ffffff !important;
    color: #000000 !important;
    font-weight: 700 !important;
    font-size: 0.85rem !important;
    letter-spacing: 0.08em !important;
    text-transform: uppercase !important;
    border: 1px solid #ffffff !important;
    border-radius: 4px !important;
    padding: 12px 24px !important;
    transition: background-color 0.15s ease, opacity 0.15s ease !important;
}

.btn-primary:hover, button.primary-btn:hover, button.primary:hover {
    background: #e4e4e7 !important;
    color: #000000 !important;
    opacity: 0.95 !important;
}

.btn-secondary, button.secondary-btn, button.secondary {
    background: transparent !important;
    color: #ffffff !important;
    font-weight: 600 !important;
    font-size: 0.82rem !important;
    letter-spacing: 0.06em !important;
    text-transform: uppercase !important;
    border: 1px solid #3f3f46 !important;
    border-radius: 4px !important;
    padding: 8px 16px !important;
    transition: all 0.15s ease !important;
}

.btn-secondary:hover, button.secondary-btn:hover, button.secondary:hover {
    background: #27272a !important;
    border-color: #52525b !important;
}

/* Status Box */
.status-box textarea {
    background-color: #09090b !important;
    border: 1px solid #27272a !important;
    color: #d4d4d8 !important;
    font-family: ui-monospace, SFMono-Regular, Menlo, Monaco, Consolas, monospace !important;
    font-size: 0.82rem !important;
}

/* Sliders */
input[type="range"] {
    accent-color: #ffffff !important;
}

/* Instructions Banner */
.instruction-card {
    background-color: #121215;
    border: 1px solid #27272a;
    border-radius: 6px;
    padding: 14px 18px;
    margin-bottom: 16px;
    font-size: 0.85rem;
    line-height: 1.6;
    color: #a1a1aa;
}

.instruction-card strong {
    color: #ffffff;
}

/* Footer */
.footer-text {
    text-align: center;
    font-size: 0.75rem;
    color: #52525b;
    letter-spacing: 0.08em;
    text-transform: uppercase;
    margin-top: 30px;
    padding-top: 15px;
    border-top: 1px solid #1f1f23;
}
"""

def build_app() -> gr.Blocks:
    theme = gr.themes.Monochrome(
        primary_hue="neutral",
        secondary_hue="neutral",
        neutral_hue="neutral",
        font=[gr.themes.GoogleFont("Inter"), "sans-serif"],
    ).set(
        body_background_fill="#09090b",
        block_background_fill="#121215",
        block_border_color="#27272a",
        input_background_fill="#18181b",
        input_border_color="#27272a",
        button_primary_background_fill="#ffffff",
        button_primary_text_color="#000000",
        button_secondary_background_fill="transparent",
        button_secondary_text_color="#ffffff",
        button_secondary_border_color="#3f3f46",
    )

    with gr.Blocks(
        title=APP_NAME,
        theme=theme,
        css=_CUSTOM_CSS,
    ) as demo:

        gr.Markdown(f"# {APP_NAME}", elem_id="app-title")
        gr.Markdown(APP_SUBTITLE, elem_id="app-subtitle")

        # ── Tab 1: Pose Studio ───────────────────────────────────────────────
        with gr.Tab("Pose Studio"):
            gr.HTML(
                """
                <div class="instruction-card">
                    <strong>Workflow:</strong>
                    1. Upload a reference photo and click <strong>Extract Pose</strong>, or sketch directly on the canvas.<br/>
                    2. Provide your descriptive prompt and generation parameters.<br/>
                    3. Click <strong>Generate</strong> to synthesize the image conditioned on the target posture.
                </div>
                """
            )

            with gr.Row():
                # Left Column: Pose Inputs
                with gr.Column(scale=1):
                    ref_image = gr.Image(
                        label="Reference Photo (Optional)",
                        type="pil",
                        height=270,
                    )
                    extract_btn = gr.Button("Extract Pose from Photo", variant="secondary")

                    gr.Markdown("---")

                    sketch = gr.Sketchpad(
                        label="Pose Sketchpad (Draw Skeleton — white strokes on black)",
                        canvas_size=(512, 512),
                        type="numpy",
                        height=300,
                    )

                # Right Column: Text Prompts & Parameters
                with gr.Column(scale=1):
                    prompt_p = gr.Textbox(
                        label="Prompt",
                        value=DEFAULT_PROMPT,
                        lines=3,
                        placeholder="Detailed description of the subject, background, lighting, and style...",
                    )
                    neg_prompt_p = gr.Textbox(
                        label="Negative Prompt",
                        value=DEFAULT_NEG,
                        lines=2,
                    )

                    with gr.Row():
                        steps_p = gr.Slider(
                            label="Steps",
                            minimum=2,
                            maximum=20,
                            step=1,
                            value=8,
                            info="4-8 recommended for LCM",
                        )
                        cfg_p = gr.Slider(
                            label="CFG Scale",
                            minimum=1.0,
                            maximum=3.0,
                            step=0.1,
                            value=1.5,
                            info="1.0-2.0 recommended for LCM",
                        )

                    cn_scale = gr.Slider(
                        label="ControlNet Conditioning Scale",
                        minimum=0.1,
                        maximum=1.5,
                        step=0.05,
                        value=0.85,
                        info="Higher values enforce stricter adherence to the pose",
                    )

                    with gr.Row():
                        seed_mode_p = gr.Dropdown(
                            choices=["Random", "Fixed"],
                            value="Random",
                            label="Seed Mode",
                            scale=1,
                        )
                        manual_seed_p = gr.Number(
                            label="Manual Seed",
                            value=42,
                            precision=0,
                            visible=False,
                            scale=2,
                        )

                    generate_btn_p = gr.Button("Generate", variant="primary", size="lg")
                    status_p = gr.Textbox(
                        label="Console / Status",
                        interactive=False,
                        elem_classes=["status-box"],
                    )

            # Output Row
            with gr.Row():
                extracted_pose_display = gr.Image(
                    label="Active Conditioning Pose",
                    type="numpy",
                    height=320,
                    interactive=False,
                )
                output_image_p = gr.Image(
                    label="Generated Output",
                    type="numpy",
                    height=320,
                    interactive=False,
                )

            # State for storing the extracted PIL pose image
            extracted_pose_state = gr.State(value=None)

            extract_btn.click(
                fn=on_extract_pose,
                inputs=[ref_image],
                outputs=[extracted_pose_state, extracted_pose_display, status_p],
            )

            seed_mode_p.change(
                fn=_seed_interactivity,
                inputs=[seed_mode_p],
                outputs=[manual_seed_p],
            )

            generate_btn_p.click(
                fn=on_generate_pose,
                inputs=[
                    ref_image,
                    sketch,
                    extracted_pose_state,
                    prompt_p,
                    neg_prompt_p,
                    steps_p,
                    cfg_p,
                    cn_scale,
                    seed_mode_p,
                    manual_seed_p,
                ],
                outputs=[output_image_p, extracted_pose_display, status_p],
            )

        # ── Tab 2: Fast Synthesis ────────────────────────────────────────────
        with gr.Tab("Fast Synthesis"):
            gr.HTML(
                """
                <div class="instruction-card">
                    <strong>Text-to-Image Mode:</strong>
                    Direct latent consistency generation from text prompts. Optimized for 4-step rapid synthesis without pose conditioning.
                </div>
                """
            )

            with gr.Row():
                with gr.Column(scale=1):
                    prompt_t = gr.Textbox(
                        label="Prompt",
                        value="a minimalist architectural interior, brutalist design, soft daylight, monochrome tone, 8k",
                        lines=4,
                        placeholder="Describe the image to synthesize...",
                    )
                    neg_prompt_t = gr.Textbox(
                        label="Negative Prompt",
                        value=DEFAULT_NEG,
                        lines=2,
                    )

                    with gr.Row():
                        steps_t = gr.Slider(
                            label="Steps",
                            minimum=2,
                            maximum=12,
                            step=1,
                            value=4,
                            info="4 steps is optimal for LCM",
                        )
                        cfg_t = gr.Slider(
                            label="CFG Scale",
                            minimum=1.0,
                            maximum=2.0,
                            step=0.1,
                            value=1.0,
                            info="Maintain at 1.0 for LCM",
                        )

                    with gr.Row():
                        seed_mode_t = gr.Dropdown(
                            choices=["Random", "Fixed"],
                            value="Random",
                            label="Seed Mode",
                            scale=1,
                        )
                        manual_seed_t = gr.Number(
                            label="Manual Seed",
                            value=42,
                            precision=0,
                            visible=False,
                            scale=2,
                        )

                    generate_btn_t = gr.Button("Generate", variant="primary", size="lg")
                    status_t = gr.Textbox(
                        label="Console / Status",
                        interactive=False,
                        elem_classes=["status-box"],
                    )

                with gr.Column(scale=1):
                    output_image_t = gr.Image(
                        label="Generated Output",
                        type="numpy",
                        height=512,
                        interactive=False,
                    )

            seed_mode_t.change(
                fn=_seed_interactivity,
                inputs=[seed_mode_t],
                outputs=[manual_seed_t],
            )

            generate_btn_t.click(
                fn=on_generate_t2i,
                inputs=[
                    prompt_t,
                    neg_prompt_t,
                    steps_t,
                    cfg_t,
                    seed_mode_t,
                    manual_seed_t,
                ],
                outputs=[output_image_t, status_t],
            )

        # ── Minimalist Footer ────────────────────────────────────────────────
        gr.HTML(
            f"""
            <div class="footer-text">
                {APP_NAME} &middot; Stable Diffusion 1.5 &middot; OpenPose ControlNet &middot; LCM-LoRA &middot; VAE ft-mse
            </div>
            """
        )

    return demo


# ──────────────────────────────────────────────────────────────────────────────
# Entry Point
# ──────────────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    logger.info("Starting SILHOUETTE Studio ...")
    app = build_app()
    # HF Spaces handles port binding and ingress automatically when launched simply.
    # We omit inbrowser, share, and custom host logic that breaks in container environments.
    app.launch()