app.py

"""HF Spaces-first Gradio app for Flux Seamless Texture LoRA.

- Geração via HF Inference (huggingface_hub.InferenceClient) usando `ModelHandler`
- UI moderna com controles avançados, presets, gallery e history
- MCP via `mcp_server=True` + endpoints expostos com `api_name`
"""

import os
import logging
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union

import gradio as gr
import spaces

from config.settings import MCP_ENABLED
from ui_theme import build_theme
from src.model_handler import ModelHandler
from src.presets import list_presets, get_preset_prompt, get_preset_params
from src.utils import validate_prompt, generate_seed
from src.image_processor import create_zip, OUTPUT_DIR

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

logger.info("BOOT: HF Spaces-first app (InferenceClient) carregando…")

model_handler = ModelHandler()

# Prompt base: o usuário não precisa lembrar de pedir "seamless/tileable".
# Mantemos em inglês por compatibilidade com a maioria dos modelos de imagem.
BASE_TEXTURE_INSTRUCTIONS = (
    "seamless, tileable, repeatable, repeating pattern, perfectly looping texture, "
    "no visible seams, no borders, no frame, no text, no watermark"
)

# Estado simples em memória (suficiente para Spaces; persistência real via OUTPUT_DIR)
gallery_state: List[Dict[str, Any]] = []
history_state: List[Dict[str, Any]] = []


def _extract_image(result: Any) -> Any:
    """Normaliza retornos: PIL, tuple/list/dict."""
    if isinstance(result, tuple):
        return result[0]
    if isinstance(result, list) and result:
        return result[0]
    if isinstance(result, dict):
        return result.get("image") or (result.get("images", [None])[0] if result.get("images") else None)
    return result


def _augment_prompt_for_seamless(prompt: str) -> str:
    """
    Acrescenta instruções de textura tileable/seamless automaticamente.
    Se o usuário já menciona seamless/tileable/repeatable, não duplica.
    """
    import re

    p = (prompt or "").strip()
    if not p:
        return p

    # Se já tem indicação de tile/seamless/repeat, não adiciona.
    if re.search(r"\b(seamless|tileable|tiling|repeatable|repeating|repeat)\b", p, flags=re.IGNORECASE):
        return p

    return f"{BASE_TEXTURE_INSTRUCTIONS}, {p}"

def _merge_negative_prompt(preset_neg: str, user_neg: str) -> str:
    """Combina negative prompt do preset com o do usuário (sem sobrescrever)."""
    preset_neg = (preset_neg or "").strip()
    user_neg = (user_neg or "").strip()
    if not preset_neg:
        return user_neg
    if not user_neg:
        return preset_neg
    # Dedupe simples por substring (case-insensitive)
    if preset_neg.lower() in user_neg.lower():
        return user_neg
    if user_neg.lower() in preset_neg.lower():
        return preset_neg
    return f"{preset_neg}, {user_neg}"


