app.py

"""
GameForge Space Backend
=======================
HF Space using gradio.Server for custom frontend + ZeroGPU.
Serves the GameForge web app with queued, concurrent-safe API endpoints.

Deploy: push to HF Spaces with ZeroGPU hardware.
"""

import os
import sys
import json
import asyncio
from pathlib import Path
from typing import Optional, Dict, Any, List

import spaces
from gradio import Server
from gradio.data_classes import FileData
from fastapi.responses import HTMLResponse, FileResponse, JSONResponse
import yaml

# Add parent to path for gameforge imports
sys.path.insert(0, str(Path(__file__).parent.parent))

# Import GameForge modules
from gameforge.config.registry_loader import get_registry
from gameforge.config.prompts import get_template_for_asset, format_npc_dialogue, list_templates
from gameforge.engine.router import get_router
from gameforge.engine.validator import validate_asset
from gameforge.engine.converter import convert_asset, export_for_engine

# Initialize
registry = get_registry()
router = get_router()

# Storage for generated assets
STORAGE_DIR = Path("/tmp/gameforge_assets")
STORAGE_DIR.mkdir(parents=True, exist_ok=True)

# Create the Server app
app = Server()


# ============================================================================
# Static file serving
# ============================================================================

STATIC_DIR = Path(__file__).parent / "static"


@app.get("/")
async def homepage():
    """Serve the main frontend."""
    html_path = STATIC_DIR / "index.html"
    if html_path.exists():
        return HTMLResponse(html_path.read_text(encoding="utf-8"))
    return HTMLResponse("<h1>GameForge</h1><p>Frontend not found</p>")


@app.get("/static/{path:path}")
async def static_files(path: str):
    """Serve static assets (CSS, JS, fonts)."""
    file_path = STATIC_DIR / path
    if file_path.exists() and file_path.is_file():
        media_types = {
            ".css": "text/css",
            ".js": "application/javascript",
            ".png": "image/png",
            ".jpg": "image/jpeg",
            ".svg": "image/svg+xml",
            ".woff2": "font/woff2",
            ".woff": "font/woff",
            ".ttf": "font/ttf",
        }
        ext = file_path.suffix.lower()
        media_type = media_types.get(ext, "application/octet-stream")
        return FileResponse(str(file_path), media_type=media_type)
    return JSONResponse({"error": "Not found"}, status_code=404)


# ============================================================================
# API Endpoints (gradio_client compatible via @app.api)
# ============================================================================

@app.api()
def registry_info() -> Dict[str, Any]:
    """Get registry summary and all models."""
    summary = registry.summary()
    models = []
    for asset_type in registry.list_asset_types():
        asset = registry.get_asset(asset_type)
        if asset:
            for variant, model in asset.variants.items():
                models.append({
                    "asset_type": asset_type,
                    "variant": variant,
                    "model": model.model,
                    "type": model.type,
                    "license": model.license,
                    "hardware": model.hardware,
                    "status": model.status,
                    "free": model.is_free,
                    "commercial_safe": model.is_commercial_safe,
                    "space_id": model.space_id,
                })
    return {"summary": summary, "models": models}


@app.api()
def get_route(asset_type: str, variant: str = "primary") -> Dict[str, Any]:
    """Get routing info for a model."""
    decision = router.route(asset_type, variant)
    return decision.to_dict()


@app.api()
def list_pipelines() -> List[Dict[str, Any]]:
    """List available pipeline definitions."""
    pipes_dir = Path(__file__).parent.parent / "pipelines"
    result = []
    for path in sorted(pipes_dir.glob("*.yaml")):
        try:
            with open(path) as f:
                data = yaml.safe_load(f)
            result.append({
                "name": path.stem,
                "description": data.get("description", ""),
                "version": data.get("version", ""),
                "steps": len(data.get("steps", [])),
                "defaults": data.get("defaults", {}),
            })
        except Exception:
            pass
    return result


@app.api()
def get_pipeline(name: str) -> Dict[str, Any]:
    """Get full pipeline definition."""
    pipes_dir = Path(__file__).parent.parent / "pipelines"
    for path in pipes_dir.glob(f"{name}.yaml"):
        with open(path) as f:
            return yaml.safe_load(f)
    return {"error": f"Pipeline not found: {name}"}


@app.api()
def format_prompt(asset_type: str, user_prompt: str, model_family: str = "") -> Dict[str, str]:
    """Format a prompt using game-specific templates."""
    template = get_template_for_asset(asset_type, model_family)
    if template:
        return template.format(user_prompt)
    return {"prompt": user_prompt, "negative_prompt": ""}


