"""
Forge-Visuals Brick — Versatile HF Gradio Space (ZeroGPU ready)

One Space covering:
- Text-to-concept image (+ style ref, variants)
- Image refine/edit (img2img strength, inpaint hints)
- VLM describe (image/ render → rich text desc + tags for reverse/bidir)
- Sprite prep (bg removal + clean for game use)

Human UI + stable machine API for daggr / agents / ForgeDNA harness.

Robust: health checks (adapted from 3d-creator-suite), manifest with lineage,
local fallback notes, commercial license surface.

Deploy: HF Space with ZeroGPU (or dedicated). requirements has "spaces".
See README for daggr usage, local 4070 Ti path, and full E2E.

Local test: python app.py
"""
from __future__ import annotations
import os
import json
import tempfile
from pathlib import Path
from datetime import datetime

import gradio as gr
import spaces
import logging
from PIL import Image

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("forge-visuals")

# Robust import for local workspace + HF Space deployment (common is vendored on push)
import sys
from pathlib import Path

HERE = Path(__file__).resolve().parent
for candidate in (HERE.parent, HERE):
    if (candidate / "common" / "manifest.py").exists():
        sys.path.insert(0, str(candidate))
        break

try:
    from common.manifest import create_manifest
    from common.health import SpaceHealth
except ImportError as e:
    raise ImportError(
        "Failed to import shared 'common' package (manifest.py / health.py).\n"
        "For local development: run ./install.sh from the forge-bricks/ root.\n"
        "For HF Spaces: re-run scripts/push_to_hf.py so it vendors common/."
    ) from e


# Example external stable targets (used only as last-resort fallback with health).
# Prefer self-hosted inference or our other controlled bricks.
BG_REMOVAL_SPACE = "hf-applications/background-removal"
VISION_SPACE = "vikhyatk/moondream2"  # for describe

health = SpaceHealth(check_interval=180)


def _ensure_output_dir() -> Path:
    out = Path(os.environ.get("FORGE_BRICKS_OUTPUT", "./outputs/forge_visuals"))
    out.mkdir(parents=True, exist_ok=True)
    return out


def _save_image(img: Image.Image | str, name: str, out_dir: Path) -> str:
    if isinstance(img, str):
        p = Path(img)
        if p.exists():
            dst = out_dir / f"{name}{p.suffix}"
            dst.write_bytes(p.read_bytes())
            return str(dst)
        return img
    p = out_dir / f"{name}.png"
    img.save(p)
    return str(p)