@spaces.GPU(duration=300)
def generate_texture(
    prompt: str,
    negative_prompt: str,
    preset: str,
    guidance_scale: float,
    num_inference_steps: int,
    seed: float,
    width: int,
    height: int,
    cfg_scale: float,
    lora_strength: float,
    progress: gr.Progress = gr.Progress(),
) -> Tuple[Any, str, Dict[str, Any]]:
    """UI handler: gera imagem e atualiza gallery/history."""
    try:
        progress(0.0, desc="Validando prompt…")
        # Valida o prompt do usuário (antes de anexar o base prompt).
        is_valid, error = validate_prompt(prompt, max_length=1000)
        if not is_valid:
            return None, f"Erro: {error}", {}

        if preset and preset != "None":
            preset_prompt = get_preset_prompt(preset)
            preset_params = get_preset_params(preset)

            if preset_prompt:
                prompt = f"{preset_prompt}, {prompt}" if prompt else preset_prompt

            if preset_params:
                guidance_scale = float(preset_params.get("guidance_scale", guidance_scale))
                num_inference_steps = int(preset_params.get("num_inference_steps", num_inference_steps))
                width = int(preset_params.get("width", width))
                height = int(preset_params.get("height", height))
                if "negative_prompt" in preset_params:
                    negative_prompt = _merge_negative_prompt(
                        str(preset_params.get("negative_prompt") or ""),
                        negative_prompt,
                    )

        # Acrescenta instruções seamless/tileable automaticamente
        prompt = _augment_prompt_for_seamless(prompt)

        # Revalida após augment (evita estourar limite)
        is_valid, error = validate_prompt(prompt, max_length=1200)
        if not is_valid:
            # fallback: corta com segurança
            prompt = prompt[:1200]

        seed_int = int(seed) if seed is not None and seed >= 0 else generate_seed()

        progress(0.15, desc="Chamando API de inferência…")
        image, metadata = model_handler.generate(
            prompt=prompt,
            negative_prompt=negative_prompt,
            guidance_scale=float(guidance_scale),
            num_inference_steps=int(num_inference_steps),
            seed=seed_int,
            width=int(width),
            height=int(height),
            cfg_scale=float(cfg_scale),
            lora_strength=float(lora_strength),
        )

        entry = {
            "timestamp": datetime.now().timestamp(),
            "prompt": prompt,
            "negative_prompt": negative_prompt,
            "params": {
                "guidance_scale": guidance_scale,
                "num_inference_steps": num_inference_steps,
                "seed": seed_int,
                "width": width,
                "height": height,
                "cfg_scale": cfg_scale,
                "lora_strength": lora_strength,
                "preset": preset,
            },
            "image_path": metadata.get("image_path"),
            "image": image,
        }

        gallery_state.append(entry)
        history_state.append(entry)

        progress(1.0, desc="Concluído")
        return image, "Pronto — imagem na pré-visualização e na Galeria/Histórico.", metadata
    except Exception as e:
        logger.error("Falha ao gerar imagem", exc_info=True)
        return None, f"Erro: {str(e)}", {}


def get_gallery_images() -> List[Tuple[Any, str]]:
    return [
        (entry["image"], f"{entry['prompt'][:50]}...")
        for entry in reversed(gallery_state[-24:])
        if entry.get("image") is not None
    ]


def get_history_table() -> List[List[str]]:
    rows: List[List[str]] = []
    for entry in reversed(history_state[-100:]):
        ts = datetime.fromtimestamp(entry["timestamp"]).strftime("%Y-%m-%d %H:%M:%S")
        p = entry["prompt"]
        p = (p[:60] + "...") if len(p) > 60 else p
        rows.append([ts, p, str(entry["params"].get("seed", "")), entry["params"].get("preset", "None")])
    return rows


def download_gallery_zip() -> Optional[Path]:
    try:
        image_paths = [
            Path(entry["image_path"])
            for entry in gallery_state
            if entry.get("image_path") and Path(entry["image_path"]).exists()
        ]
        if not image_paths:
            return None
        zip_path = OUTPUT_DIR / f"gallery_{int(datetime.now().timestamp())}.zip"
        create_zip(image_paths, zip_path)
        return zip_path
    except Exception as e:
        logger.error(f"Erro criando ZIP: {e}", exc_info=True)
        return None


# Endpoints expostos via Gradio API / MCP (api_name)
def _sanitize_for_json(obj: Any) -> Any:
    """Garante que todas as chaves de dicionários sejam strings (ORJSON requirement)."""
    if isinstance(obj, dict):
        return {str(k): _sanitize_for_json(v) for k, v in obj.items()}
    if isinstance(obj, (list, tuple)):
        return [_sanitize_for_json(v) for v in obj]
    # Converte tipos não-serializáveis para string
    if not isinstance(obj, (str, int, float, bool, type(None))):
        return str(obj)
    return obj


