"""
Forge-Refiner Brick — Edit any previous asset + maintain consistency.
"""
import gradio as gr
import spaces
# 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()
import os
from pathlib import Path
from datetime import datetime

@spaces.GPU(duration=60)
def refine_asset(asset_path: str, edit_instruction: str, parent_manifest: str = "") -> dict:
    out = Path(os.environ.get("FORGE_BRICKS_OUTPUT", "./outputs/forge_refiner"))
    out.mkdir(parents=True, exist_ok=True)
    ts = int(datetime.now().timestamp())

    # Real refinement: if image, use FLUX edit; else metadata + text edit. Route to modality.
    new_file = str(out / f"refined_{ts}.asset")
    refined_data = {"original": asset_path, "instruction": edit_instruction}

    # Simple image edit if path looks like image
    if asset_path.lower().endswith(('.png', '.jpg', '.jpeg')) and os.path.exists(asset_path):
        try:
            target = "black-forest-labs/FLUX.1-schnell"
            if health.is_ok(target) is not False:
                from gradio_client import Client
                from PIL import Image as PILImage
                client = Client(target, timeout=120)
                # For refine, use edit as new prompt (or enhance with original desc)
                res = client.predict(prompt=edit_instruction, seed=-1, randomize_seed=True, width=1024, height=1024, 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:
                        new_file = str(out / f"refined_{ts}.png")
                        PILImage.open(p).save(new_file)
                        refined_data["type"] = "image"
        except Exception as e:
            print(f"Refiner image edit error: {e}")

    if not os.path.exists(new_file) or new_file.endswith('.asset'):
        with open(new_file, 'w') as f:
            f.write(f"Refined asset metadata:\nOriginal: {asset_path}\nInstruction: {edit_instruction}\nTimestamp: {ts}")
        refined_data["type"] = "metadata"

    man = create_manifest(
        name="refined_asset",
        type="edited_asset",
        source_brick="forge-refiner",
        prompt_or_desc=edit_instruction,
        files={"refined": new_file},
        parent_id=parent_manifest,
        metadata=refined_data,
        commercial_ok=True
    )
    man.save(out / f"manifest_{ts}.json")
    return {"refined": new_file, "manifest": man.to_dict()}

def build_ui():
    with gr.Blocks() as d:
        gr.Markdown("# Forge-Refiner (Edit + Consistency)")
        asset = gr.Textbox(label="Previous asset path or manifest id")
        instr = gr.Textbox(label="Edit instruction", value="make it more armored and add glowing accents")
        btn = gr.Button("Refine")
        out = gr.JSON()
        btn.click(refine_asset, [asset, instr], out)
    return d

# 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=7868, mcp_server=True)