src/apps/p5_memory_quilt/app.py

from __future__ import annotations

from dataclasses import asdict
from datetime import date, datetime, timezone
from pathlib import Path
import argparse
import os

from app_kit.tracing import write_trace_artifact

from .demo_pack_loader import load_default_demo_pack
from .pipeline import QuiltState, add_memory_to_state, empty_state, render_artifacts, seed_state_from_pack, state_summary
from .prompt_rewriter import rewrite_memory
from .quilt import render_tile, save_image

ROOT = Path(__file__).resolve().parents[3]
THEME_CSS_PATH = ROOT / "assets" / "theme.css"
DEFAULT_PORT = 7860


def _export_root(kind: str = "exports") -> Path:
    return ROOT / "artifacts" / kind / date.today().isoformat()


def _runtime_root(kind: str = "runtime") -> Path:
    return ROOT / "artifacts" / "verification" / date.today().isoformat() / kind


def _sample_gallery(pack) -> list[str]:
    return [str(photo) for photo in pack.photos]


def _sample_memories(pack) -> list[str]:
    cards: list[str] = []
    for memory in pack.memories:
        location = memory.location.strip() if getattr(memory, "location", "") else "Neighborhood memory"
        cards.append(f"- **{location}** — {memory.text}")
    return cards


def _state_to_details(state: QuiltState) -> list[dict[str, object]]:
    details: list[dict[str, object]] = []
    for index, card in enumerate(state.cards, start=1):
        row = asdict(card)
        row["index"] = index
        details.append(row)
    return details


def _inference_metadata(card, *, artifact_path: str | None = None, trace_path: str | None = None) -> dict[str, object]:
    meta = dict(getattr(card, "inference_meta", {}) or {})
    meta.update(
        {
            "caption": getattr(card, "caption", ""),
            "story": getattr(card, "story", ""),
            "flux_prompt": getattr(card, "flux_prompt", ""),
            "style": getattr(card, "style", ""),
            "selected_model_id": getattr(card, "selected_model_id", ""),
            "prompt_source": getattr(card, "prompt_source", ""),
            "checkpoint_path": getattr(card, "checkpoint_path", ""),
            "artifact_path": artifact_path,
            "trace_path": trace_path,
        }
    )
    return meta


def _trace_payload(kind: str, inputs: dict[str, object], parsed_outputs: dict[str, object], card, *, pack_id: str = "", pack_path: str = "") -> dict[str, object]:
    meta = dict(getattr(card, "inference_meta", {}) or {})
    return {
        "kind": kind,
        "project": "p5",
        "pack_id": pack_id,
        "pack_path": pack_path,
        "timestamp": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
        "inputs": inputs,
        "parsed_outputs": parsed_outputs,
        "model_name": getattr(card, "selected_model_id", ""),
        "model_id": str(meta.get("model_id", getattr(card, "selected_model_id", ""))),
        "adapter_name": str(meta.get("adapter_name", getattr(card, "prompt_source", ""))),
        "checkpoint_path": str(meta.get("checkpoint_path", getattr(card, "checkpoint_path", ""))),
        "checkpoint_source": str(meta.get("checkpoint_source", "unknown")),
        "generation_stats": meta.get("generation_stats", {}),
    }


def _write_runtime_trace(kind: str, inputs: dict[str, object], parsed_outputs: dict[str, object], card, *, pack_id: str = "", pack_path: str = "") -> Path:
    return write_trace_artifact(_runtime_root(kind), _trace_payload(kind, inputs, parsed_outputs, card, pack_id=pack_id, pack_path=pack_path))


def _gradio_available():
    try:
        import gradio as gr
    except Exception as exc:  # pragma: no cover - only hit when runtime deps are missing
        raise RuntimeError(
            "Gradio is required to launch the app. Run scripts/bootstrap_venv.sh or install requirements.txt first."
        ) from exc
    return gr