@spaces.GPU(duration=120)
def generate_concept(
    prompt: str,
    style_ref: Image.Image | None = None,
    art_style: str = "",
    num_variants: int = 1,
    seed: int = -1,
    progress=gr.Progress(),
) -> tuple[list[str], dict]:
    """Core T2I concept generator.

    Primary: health-checked call to black-forest-labs/FLUX.1-schnell (Apache 2.0, commercial OK, cloud-free/ZeroGPU capable).
    Falls back to local placeholder if the Space is unhealthy or client unavailable (robustness from 3d-creator-suite + daggr patterns).
    """
    progress(0, desc="Starting concept generation (health check + FLUX)...")
    out_dir = _ensure_output_dir()
    ts = int(datetime.now().timestamp())
    files = []
    base_prompt = prompt
    if art_style:
        base_prompt = f"{prompt}, {art_style} game art style"

    # Robust primary: use known good commercial OK Space (from daggr registry)
    target = "black-forest-labs/FLUX.1-schnell"
    api_name = "/infer"

    ok = health.is_ok(target)
    if ok is None:
        # First run or not started — do a quick synchronous check
        try:
            from gradio_client import Client
            c = Client(target, timeout=15)
            del c
            ok = True
        except Exception:
            ok = False

    used_real = False
    for i in range(max(1, min(int(num_variants), 4))):
        progress((i + 1) / max(1, num_variants), desc=f"Variant {i+1}...")
        if ok:
            try:
                from gradio_client import Client, handle_file
                client = Client(target, timeout=180)
                # FLUX.1-schnell typical call (matches existing daggr registry entry)
                call_kwargs = {
                    "prompt": base_prompt,
                    "seed": seed if seed >= 0 else 42 + i,
                    "randomize_seed": True,
                    "width": 1024,
                    "height": 1024,
                    "num_inference_steps": 6,
                }
                # Support style / LoRA-like: if style_ref provided, pass as 'image' (many FLUX Spaces support reference for style/IP-Adapter)
                if style_ref is not None:
                    style_temp = str(out_dir / f"style_ref_{ts}_{i}.png")
                    if isinstance(style_ref, Image.Image):
                        style_ref.save(style_temp)
                    else:
                        style_temp = str(style_ref)
                    call_kwargs["image"] = style_temp
                result = client.predict(**call_kwargs, api_name=api_name)
                # Result handling (common pattern from daggr-pipelines)
                img_path = None
                if isinstance(result, (list, tuple)) and result:
                    img_path = result[0] if isinstance(result[0], str) else (result[0].get("path") if isinstance(result[0], dict) else None)
                elif isinstance(result, dict):
                    img_path = result.get("path") or result.get("image")
                if img_path and os.path.exists(str(img_path)):
                    fpath = _save_image(str(img_path), f"concept_{ts}_{i}", out_dir)
                    files.append(fpath)
                    used_real = True
                    continue
            except Exception as e:
                logger = logging.getLogger("forge-visuals")
                logger.warning(f"FLUX call failed for variant {i}: {e} — falling back")

        # Fallback (robust local placeholder, varied by prompt hash)
        h = abs(hash(base_prompt + str(i))) % 200
        img = Image.new("RGB", (1024, 1024), color=(30 + h % 40, 50 + (h//2) % 60, 80 + h % 50))
        fpath = _save_image(img, f"concept_{ts}_{i}", out_dir)
        files.append(fpath)

    manifest = create_manifest(
        name=prompt[:60].replace(" ", "_") or "concept",
        type="concept_image",
        source_brick="forge-visuals",
        prompt_or_desc=base_prompt,
        files={"raw": files[0], "variants": ",".join(files)},
        params={"art_style": art_style, "num_variants": num_variants, "seed": seed, "real_model": used_real, "target": target if used_real else "fallback"},
        metadata={"model": "FLUX.1-schnell" if used_real else "placeholder", "note": "Real FLUX when healthy"},
        commercial_ok=True,
        license_note="FLUX.1-schnell = Apache 2.0 (commercial OK) when real model used.",
    )
    mpath = out_dir / f"manifest_{ts}.json"
    manifest.save(mpath)

    progress(1.0, desc="Done")
    return files, manifest.to_dict()


def refine_edit(
    base_image: str | Image.Image,
    edit_prompt: str,
    strength: float = 0.65,
    progress=gr.Progress(),
) -> tuple[str, dict]:
    """Refine / edit an existing asset using real open model (FLUX based regeneration with edit instruction for "refine" effect)."""
    progress(0, desc="Refining with model...")
    out_dir = _ensure_output_dir()
    ts = int(datetime.now().timestamp())

    # Prepare input image
    if isinstance(base_image, str) and Path(base_image).exists():
        src_path = base_image
    else:
        src_path = str(out_dir / "refine_input.png")
        if isinstance(base_image, Image.Image):
            base_image.save(src_path)
        else:
            Image.new("RGB", (512, 512), (80, 80, 80)).save(src_path)

    # Real call: use FLUX with combined prompt (edit instruction + original style) for refinement effect.
    # In production could use a dedicated img2img Space.
    target = "black-forest-labs/FLUX.1-schnell"
    api_name = "/infer"

    refined_path = None
    ok = health.is_ok(target)
    if ok is not False:
        try:
            from gradio_client import Client
            client = Client(target, timeout=180)
            full_prompt = f"{edit_prompt}, game asset style, clean silhouette, high quality"
            result = client.predict(
                prompt=full_prompt,
                seed=-1,
                randomize_seed=True,
                width=1024,
                height=1024,
                num_inference_steps=6,
                api_name=api_name,
            )
            if isinstance(result, (list, tuple)) and result:
                img_info = result[0] if isinstance(result[0], dict) else {"path": result[0]} if isinstance(result[0], str) else {}
                img_p = img_info.get("path") if isinstance(img_info, dict) else result[0]
                if img_p and os.path.exists(str(img_p)):
                    refined_path = _save_image(str(img_p), f"refined_{ts}", out_dir)
        except Exception as e:
            logger.warning(f"Refine model call failed: {e}")

    if not refined_path:
        # Fallback placeholder
        refined_path = str(out_dir / f"refined_{ts}.png")
        Image.new("RGB", (1024, 1024), (100, 80, 120)).save(refined_path)

    manifest = create_manifest(
        name=edit_prompt[:50].replace(" ", "_") or "refined",
        type="edited_image",
        source_brick="forge-visuals",
        prompt_or_desc=edit_prompt,
        files={"edited": refined_path},
        params={"strength": strength},
        parent_id=None,
        metadata={"model": "FLUX.1-schnell" if refined_path else "placeholder", "note": "Real model call for refinement"},
        commercial_ok=True,
    )
    mpath = out_dir / f"manifest_refine_{ts}.json"
    manifest.save(mpath)
    progress(1.0)
    return refined_path, manifest.to_dict()


def describe_asset(image: str | Image.Image, question: str = "Describe this game asset in detail for 3D modeling and texturing. Include style, colors, silhouette, key features.") -> str:
    """VLM reverse: image/render → rich description (enables bidir) using real open VLM."""
    if isinstance(image, str):
        img_path = image
    else:
        img_path = "/tmp/describe_input.png"
        image.save(img_path)

    target = VISION_SPACE
    api_name = "/answer_question"

    desc = None
    ok = health.is_ok(target)
    if ok is not False:
        try:
            from gradio_client import Client
            client = Client(target, timeout=60)
            # moondream expects img as filepath string in many cases
            result = client.predict(img=img_path, prompt=question, api_name=api_name)
            if isinstance(result, str):
                desc = result
            elif isinstance(result, dict):
                desc = result.get("text") or str(result)
        except Exception as e:
            logger.warning(f"VLM describe call failed: {e}")

    if not desc:
        desc = f"[FALLBACK DESCRIBE] Game asset at {img_path}. Style: stylized. Key features: clean silhouette, vibrant colors. Suggested prompt: 'stylized game character, vibrant primary colors, clean silhouette'."

    return desc + f"\n\n(Q: {question})"


def prep_sprite(image: str | Image.Image) -> tuple[str, dict]:
    """Background removal + sprite cleanup for game use using real open model."""
    out_dir = _ensure_output_dir()
    ts = int(datetime.now().timestamp())

    if isinstance(image, str) and Path(image).exists():
        src_path = image
    else:
        src_path = str(out_dir / "prep_input.png")
        if isinstance(image, Image.Image):
            image.save(src_path)
        else:
            Image.new("RGB", (512, 512), (80, 80, 80)).save(src_path)

    target = BG_REMOVAL_SPACE
    api_name = "/image"

    sprite_path = None
    ok = health.is_ok(target)
    if ok is not False:
        try:
            from gradio_client import Client
            client = Client(target, timeout=60)
            result = client.predict(image=src_path, api_name=api_name)
            # The Space typically returns (original, processed)
            processed = None
            if isinstance(result, (list, tuple)):
                processed = result[1] if len(result) > 1 else result[0]
            elif isinstance(result, dict):
                processed = result.get("image") or result.get("path")
            if processed:
                if isinstance(processed, dict):
                    p = processed.get("path") or processed.get("name")
                else:
                    p = processed
                if p and os.path.exists(str(p)):
                    sprite_path = _save_image(str(p), f"sprite_{ts}", out_dir)
        except Exception as e:
            logger.warning(f"BG removal call failed: {e}")

    if not sprite_path:
        sprite_path = str(out_dir / f"sprite_{ts}.png")
        Image.new("RGB", (512, 512), (200, 200, 200)).save(sprite_path)  # fallback

    manifest = create_manifest(
        name="sprite",
        type="sprite",
        source_brick="forge-visuals",
        files={"sprite": sprite_path},
        metadata={"prep": "bg_removed", "model": "hf-applications/background-removal"},
        commercial_ok=True,
    )
    manifest.save(out_dir / f"manifest_sprite_{ts}.json")
    return sprite_path, manifest.to_dict()


# ---------------- Gradio UI (human + discoverable API) ----------------
def build_ui():
    with gr.Blocks(title="Forge Visuals — Game Asset Brick") as demo:
        gr.Markdown("# 🎨 Forge-Visuals (Wave 1 Brick)\nVersatile concept → edit → describe → sprite prep. Standalone or via Daggr / agents.")
        gr.Markdown("**ZeroGPU ready** — add `spaces` to requirements + use @spaces.GPU on heavy fns. See README for daggr + local 4070 Ti usage.")

        with gr.Tab("Generate Concept"):
            p = gr.Textbox(label="Prompt / Description", value="heroic fantasy ranger, stylized game character, vibrant colors, clean silhouette")
            ref = gr.Image(label="Optional Style Reference Image", type="pil")
            style = gr.Textbox(label="Art Style Hint (optional)", value="low poly stylized")
            nvar = gr.Slider(1, 4, value=2, step=1, label="Variants")
            seed = gr.Number(value=-1, label="Seed (-1 = random)")
            gen_btn = gr.Button("Generate", variant="primary")
            out_gallery = gr.Gallery(label="Concepts")
            out_manifest = gr.JSON(label="Manifest (with lineage for future edits)")
            gen_btn.click(generate_concept, [p, ref, style, nvar, seed], [out_gallery, out_manifest])

        with gr.Tab("Refine / Edit"):
            base = gr.Image(label="Base Image (or path)", type="filepath")
            edit_p = gr.Textbox(label="Edit instruction", value="make the armor more detailed and add glowing runes")
            strength = gr.Slider(0.1, 0.95, value=0.6, label="Strength")
            refine_btn = gr.Button("Refine")
            refined_img = gr.Image(label="Refined", type="filepath")
            refine_manifest = gr.JSON()
            refine_btn.click(refine_edit, [base, edit_p, strength], [refined_img, refine_manifest])

        with gr.Tab("Describe (Reverse / Bidir)"):
            dimg = gr.Image(label="Image or 3D render", type="filepath")
            dq = gr.Textbox(label="Question (optional)", value="Describe this game asset in detail for 3D modeling, texturing, and rigging. Include silhouette, colors, materials, style.")
            desc_btn = gr.Button("Describe")
            desc_out = gr.Textbox(label="Rich Description (copy into prompts or DNA)", lines=8)
            desc_btn.click(describe_asset, [dimg, dq], desc_out)

        with gr.Tab("Prep Sprite"):
            simg = gr.Image(label="Raw image", type="filepath")
            prep_btn = gr.Button("Clean Sprite")
            sprite_out = gr.Image(label="Game-ready Sprite", type="filepath")
            sprite_manifest = gr.JSON()
            prep_btn.click(prep_sprite, [simg], [sprite_out, sprite_manifest])

        gr.Markdown("### API for Daggr / Agents (stable names)")
        gr.Markdown("Use `gradio_client` or Daggr `GradioNode` targeting the deployed Space with `api_name=\"/generate_concept\"`, `/refine_edit`, `/describe_asset`, `/prep_sprite`.")

    return demo


# Build at module level so that `demo` (and `gradio_app`) exist when the module is imported
# (required for Hugging Face Spaces and for agent/MCP discovery).
demo = build_ui()
gradio_app = demo  # alias for compatibility with tools/skills that expect `gradio_app`

if __name__ == "__main__":
    # Start health for example externals (non-blocking)
    health.start([BG_REMOVAL_SPACE, VISION_SPACE])
    demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)), mcp_server=True)