def api_generate_texture(
    prompt: str,
    negative_prompt: str = "",
    preset: str = "None",
    guidance_scale: float = 7.5,
    num_inference_steps: int = 50,
    seed: int = -1,
    width: int = 1024,
    height: int = 1024,
    cfg_scale: float = 7.5,
    lora_strength: float = 1.0,
) -> Tuple[Any, Dict[str, Any]]:
    """
    Gera uma textura seamless.
    
    Retorna: (imagem_gerada, metadata_json)
    - A imagem é servida automaticamente pelo Gradio para download.
    - O metadata contém informações sobre a geração.
    """
    image, status, metadata = generate_texture(
        prompt=prompt,
        negative_prompt=negative_prompt,
        preset=preset,
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
        seed=seed,
        width=width,
        height=height,
        cfg_scale=cfg_scale,
        lora_strength=lora_strength,
    )
    if image is None:
        return None, _sanitize_for_json({"success": False, "error": status})
    return image, _sanitize_for_json({"success": True, "metadata": metadata})


def api_get_presets() -> Dict[str, Any]:
    from src.presets import TEXTURE_PRESETS

    return {"presets": list_presets(), "details": TEXTURE_PRESETS}


def api_get_history(limit: int = 10) -> Dict[str, Any]:
    recent = history_state[-limit:] if len(history_state) > limit else history_state
    # Evitar enviar a imagem inteira via API
    safe = []
    for e in recent:
        safe.append(
            {
                "timestamp": e["timestamp"],
                "prompt": e["prompt"],
                "negative_prompt": e.get("negative_prompt", ""),
                "params": _sanitize_for_json(e.get("params", {})),
                "image_path": e.get("image_path"),
            }
        )
    return _sanitize_for_json({"total": len(history_state), "entries": safe})


SEAMLESS_CSS = """
.seamless-hero {
  padding: 1.35rem 1.5rem 1.25rem;
  border-radius: 14px;
  margin-bottom: 0.5rem;
  background: linear-gradient(115deg, rgba(245, 158, 11, 0.14) 0%, rgba(148, 163, 184, 0.12) 45%, rgba(241, 245, 249, 0.65) 100%);
  border: 1px solid rgba(148, 163, 184, 0.45);
  box-shadow: 0 12px 40px rgba(15, 23, 42, 0.06);
}
.seamless-hero h1 {
  margin: 0 0 0.35rem 0;
  font-size: 1.65rem;
  letter-spacing: -0.02em;
  line-height: 1.2;
}
.seamless-hero p {
  margin: 0;
  opacity: 0.88;
  font-size: 0.98rem;
  line-height: 1.45;
}
.seamless-badge {
  display: inline-block;
  font-size: 0.72rem;
  font-weight: 600;
  text-transform: uppercase;
  letter-spacing: 0.12em;
  color: var(--color-accent);
  margin-bottom: 0.5rem;
}
footer.seamless-foot {
  margin-top: 1.25rem;
  padding-top: 0.75rem;
  font-size: 0.85rem;
  opacity: 0.75;
  border-top: 1px solid rgba(148, 163, 184, 0.35);
}
"""


def apply_preset(
    preset: str,
    prompt: str,
    negative_prompt: str,
    guidance_scale: float,
    num_steps: int,
    width: int,
    height: int,
):
    if not preset or preset == "None":
        return prompt, negative_prompt, guidance_scale, num_steps, width, height
    preset_prompt = get_preset_prompt(preset)
    preset_params = get_preset_params(preset)
    if preset_prompt:
        prompt = f"{preset_prompt}, {prompt}" if prompt else preset_prompt
    if preset_params:
        guidance_scale = float(preset_params.get("guidance_scale", guidance_scale))
        num_steps = int(preset_params.get("num_inference_steps", num_steps))
        width = int(preset_params.get("width", width))
        height = int(preset_params.get("height", height))
        if "negative_prompt" in preset_params:
            negative_prompt = _merge_negative_prompt(
                str(preset_params.get("negative_prompt") or ""),
                negative_prompt,
            )
    return prompt, negative_prompt, guidance_scale, num_steps, width, height