def _initial_outputs():
    pack = load_default_demo_pack()
    first = pack.memories[0] if pack.memories else None
    status = (
        "Ready to stitch. Load example memories or type your own memory below. "
        "Generation now requires a mounted local checkpoint; if none is present, the UI will show a clear error."
    )
    sample_text = "\n\n".join(_sample_memories(pack))
    return (
        empty_state(pack.style),
        first.text if first else "",
        first.location if first else "",
        pack.style,
        _sample_gallery(pack),
        None,
        None,
        [],
        {
            "model_ready": False,
            "message": status,
            "checkpoint_hint": "Set P5_MEMORY_QUILT_PRIMARY_MODEL_PATH or P5_MEMORY_QUILT_FALLBACK_MODEL_PATH to a local checkpoint directory.",
        },
        status,
        sample_text,
    )


def create_app():
    gr = _gradio_available()

    with gr.Blocks(title="FLUX Memory Quilt", css_paths=THEME_CSS_PATH) as demo:
        with gr.Row():
            gr.Markdown(
                "<div style='text-align: center; margin-bottom: 20px;'>"
                "<h1 style='font-size: 2.5em; margin-bottom: 0.2em;'>🧵 Memory Quilt Builder</h1>"
                "<p style='font-size: 1.1em; color: var(--p5-ink-soft);'>Turn tiny neighborhood memories into stitched quilt panels with a local model checkpoint.</p>"
                "</div>\n"
                "<div class='status-bar' style='text-align: center; max-width: 800px; margin: 0 auto; font-weight: 500;'>If no local checkpoint is mounted, generation actions fail clearly instead of falling back to synthetic output.</div>"
            )
        state = gr.State(empty_state())
        with gr.Tabs():
            with gr.Tab("🎨 Quilt Builder"):
                with gr.Row():
                    with gr.Column(scale=1, elem_classes=["quilt-card"]):
                        gr.Markdown(
                            "#### Describe a memory\n"
                            "Use 1–3 sentences and an optional location tag. The sample gallery below shows the bundled demo pack.")
                        memory_input = gr.Textbox(
                            label="Tiny neighborhood memory",
                            lines=3,
                            placeholder="Every Friday the tamale cart parked outside the blue house.",
                            info="Required: a short memory to transform into a quilt panel.",
                        )
                        location_input = gr.Textbox(label="Location tag", placeholder="corner store", info="Optional: a place name or neighborhood tag.")
                        style_input = gr.Dropdown(
                            label="Quilt style",
                            choices=["Fabric Quilt", "Watercolor Map", "Polaroid Collage", "Linocut Print"],
                            value="Fabric Quilt",
                            info="Choose the visual language for the generated panel.",
                        )
                        photo_input = gr.Image(type="filepath", label="Photo reference (optional)", elem_classes=["upload-area"])
                        with gr.Row():
                            load_button = gr.Button("Load example memories", variant="secondary")
                            generate_button = gr.Button("Generate quilt", variant="primary")
                        gr.Markdown(
                            "<div class='quilt-note'>The example memories are synthetic and bundled locally. Generation requires a mounted local checkpoint.</div>"
                        )
                    with gr.Column(scale=1, elem_classes=["quilt-card"]):
                        quilt_image = gr.Image(label="Generated quilt", type="filepath", elem_classes=["quilt-preview"])
                        download_file = gr.File(label="Download quilt PNG")
                        status_output = gr.Markdown("Ready to stitch.", elem_classes=["quilt-note"])
                        inference_output = gr.JSON(label="Inference log")
                        gr.Markdown("#### 🖼️ Generated tiles")
                        details_output = gr.JSON(label="Tile details")
                        gr.Markdown("#### 📖 Example gallery")
                        sample_gallery = gr.Gallery(label="Example photos", columns=3, height=240)
                        sample_text = gr.Markdown()

                def load_pack():
                    pack = load_default_demo_pack()
                    first = pack.memories[0] if pack.memories else None
                    sample_text_value = "\n\n".join(_sample_memories(pack))
                    try:
                        seeded_state = seed_state_from_pack(pack, count=min(6, len(pack.memories)))
                        render = render_artifacts(seeded_state, _export_root("exports") / "sample_pack", stem="sample_pack")
                        last_card = seeded_state.cards[-1]
                        trace_path = _write_runtime_trace(
                            "sample_pack",
                            {
                                "memory_count": len(pack.memories),
                                "style": pack.style,
                                "action": "load_example_memories",
                            },
                            {
                                "quilt_path": str(render.quilt_path),
                                "tile_path": str(render.tile_path),
                                "log_path": str(render.log_path),
                                "card_count": len(seeded_state.cards),
                            },
                            last_card,
                            pack_id=pack.pack_id,
                            pack_path=str(pack.path),
                        )
                        status = f"✅ Loaded {len(pack.memories)} example memories — {state_summary(seeded_state)}"
                        return (
                            seeded_state,
                            first.text if first else "",
                            first.location if first else "",
                            pack.style,
                            _sample_gallery(pack),
                            str(render.quilt_path),
                            str(render.quilt_path),
                            _state_to_details(seeded_state),
                            _inference_metadata(last_card, artifact_path=str(render.log_path), trace_path=str(trace_path)),
                            status,
                            sample_text_value,
                        )
                    except Exception as exc:
                        status = f"❌ {exc}"
                        return (
                            empty_state(pack.style),
                            first.text if first else "",
                            first.location if first else "",
                            pack.style,
                            _sample_gallery(pack),
                            None,
                            None,
                            [],
                            {
                                "model_ready": False,
                                "error": str(exc),
                                "checkpoint_hint": "Set P5_MEMORY_QUILT_PRIMARY_MODEL_PATH or P5_MEMORY_QUILT_FALLBACK_MODEL_PATH to a local checkpoint directory.",
                            },
                            status,
                            sample_text_value,
                        )

                def generate_from_inputs(current_state: QuiltState, memory_text: str, location_tag: str, style: str, photo: str | None):
                    try:
                        next_state, card = add_memory_to_state(current_state, memory_text, location_tag, style, photo_path=photo)
                        render = render_artifacts(next_state, _export_root("exports") / "sessions", stem="memory_quilt")
                        trace_path = _write_runtime_trace(
                            "quilt_generation",
                            {
                                "memory_text": memory_text,
                                "location_tag": location_tag,
                                "style": style,
                                "has_photo": bool(photo),
                            },
                            {
                                "quilt_path": str(render.quilt_path),
                                "tile_path": str(render.tile_path),
                                "log_path": str(render.log_path),
                                "card_count": len(next_state.cards),
                            },
                            card,
                            pack_id="",
                            pack_path="",
                        )
                        meta = _inference_metadata(card, artifact_path=str(render.log_path), trace_path=str(trace_path))
                        stats = meta.get("generation_stats") or {}
                        stats_text = ""
                        if isinstance(stats, dict) and stats:
                            tokens = stats.get("generated_tokens")
                            elapsed = stats.get("elapsed_ms")
                            stats_text = f" · {tokens} tokens · {elapsed} ms" if tokens is not None and elapsed is not None else ""
                        status = f"Added {card.caption} with {card.selected_model_id} via {card.prompt_source}{stats_text}. {state_summary(next_state)}"
                        return (
                            next_state,
                            str(render.quilt_path),
                            str(render.quilt_path),
                            _state_to_details(next_state),
                            meta,
                            status,
                        )
                    except Exception as exc:
                        return (
                            current_state,
                            None,
                            None,
                            [],
                            {
                                "model_ready": False,
                                "error": str(exc),
                                "selected_style": style,
                                "has_photo": bool(photo),
                            },
                            f"❌ {exc}",
                        )

                load_button.click(
                    load_pack,
                    outputs=[state, memory_input, location_input, style_input, sample_gallery, quilt_image, download_file, details_output, inference_output, status_output, sample_text],
                )
                generate_button.click(
                    generate_from_inputs,
                    inputs=[state, memory_input, location_input, style_input, photo_input],
                    outputs=[state, quilt_image, download_file, details_output, inference_output, status_output],
                )

            with gr.Tab("🖼️ Single Tile Playground"):
                with gr.Row():
                    with gr.Column(scale=1, elem_classes=["quilt-card"]):
                        gr.Markdown(
                            "#### Generate a single tile\n"
                            "Prompt the model directly for a one-off quilt patch. The result includes a log with model ID, adapter, and generation stats."
                        )
                        single_prompt = gr.Textbox(label="Free-form prompt", lines=3, placeholder="A winter bus stop, warm light, friends waiting together.")
                        single_location = gr.Textbox(label="Location tag", placeholder="bus stop")
                        single_style = gr.Dropdown(
                            label="Quilt style",
                            choices=["Fabric Quilt", "Watercolor Map", "Polaroid Collage", "Linocut Print"],
                            value="Fabric Quilt",
                        )
                        single_photo = gr.Image(type="filepath", label="Photo reference (optional)", elem_classes=["upload-area"])
                        single_button = gr.Button("Generate single tile", variant="primary")
                    with gr.Column(scale=1, elem_classes=["quilt-card"]):
                        single_image = gr.Image(label="Single tile", type="filepath", elem_classes=["quilt-preview"])
                        single_file = gr.File(label="Download tile PNG")
                        single_details = gr.JSON(label="Tile details")
                        single_inference = gr.JSON(label="Inference log")
                        single_status = gr.Markdown("Waiting for a prompt.", elem_classes=["quilt-note"])

                def generate_single_tile(prompt: str, location_tag: str, style: str, photo: str | None):
                    try:
                        card = rewrite_memory(prompt, location_tag, style, photo_path=photo)
                        tile_image = render_tile(card)
                        export_dir = _export_root("exports") / "single_tile"
                        export_dir.mkdir(parents=True, exist_ok=True)
                        tile_path = save_image(tile_image, export_dir / "single_tile.png")
                        trace_path = _write_runtime_trace(
                            "single_tile",
                            {
                                "prompt": prompt,
                                "location_tag": location_tag,
                                "style": style,
                                "has_photo": bool(photo),
                            },
                            {
                                "tile_path": str(tile_path),
                            },
                            card,
                            pack_id="",
                            pack_path="",
                        )
                        meta = _inference_metadata(card, artifact_path=str(tile_path), trace_path=str(trace_path))
                        stats = meta.get("generation_stats") or {}
                        stats_text = ""
                        if isinstance(stats, dict) and stats:
                            tokens = stats.get("generated_tokens")
                            elapsed = stats.get("elapsed_ms")
                            stats_text = f" · {tokens} tokens · {elapsed} ms" if tokens is not None and elapsed is not None else ""
                        return (
                            str(tile_path),
                            str(tile_path),
                            [asdict(card)],
                            meta,
                            f"Rendered {card.caption} with {card.selected_model_id} via {card.prompt_source}{stats_text}.",
                        )
                    except Exception as exc:
                        return (
                            None,
                            None,
                            [],
                            {
                                "model_ready": False,
                                "error": str(exc),
                                "selected_style": style,
                                "has_photo": bool(photo),
                            },
                            f"❌ {exc}",
                        )

                single_button.click(
                    generate_single_tile,
                    inputs=[single_prompt, single_location, single_style, single_photo],
                    outputs=[single_image, single_file, single_details, single_inference, single_status],
                )

            with gr.Tab("📖 How It Works"):
                gr.Markdown(
                    """
                    ### How to use the Memory Quilt

                    1. **Start with an example**: Click **Load example memories** to generate a starter quilt from the bundled demo pack.
                    2. **Write a memory**: Describe a tiny neighborhood moment and add a location tag if it helps.
                    3. **Choose a style**: Pick a quilt aesthetic that matches the mood of the memory.
                    4. **Generate locally**: The app now requires a mounted local checkpoint and fails clearly if it is missing.
                    5. **Inspect the logs**: Each render returns the model ID, adapter, and generation stats alongside the exported image paths.
                    """
                )
        demo.load(
            lambda: _initial_outputs(),
            outputs=[state, memory_input, location_input, style_input, sample_gallery, quilt_image, download_file, details_output, inference_output, status_output, sample_text],
        )
    return demo


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Launch the FLUX memory quilt Gradio app")
    parser.add_argument("--host", default=os.environ.get("HOST", "0.0.0.0"), help="Host interface for the Gradio server")
    parser.add_argument("--port", type=int, default=int(os.environ.get("PORT", str(DEFAULT_PORT))), help="Port for the Gradio server")
    parser.add_argument("--share", action="store_true", help="Enable Gradio share links")
    args = parser.parse_args(argv)
    demo = create_app()
    demo.queue()
    demo.launch(server_name=args.host, share=args.share)
    return 0