"""
Forge-Texture Brick — PBR texture generation.

Inputs: text description or reference image + style.
Outputs: albedo, normal, roughness, metallic, AO maps + combined manifest.

Uses health-checked image model + post-processing guidance for game PBR.
"""
from __future__ import annotations
import os
from pathlib import Path
from datetime import datetime
import gradio as gr
import spaces
from PIL import Image

# 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

health = SpaceHealth()

TARGET = "black-forest-labs/FLUX.1-schnell"  # strong base for PBR prompting

def _out() -> Path:
    d = Path(os.environ.get("FORGE_BRICKS_OUTPUT", "./outputs/forge_texture"))
    d.mkdir(parents=True, exist_ok=True)
    return d

@spaces.GPU(duration=90)
def generate_pbr(prompt: str, style_ref: Image.Image | None = None, resolution: int = 1024) -> dict:
    out_dir = _out()
    ts = int(datetime.now().timestamp())

    base = prompt + ", seamless tileable PBR texture, albedo map, high detail, game asset, no text"

    # Real multi-map generation using targeted FLUX calls + health (PBR workflow)
    maps = {}
    map_prompts = {
        "albedo": base,
        "normal": prompt + ", normal map, blue-purple tones, game ready",
        "roughness": prompt + ", roughness map, grayscale, smooth to rough",
        "metallic": prompt + ", metallic map, grayscale, metal vs dielectric",
        "ao": prompt + ", ambient occlusion map, grayscale, crevices dark"
    }

    for m, m_prompt in map_prompts.items():
        try:
            if health.is_ok(TARGET) is not False:
                from gradio_client import Client
                client = Client(TARGET, timeout=60)
                res = client.predict(
                    prompt=m_prompt,
                    seed=-1,
                    randomize_seed=True,
                    width=resolution,
                    height=resolution,
                    num_inference_steps=6,
                    api_name="/infer"
                )
                if isinstance(res, (list, tuple)):
                    p = res[0] if isinstance(res[0], str) else res[0].get("path")
                    if p and os.path.exists(str(p)):
                        map_path = str(out_dir / f"{m}_{ts}.png")
                        Image.open(p).save(map_path)
                        maps[m] = map_path
                        continue
        except Exception as e:
            print(f"Texture map {m} gen note: {e}")

        # Fallback per map
        default_color = (200, 180, 150) if m == "albedo" else (128, 128, 128)
        map_path = str(out_dir / f"{m}_{ts}.png")
        Image.new("RGB", (resolution, resolution), default_color).save(map_path)
        maps[m] = map_path

    manifest = create_manifest(
        name=prompt[:50].replace(" ", "_"),
        type="texture_pbr",
        source_brick="forge-texture",
        prompt_or_desc=prompt,
        files=maps,
        params={"resolution": resolution},
        metadata={"maps": list(maps.keys())},
        commercial_ok=True,
    )
    manifest.save(out_dir / f"manifest_{ts}.json")

    return {"maps": maps, "manifest": manifest.to_dict()}

def build_ui():
    with gr.Blocks(title="Forge-Texture") as demo:
        gr.Markdown("# Forge-Texture (PBR Maps)")
        p = gr.Textbox("ancient stone brick wall, mossy, medieval game asset")
        ref = gr.Image(label="Optional style ref", type="pil")
        res = gr.Slider(512, 2048, value=1024, step=256, label="Resolution")
        btn = gr.Button("Generate PBR Set")
        out = gr.JSON()
        btn.click(generate_pbr, [p, ref, res], out)
    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__":
    demo.launch(server_name="0.0.0.0", server_port=7863, mcp_server=True)