@app.api()
def format_npc(text: str, emotion: str = "neutral", speaker: str = "") -> Dict[str, str]:
    """Format NPC dialogue with emotion."""
    return {"formatted": format_npc_dialogue(text, emotion, speaker)}


@app.api()
def list_templates_api() -> List[Dict[str, str]]:
    """List all prompt templates."""
    templates = list_templates()
    return [
        {
            "name": t.name,
            "asset_type": t.asset_type,
            "model_family": t.model_family,
            "examples": t.example_prompts[:3],
        }
        for t in templates
    ]


# ============================================================================
# GPU-accelerated generation endpoints (ZeroGPU)
# ============================================================================

@spaces.GPU(duration=60)
def _generate_image(prompt: str, negative_prompt: str = "", steps: int = 4) -> str:
    """Generate an image using the HF Inference API."""
    from huggingface_hub import InferenceClient
    import tempfile

    token = os.environ.get("HF_TOKEN", "")
    if not token:
        for tp in [os.path.expanduser("~/.cache/huggingface/token"),
                    os.path.expanduser("~/.huggingface/token")]:
            if os.path.isfile(tp):
                token = open(tp).read().strip()
                break

    client = InferenceClient(token=token, provider="hf-inference")
    image = client.text_to_image(
        prompt,
        model="black-forest-labs/FLUX.1-schnell",
        negative_prompt=negative_prompt or None,
        num_inference_steps=steps,
    )

    out_path = str(STORAGE_DIR / f"img_{hash(prompt) % 100000}.png")
    image.save(out_path)
    return out_path


@app.api()
def generate_image(prompt: str, negative_prompt: str = "", steps: int = 4) -> FileData:
    """Generate an image from text. Returns PNG."""
    out_path = _generate_image(prompt, negative_prompt, steps)
    return FileData(path=out_path)


@spaces.GPU(duration=120)
def _generate_3d(image_path: str) -> str:
    """Generate 3D mesh from image via TRELLIS.2 Space."""
    from gradio_client import Client
    import shutil

    client = Client("microsoft/TRELLIS.2")
    result = client.predict(image_path, api_name="/generate")

    # Result is typically a file path or tuple
    if isinstance(result, tuple):
        mesh_path = result[0]
    elif isinstance(result, str) and os.path.isfile(result):
        mesh_path = result
    else:
        return None

    out_path = str(STORAGE_DIR / f"mesh_{hash(image_path) % 100000}.glb")
    if isinstance(mesh_path, str) and os.path.isfile(mesh_path):
        shutil.copy2(mesh_path, out_path)
        return out_path
    return None


@app.api()
def generate_3d(image_path: FileData) -> Optional[FileData]:
    """Generate 3D mesh from an image. Returns GLB."""
    result = _generate_3d(image_path["path"])
    if result:
        return FileData(path=result)
    return None


@spaces.GPU(duration=60)
def _generate_voice(text: str) -> str:
    """Generate voice via MeloTTS."""
    from gradio_client import Client
    import shutil

    client = Client("mrfakename/MeloTTS")
    result = client.predict(text, api_name="/synthesize")

    if isinstance(result, tuple):
        audio_path = result[0]
    elif isinstance(result, str) and os.path.isfile(result):
        audio_path = result
    else:
        return None

    out_path = str(STORAGE_DIR / f"voice_{hash(text) % 100000}.wav")
    if os.path.isfile(audio_path):
        shutil.copy2(audio_path, out_path)
        return out_path
    return None


@app.api()
def generate_voice(text: str) -> Optional[FileData]:
    """Generate NPC voice from text. Returns WAV."""
    result = _generate_voice(text)
    if result:
        return FileData(path=result)
    return None


@spaces.GPU(duration=120)
def _generate_video(prompt: str) -> str:
    """Generate video via LTX-2 Turbo."""
    from gradio_client import Client
    import shutil

    client = Client("alexnasa/ltx-2-TURBO")
    result = client.predict(prompt, api_name="/generate")

    if isinstance(result, tuple):
        video_path = result[0]
    elif isinstance(result, str) and os.path.isfile(result):
        video_path = result
    else:
        return None

    out_path = str(STORAGE_DIR / f"video_{hash(prompt) % 100000}.mp4")
    if os.path.isfile(video_path):
        shutil.copy2(video_path, out_path)
        return out_path
    return None


@app.api()
def generate_video(prompt: str) -> Optional[FileData]:
    """Generate a video from text. Returns MP4."""
    result = _generate_video(prompt)
    if result:
        return FileData(path=result)
    return None