with gr.Blocks(title="Seamless Texture Studio") as demo:
    gr.HTML(
        """
        <div class="seamless-hero">
          <span class="seamless-badge">HF Inference · Flux LoRA</span>
          <h1>Seamless Texture Studio</h1>
          <p>Texturas repetíveis em alta resolução. Descreva o material; o app reforça <em>seamless/tileable</em> automaticamente. API e MCP em <code>/gradio_api/mcp/</code>.</p>
        </div>
        """
    )

    with gr.Tabs():
        with gr.Tab("Gerar", id="tab-generate"):
            with gr.Row(equal_height=True):
                with gr.Column(scale=5):
                    prompt_input = gr.Textbox(
                        label="Prompt",
                        placeholder="ex.: madeira clara com veios finos, desgaste suave",
                        lines=4,
                        info="Inglês costuma funcionar melhor com a maioria dos modelos.",
                    )
                    negative_prompt_input = gr.Textbox(
                        label="Prompt negativo",
                        placeholder="O que evitar (opcional)",
                        lines=2,
                    )
                    with gr.Row():
                        preset_dropdown = gr.Dropdown(
                            choices=["None"] + list_presets(),
                            value="None",
                            label="Preset",
                            info="Aplica prompt e parâmetros sugeridos para o material.",
                        )
                        seed_input = gr.Number(
                            label="Seed",
                            value=-1,
                            precision=0,
                            info="−1 = aleatório",
                        )
                    with gr.Accordion("Parâmetros avançados", open=False):
                        guidance_scale_slider = gr.Slider(
                            1.0, 20.0, value=7.5, step=0.5, label="Guidance scale"
                        )
                        num_steps_slider = gr.Slider(
                            10, 100, value=50, step=5, label="Passos de inferência"
                        )
                        with gr.Row():
                            width_slider = gr.Slider(256, 2048, value=1024, step=64, label="Largura")
                            height_slider = gr.Slider(256, 2048, value=1024, step=64, label="Altura")
                        cfg_scale_slider = gr.Slider(1.0, 20.0, value=7.5, step=0.5, label="CFG scale")
                        lora_strength_slider = gr.Slider(
                            0.0, 2.0, value=1.0, step=0.1, label="Força do LoRA"
                        )

                    with gr.Row():
                        apply_preset_btn = gr.Button("Aplicar preset", variant="secondary", size="lg")
                        generate_btn = gr.Button("Gerar textura", variant="primary", size="lg")

                    apply_preset_btn.click(
                        fn=apply_preset,
                        inputs=[
                            preset_dropdown,
                            prompt_input,
                            negative_prompt_input,
                            guidance_scale_slider,
                            num_steps_slider,
                            width_slider,
                            height_slider,
                        ],
                        outputs=[
                            prompt_input,
                            negative_prompt_input,
                            guidance_scale_slider,
                            num_steps_slider,
                            width_slider,
                            height_slider,
                        ],
                        api_name=False,
                    )

                with gr.Column(scale=5):
                    image_output = gr.Image(
                        label="Pré-visualização",
                        type="pil",
                        height=420,
                        show_label=True,
                    )
                    status_output = gr.Textbox(label="Estado", interactive=False, lines=2)
                    with gr.Accordion("Metadados técnicos", open=False):
                        metadata_output = gr.JSON(label="Metadata")

            generate_btn.click(
                fn=generate_texture,
                inputs=[
                    prompt_input,
                    negative_prompt_input,
                    preset_dropdown,
                    guidance_scale_slider,
                    num_steps_slider,
                    seed_input,
                    width_slider,
                    height_slider,
                    cfg_scale_slider,
                    lora_strength_slider,
                ],
                outputs=[image_output, status_output, metadata_output],
                api_name=False,
                show_progress="full",
            )

        with gr.Tab("Galeria"):
            gr.Markdown("Últimas gerações desta sessão (memória do processo).")
            with gr.Row():
                gallery_refresh_btn = gr.Button("Atualizar galeria", variant="secondary")
                download_zip_btn = gr.Button("Baixar tudo (ZIP)", variant="primary")
            gallery_display = gr.Gallery(
                label="Texturas geradas",
                columns=4,
                rows=2,
                height="auto",
                object_fit="contain",
                show_label=True,
            )
            zip_download = gr.File(label="Arquivo ZIP", visible=False)

            gallery_refresh_btn.click(fn=get_gallery_images, outputs=gallery_display, api_name=False)
            download_zip_btn.click(fn=download_gallery_zip, outputs=zip_download, api_name=False).then(
                fn=lambda x: gr.update(visible=True) if x else gr.update(),
                inputs=zip_download,
                outputs=zip_download,
                api_name=False,
            )

        with gr.Tab("Histórico"):
            history_table = gr.Dataframe(
                label="Últimas execuções",
                headers=["Data/hora", "Prompt", "Seed", "Preset"],
                interactive=False,
                wrap=True,
            )
            history_refresh_btn = gr.Button("Atualizar tabela")
            history_refresh_btn.click(fn=get_history_table, outputs=history_table, api_name=False)

        with gr.Tab("MCP / API", visible=MCP_ENABLED):
            gr.Markdown(
                """
### Model Context Protocol (MCP)

| | |
| --- | --- |
| **Endpoint** | `/gradio_api/mcp/` |
| **Ferramentas** | `generate_texture`, `get_presets`, `get_history` |

Integração útil para agentes e pipelines que chamam o Space remotamente.
                """
            )

        gr.HTML(
            """
            <footer class="seamless-foot">
              Geração via Hugging Face Inference · Modelo configurável por variável <code>MODEL_ID</code>
            </footer>
            """
        )

        # Hidden: endpoints para MCP/Gradio API
        with gr.Accordion("🔌 API/MCP (hidden)", open=False, visible=False):
            api_prompt = gr.Textbox(label="prompt")
            api_negative_prompt = gr.Textbox(label="negative_prompt", value="")
            api_preset = gr.Dropdown(choices=["None"] + list_presets(), value="None", label="preset")
            api_guidance = gr.Slider(1.0, 20.0, value=7.5, label="guidance_scale")
            api_steps = gr.Slider(10, 100, value=50, step=5, label="num_inference_steps")
            api_seed = gr.Number(value=-1, precision=0, label="seed")
            api_width = gr.Slider(256, 2048, value=1024, step=64, label="width")
            api_height = gr.Slider(256, 2048, value=1024, step=64, label="height")
            api_cfg = gr.Slider(1.0, 20.0, value=7.5, step=0.5, label="cfg_scale")
            api_lora = gr.Slider(0.0, 2.0, value=1.0, step=0.1, label="lora_strength")

            # Output: imagem (servida automaticamente pelo Gradio) + metadata JSON
            api_out_image = gr.Image(label="generated_image", type="pil")
            api_out = gr.JSON(label="result")
            gr.Button("generate_texture", visible=False).click(
                fn=api_generate_texture,
                inputs=[
                    api_prompt,
                    api_negative_prompt,
                    api_preset,
                    api_guidance,
                    api_steps,
                    api_seed,
                    api_width,
                    api_height,
                    api_cfg,
                    api_lora,
                ],
                outputs=[api_out_image, api_out],
                api_name="generate_texture",
            )

            api_presets_out = gr.JSON(label="presets")
            gr.Button("get_presets", visible=False).click(
                fn=api_get_presets,
                inputs=[],
                outputs=[api_presets_out],
                api_name="get_presets",
            )

            api_history_limit = gr.Number(value=10, precision=0, label="limit")
            api_history_out = gr.JSON(label="history")
            gr.Button("get_history", visible=False).click(
                fn=api_get_history,
                inputs=[api_history_limit],
                outputs=[api_history_out],
                api_name="get_history",
            )


if __name__ == "__main__":
    # Spaces-friendly: queue + desabilitar SSR experimental
    # Gradio 6: theme/css em launch(), não no construtor de Blocks
    demo.queue(default_concurrency_limit=int(os.getenv("GRADIO_CONCURRENCY_LIMIT", "1")))
    demo.launch(
        theme=build_theme(),
        css=SEAMLESS_CSS,
        mcp_server=bool(MCP_ENABLED),
        ssr_mode=False,
        share=False,
    )