import gradio as gr
import base64
import os
import traceback
from pathlib import Path
from gradio_client import Client

HF_TOKEN = os.environ.get("HF_TOKEN")

def generate_b64(prompt: str, seed: int = 42, steps: int = 8) -> str:
    """Generate an image with Z-Image-Turbo and return it as a base64-encoded WebP string.

    Args:
        prompt: Text prompt describing the desired image.
        seed: Random seed for reproducibility. Default 42.
        steps: Number of inference steps. Default 8.

    Returns:
        Base64-encoded WebP image bytes as an ASCII string, or ERROR: message.
    """
    try:
        if not HF_TOKEN:
            return "ERROR: HF_TOKEN not set in Space secrets"
        client = Client("mrfakename/Z-Image-Turbo", token=HF_TOKEN)

        if prompt == "__DEBUG__":
            info = client.view_api(return_format="dict", print_info=False)
            return f"DEBUG: {str(info)[:3000]}"

        result = client.predict(
            prompt=prompt,
            height=1024,
            width=1024,
            num_inference_steps=int(steps),
            seed=int(seed),
            randomize_seed=False,
            api_name="/generate_image",
        )
        filepath = result[0] if isinstance(result, (tuple, list)) else result
        if isinstance(filepath, dict):
            filepath = filepath.get("path") or filepath.get("url")
        data = Path(filepath).read_bytes()
        return base64.b64encode(data).decode("ascii")
    except Exception as e:
        return f"ERROR: {type(e).__name__}: {str(e)}\n{traceback.format_exc()[:1500]}"

demo = gr.Interface(
    fn=generate_b64,
    inputs=[
        gr.Textbox(label="prompt"),
        gr.Number(label="seed", value=42),
        gr.Number(label="steps", value=8),
    ],
    outputs=gr.Textbox(label="base64_webp"),
    title="Z-Image-Turbo Base64 Proxy",
)

if __name__ == "__main__":
    demo.launch(mcp_server=True, show_error=True)