@spaces.GPU(duration=120)
def _generate_music(prompt: str) -> str:
    """Generate music via ACE-Step."""
    from gradio_client import Client
    import shutil

    client = Client("victor/ace-step-jam")
    result = client.predict(prompt, api_name="/predict")

    if isinstance(result, tuple):
        audio_path = result[0]
    elif isinstance(result, str) and os.path.isfile(result):
        audio_path = result
    else:
        return None

    out_path = str(STORAGE_DIR / f"music_{hash(prompt) % 100000}.wav")
    if os.path.isfile(audio_path):
        shutil.copy2(audio_path, out_path)
        return out_path
    return None


@app.api()
def generate_music(prompt: str) -> Optional[FileData]:
    """Generate music from text. Returns WAV."""
    result = _generate_music(prompt)
    if result:
        return FileData(path=result)
    return None


@spaces.GPU(duration=60)
def _generate_sfx(prompt: str) -> str:
    """Generate sound effect via TangoFlux."""
    from gradio_client import Client
    import shutil

    client = Client("declare-lab/TangoFlux")
    result = client.predict(prompt, api_name="/predict")

    if isinstance(result, tuple):
        audio_path = result[0]
    elif isinstance(result, str) and os.path.isfile(result):
        audio_path = result
    else:
        return None

    out_path = str(STORAGE_DIR / f"sfx_{hash(prompt) % 100000}.wav")
    if os.path.isfile(audio_path):
        shutil.copy2(audio_path, out_path)
        return out_path
    return None


@app.api()
def generate_sfx(prompt: str) -> Optional[FileData]:
    """Generate sound effect from text. Returns WAV."""
    result = _generate_sfx(prompt)
    if result:
        return FileData(path=result)
    return None


# ============================================================================
# Asset management
# ============================================================================

@app.api()
def list_assets(folder: str = "") -> List[Dict[str, Any]]:
    """List generated assets."""
    search_dir = STORAGE_DIR / folder if folder else STORAGE_DIR
    if not search_dir.exists():
        return []
    assets = []
    for f in sorted(search_dir.rglob("*")):
        if f.is_file():
            stat = f.stat()
            assets.append({
                "name": f.name,
                "path": str(f),
                "size": stat.st_size,
                "format": f.suffix.lower(),
                "modified": stat.st_mtime,
            })
    return assets


@app.api()
def delete_asset(path: str) -> Dict[str, bool]:
    """Delete a generated asset."""
    p = Path(path)
    if p.exists() and p.is_file() and p.is_relative_to(STORAGE_DIR):
        p.unlink()
        return {"deleted": True}
    return {"deleted": False}


@app.api()
def validate_asset_api(path: str, checks: str = "file_exists,non_empty") -> Dict[str, Any]:
    """Validate an asset file."""
    check_list = [c.strip() for c in checks.split(",")]
    return validate_asset(path, check_list)


@app.api()
def convert_asset_api(input_path: str, target_format: str, engine: str = "") -> Optional[FileData]:
    """Convert an asset format."""
    eng = engine if engine else None
    result = convert_asset(input_path, target_format, engine=eng)
    if result:
        return FileData(path=result)
    return None


# ============================================================================
# MCP Tool Registration (for HF Spaces MCP support)
# ============================================================================

@app.mcp.tool()
def gameforge_generate(
    asset_type: str,
    prompt: str,
) -> str:
    """Generate a game asset (image, 3D, voice, music, video, SFX).

    Args:
        asset_type: Type of asset - "image", "3d", "voice", "music", "video", "sfx"
        prompt: Description of the asset to generate
    """
    generators = {
        "image": _generate_image,
        "3d": lambda p: _generate_3d(p),  # Takes image path, would need intermediate
        "voice": _generate_voice,
        "music": _generate_music,
        "video": _generate_video,
        "sfx": _generate_sfx,
    }
    fn = generators.get(asset_type)
    if fn:
        result = fn(prompt)
        return result or "Generation failed"
    return f"Unknown asset type: {asset_type}"


@app.mcp.tool()
def gameforge_list_models() -> str:
    """List all available AI models for game asset generation."""
    models = []
    for asset_type in registry.list_asset_types():
        asset = registry.get_asset(asset_type)
        if asset:
            for variant, model in asset.variants.items():
                free = "FREE" if model.is_free else "PAID"
                models.append(f"{asset_type}/{variant}: {model.model} [{free}]")
    return "\n".join(models)


# ============================================================================
# Launch
# ============================================================================

if __name__ == "__main__":
    app.launch(show_error=True)