"""
Spaces as Agent Tools — Interactive Demo & Reference

Demonstrates how any coding agent can call Hugging Face Spaces
via their agents.md endpoint, and how to chain Spaces together
(text → image → 3D model) with zero client libraries.
"""

import gradio as gr
import httpx
import json
import os

HF_TOKEN = os.environ.get("HF_TOKEN", "")
AUTH_HEADER = {"Authorization": f"Bearer {HF_TOKEN}"} if HF_TOKEN else {}


# ── helpers for Explore tab (raw REST — same protocol agents use) ────

def fetch_agents_md(space_id: str) -> str:
    """Fetch the agents.md for any Space."""
    url = f"https://huggingface.co/spaces/{space_id}/agents.md"
    r = httpx.get(url, timeout=15)
    r.raise_for_status()
    return r.text


def fetch_api_schema(space_id: str) -> dict:
    """Fetch the /gradio_api/info schema for a Space."""
    subdomain = space_id.replace("/", "-").replace(".", "-").lower()
    url = f"https://{subdomain}.hf.space/gradio_api/info"
    r = httpx.get(url, timeout=15, headers=AUTH_HEADER)
    r.raise_for_status()
    return r.json()


def format_schema_summary(schema: dict) -> str:
    """Pretty-print the API schema as a Markdown summary."""
    lines = []
    for name, ep in schema.get("named_endpoints", {}).items():
        params = ep.get("parameters", [])
        returns = ep.get("returns", [])
        lines.append(f"### `{name}`")
        if params:
            lines.append("**Parameters:**")
            for p in params:
                pname = p.get("parameter_name", p.get("label", "?"))
                ptype = p.get("python_type", {}).get("type", "?")
                default = p.get("parameter_default", "—")
                lines.append(f"- `{pname}`: `{ptype}` — default: `{default}`")
        if returns:
            lines.append("**Returns:**")
            for r in returns:
                rtype = r.get("python_type", {}).get("type", "?")
                lines.append(f"- `{r.get('label', '?')}`: `{rtype}`")
        lines.append("")
    return "\n".join(lines)


def call_space_rest(space_id: str, endpoint: str, data: list, timeout: int = 120):
    """
    Call a Space via raw REST (POST + SSE poll).
    Works well for Spaces that don't need file uploads.
    """
    subdomain = space_id.replace("/", "-").replace(".", "-").lower()
    base = f"https://{subdomain}.hf.space/gradio_api/call"

    r = httpx.post(
        f"{base}/{endpoint}",
        json={"data": data},
        headers={**AUTH_HEADER, "Content-Type": "application/json"},
        timeout=30,
    )
    r.raise_for_status()
    event_id = r.json().get("event_id")
    if not event_id:
        raise RuntimeError(f"No event_id: {r.text}")

    final_data = None
    error_data = None
    with httpx.stream("GET", f"{base}/{endpoint}/{event_id}",
                      headers=AUTH_HEADER, timeout=timeout) as stream:
        current_event = None
        for line in stream.iter_lines():
            if line.startswith("event: "):
                current_event = line[7:].strip()
            elif line.startswith("data: "):
                payload = line[6:]
                if current_event == "error":
                    error_data = payload
                else:
                    try:
                        final_data = json.loads(payload)
                    except json.JSONDecodeError:
                        final_data = payload

    if error_data:
        raise RuntimeError(f"Space returned error: {error_data}")
    if final_data is None:
        raise RuntimeError("No data received from SSE stream")
    return final_data


# ── Tab 1: Explore any Space ────────────────────────────────────────

def explore_space(space_id: str):
    """Fetch agents.md + schema for a given Space."""
    space_id = space_id.strip()
    if not space_id:
        return "⚠️ Enter a Space ID (e.g. `black-forest-labs/flux-klein-9b-kv`)"

    try:
        agents_md = fetch_agents_md(space_id)
    except Exception as e:
        agents_md = f"❌ Could not fetch agents.md: {e}"

    try:
        schema = fetch_api_schema(space_id)
        schema_md = format_schema_summary(schema)
    except Exception as e:
        schema_md = f"❌ Could not fetch schema: {e}"

    curl_cmd = f'curl "https://huggingface.co/spaces/{space_id}/agents.md"'

    full_output = f"""## agents.md

```
{agents_md}
```

---

## curl command (copy this into your coding agent)

```bash
{curl_cmd}
```

---

## API Endpoints

{schema_md}
"""
    return full_output


# ── Tab 2: Chain demo ───────────────────────────────────────────────

def chain_prompt_to_3d(prompt: str, image_source: str, progress=gr.Progress(track_tqdm=False)):
    """
    Chain two Spaces together:
      1. Image generator  →  creates an image from text
      2. TRELLIS.2        →  turns it into a 3D GLB model
    """
    from gradio_client import Client, handle_file

    if not HF_TOKEN:
        return (
            None, None, None,
            "❌ **HF_TOKEN not set.**\n\n"
            "**To fix this (one-time setup, takes 2 minutes):**\n\n"
            "1. Go to [huggingface.co/settings/tokens](https://huggingface.co/settings/tokens) "
            "and create a token (type: **Read**)\n"
            "2. Go to [⚙️ this Space's Settings](https://huggingface.co/spaces/Crealink/spaces-as-agent-tools/settings)\n"
            "3. Scroll to **Variables and secrets** → click **New secret**\n"
            "4. Name: `HF_TOKEN` — Value: paste your token\n"
            "5. Click **Save** — the Space restarts automatically\n"
            "6. Come back here and click **Generate** again"
        )

    logs = []
    trellis_space = "microsoft/TRELLIS.2"
    image_path = None

    # ── Step 1: Generate image ──
    if "ERNIE" in image_source:
        # ERNIE-Image: runs on CPU (calls Baidu API) — NO GPU quota used
        ernie_space = "baidu/ERNIE-Image-Turbo"
        logs.append(f"### Step 1 — Image generation with [{ernie_space}](https://huggingface.co/spaces/{ernie_space})")
        logs.append(f"*Using ERNIE-Image (API-backed — no GPU quota used)*\n")
        logs.append(f"Prompt: *\"{prompt}\"*\n")
        progress(0.05, desc="Generating image with ERNIE-Image (no GPU quota)...")

        try:
            result = call_space_rest(
                ernie_space,
                "generate_image",
                data=[prompt, "1024x1024", -1, True],
                timeout=120,
            )
            # Result: [image_file_data, revised_prompt_string]
            image_data = result[0] if isinstance(result, list) else result
            if isinstance(image_data, dict):
                image_path = image_data.get("url") or image_data.get("path")
            else:
                image_path = str(image_data)

            revised = result[1] if isinstance(result, list) and len(result) > 1 else None
            logs.append(f"✅ Image generated!\n")
            if revised:
                logs.append(f"Revised prompt: *\"{revised[:150]}...\"*\n")
            logs.append(f"Image: `{image_path}`\n")

        except Exception as e:
            logs.append(f"❌ ERNIE-Image failed: `{e}`\n")
            return None, None, None, "\n".join(logs)

    else:
        # FLUX Klein 9B: uses ZeroGPU — counts against quota
        flux_space = "black-forest-labs/flux-klein-9b-kv"
        logs.append(f"### Step 1 — Image generation with [{flux_space}](https://huggingface.co/spaces/{flux_space})")
        logs.append(f"*Using FLUX Klein 9B (ZeroGPU — uses GPU quota)*\n")
        logs.append(f"Prompt: *\"{prompt}\"*\n")
        progress(0.05, desc="Connecting to FLUX Klein 9B...")

        try:
            flux = Client(flux_space)
            progress(0.1, desc="Generating image with FLUX Klein 9B...")

            flux_result = flux.predict(
                prompt=prompt,
                input_images=[],
                seed=0,
                randomize_seed=True,
                width=1024,
                height=1024,
                num_inference_steps=4,
                prompt_upsampling=False,
                api_name="/generate",
            )

            if isinstance(flux_result, (tuple, list)):
                image_path = flux_result[0]
                seed_used = flux_result[1] if len(flux_result) > 1 else "?"
            else:
                image_path = flux_result
                seed_used = "?"

            if isinstance(image_path, dict):
                image_path = image_path.get("url") or image_path.get("path")

            logs.append(f"✅ Image generated! (seed: {seed_used})\n")
            logs.append(f"Image: `{image_path}`\n")

        except Exception as e:
            error_msg = str(e)
            if "GPU quota" in error_msg:
                logs.append(f"❌ **GPU quota exceeded.**\n")
                logs.append("💡 **Tip:** Switch to **ERNIE-Image** in the dropdown — it uses no GPU quota.\n")
            else:
                logs.append(f"❌ FLUX failed: `{error_msg}`\n")
            return None, None, None, "\n".join(logs)

    progress(0.35, desc="Image ready — connecting to TRELLIS.2...")

    # ── Step 2: Send image to TRELLIS.2 for 3D generation ──
    logs.append(f"### Step 2 — 3D generation with [{trellis_space}](https://huggingface.co/spaces/{trellis_space})")
    logs.append(f"Passing image to TRELLIS.2...\n")

    try:
        trellis = Client(trellis_space)

        progress(0.4, desc="Starting TRELLIS session...")
        trellis.predict(api_name="/start_session")
        logs.append("✅ Session started\n")

        progress(0.45, desc="Preprocessing image (background removal)...")
        preprocessed = trellis.predict(
            input=handle_file(image_path),
            api_name="/preprocess_image",
        )
        logs.append("✅ Image preprocessed (background removed)\n")

        if isinstance(preprocessed, (tuple, list)):
            preprocessed_path = preprocessed[0]
        else:
            preprocessed_path = preprocessed
        if isinstance(preprocessed_path, dict):
            preprocessed_path = preprocessed_path.get("url") or preprocessed_path.get("path")

    except Exception as e:
        error_msg = str(e)
        if "GPU quota" in error_msg:
            logs.append(f"❌ **GPU quota exceeded** on TRELLIS.2.\n")
            logs.append("The image was generated (see left). 3D conversion needs GPU time — try again when quota resets.\n")
        else:
            logs.append(f"❌ Preprocess failed: `{error_msg}`\n")
        return image_path, None, None, "\n".join(logs)

    # ── Step 3: Generate 3D model ──
    progress(0.55, desc="Generating 3D model (this takes 1–2 min)...")

    try:
        result_3d = trellis.predict(
            image=handle_file(preprocessed_path),
            seed=0,
            resolution="512",
            ss_guidance_strength=7.5,
            ss_guidance_rescale=0.7,
            ss_sampling_steps=12,
            ss_rescale_t=5.0,
            shape_slat_guidance_strength=7.5,
            shape_slat_guidance_rescale=0.5,
            shape_slat_sampling_steps=12,
            shape_slat_rescale_t=3.0,
            tex_slat_guidance_strength=1.0,
            tex_slat_guidance_rescale=0.0,
            tex_slat_sampling_steps=12,
            tex_slat_rescale_t=3.0,
            api_name="/image_to_3d",
        )
        logs.append("✅ 3D model generated!\n")

        if isinstance(result_3d, (tuple, list)):
            preview_path = result_3d[0]
        else:
            preview_path = result_3d
        if isinstance(preview_path, dict):
            preview_path = preview_path.get("url") or preview_path.get("path")
        if preview_path:
            logs.append(f"Preview: `{preview_path}`\n")

    except Exception as e:
        error_msg = str(e)
        if "GPU quota" in error_msg:
            logs.append(f"❌ **GPU quota exceeded** during 3D generation.\n")
        else:
            logs.append(f"❌ 3D generation failed: `{error_msg}`\n")
        return image_path, None, None, "\n".join(logs)

    # ── Step 4: Extract GLB mesh ──
    progress(0.85, desc="Extracting GLB mesh...")
    glb_path = None

    try:
        glb_result = trellis.predict(
            decimation_target=300000,
            texture_size=2048,
            api_name="/extract_glb",
        )
        logs.append("✅ GLB mesh extracted!\n")

        if isinstance(glb_result, (tuple, list)):
            glb_path = glb_result[0]
        else:
            glb_path = glb_result
        if isinstance(glb_path, dict):
            glb_path = glb_path.get("url") or glb_path.get("path")
        if glb_path:
            logs.append(f"GLB file: `{glb_path}`\n")

    except Exception as e:
        logs.append(f"⚠️ GLB extraction failed: `{e}`\n")

    progress(1.0, desc="Done!")
    logs.append("---\n### 🎉 Pipeline complete!")
    return image_path, preview_path, glb_path, "\n".join(logs)


# ── Tab 3: curl reference ────────────────────────────────────────────

CURL_REFERENCE = """
## Spaces as Agent Tools — curl Reference

Every Gradio Space on Hugging Face exposes a plain-text `agents.md` that any coding agent
(Claude Code, Codex, OpenCode, Pi, etc.) can read and act on. **No client library required.**

---

### 1️⃣ Discover — Read agents.md

```bash
curl "https://huggingface.co/spaces/black-forest-labs/flux-klein-9b-kv/agents.md"
```

Returns 4 lines:
```
To use this application (black-forest-labs/flux-klein-9b-kv: Generate or edit images...):
API schema: GET https://black-forest-labs-flux-klein-9b-kv.hf.space/gradio_api/info
Call endpoint: POST https://black-forest-labs-flux-klein-9b-kv.hf.space/gradio_api/call/{endpoint} {"data": [...]}
Poll result: GET https://black-forest-labs-flux-klein-9b-kv.hf.space/gradio_api/call/{endpoint}/{event_id}
Auth: Bearer $HF_TOKEN
```

---

### 2️⃣ Inspect — Get the API schema

```bash
curl "https://black-forest-labs-flux-klein-9b-kv.hf.space/gradio_api/info" | python3 -m json.tool
```

---

### 3️⃣ Call — Submit a job

```bash
EVENT_ID=$(curl -s -X POST \\
  "https://baidu-ernie-image-turbo.hf.space/gradio_api/call/generate_image" \\
  -H "Authorization: Bearer $HF_TOKEN" \\
  -H "Content-Type: application/json" \\
  -d '{"data": ["a crystal dragon on a mountain", "1024x1024", -1, true]}' \\
  | python3 -c "import sys,json; print(json.load(sys.stdin)['event_id'])")

echo "Event ID: $EVENT_ID"
```

### 4️⃣ Poll — Get the result

```bash
curl -N "https://baidu-ernie-image-turbo.hf.space/gradio_api/call/generate_image/$EVENT_ID" \\
  -H "Authorization: Bearer $HF_TOKEN"
```

---

### 5️⃣ Chain — pass image into TRELLIS.2

When chaining Spaces that accept **file inputs**, the image must be uploaded to the target Space first:

**Option A — Upload + REST (curl-based agents):**
```bash
# 1. Download the image
curl -o image.webp "$IMAGE_URL"

# 2. Upload to TRELLIS.2
UPLOADED=$(curl -s -X POST \\
  "https://microsoft-trellis-2.hf.space/gradio_api/upload" \\
  -H "Authorization: Bearer $HF_TOKEN" \\
  -F "files=@image.webp")
TRELLIS_PATH=$(echo $UPLOADED | python3 -c "import sys,json; print(json.load(sys.stdin)[0])")

# 3. Call preprocess + image_to_3d + extract_glb
EVENT_ID=$(curl -s -X POST \\
  "https://microsoft-trellis-2.hf.space/gradio_api/call/preprocess_image" \\
  -H "Authorization: Bearer $HF_TOKEN" \\
  -H "Content-Type: application/json" \\
  -d '{"data": [{"path": "'$TRELLIS_PATH'", "meta": {"_type": "gradio.FileData"}}]}' \\
  | python3 -c "import sys,json; print(json.load(sys.stdin)['event_id'])")
curl -N "https://microsoft-trellis-2.hf.space/gradio_api/call/preprocess_image/$EVENT_ID" \\
  -H "Authorization: Bearer $HF_TOKEN"
```

**Option B — gradio_client (recommended for Python):**
```python
from gradio_client import Client, handle_file

# Step 1: generate image (ERNIE = no GPU quota)
ernie = Client("baidu/ERNIE-Image-Turbo")
image, _ = ernie.predict(
    prompt="a crystal dragon",
    size="1024x1024", seed=-1, use_pe=True,
    api_name="/generate_image"
)

# Step 2: image → 3D
trellis = Client("microsoft/TRELLIS.2")
trellis.predict(api_name="/start_session")
preprocessed = trellis.predict(
    input=handle_file(image),
    api_name="/preprocess_image"
)
trellis.predict(
    image=handle_file(preprocessed), seed=0, resolution="512",
    ss_guidance_strength=7.5, ss_guidance_rescale=0.7,
    ss_sampling_steps=12, ss_rescale_t=5.0,
    shape_slat_guidance_strength=7.5, shape_slat_guidance_rescale=0.5,
    shape_slat_sampling_steps=12, shape_slat_rescale_t=3.0,
    tex_slat_guidance_strength=1.0, tex_slat_guidance_rescale=0.0,
    tex_slat_sampling_steps=12, tex_slat_rescale_t=3.0,
    api_name="/image_to_3d"
)
glb, _ = trellis.predict(
    decimation_target=300000, texture_size=2048,
    api_name="/extract_glb"
)
print(f"3D model: {glb}")
```

---

### 🔑 Authentication & GPU quotas

| Tier | ZeroGPU quota | Spaces affected |
|------|---------------|-----------------|
| No token | 2 min/day | FLUX, TRELLIS, Z-Image |
| Free account | 3.5 min/day | FLUX, TRELLIS, Z-Image |
| HF Pro | 25 min/day | FLUX, TRELLIS, Z-Image |

**Tip:** [ERNIE-Image](https://huggingface.co/spaces/baidu/ERNIE-Image-Turbo) runs on CPU (calls Baidu's API) — **no GPU quota used**. Use it for image generation to save all your quota for 3D.

Get your token at: [huggingface.co/settings/tokens](https://huggingface.co/settings/tokens)

---

### 📋 Popular Agent-Compatible Spaces

| Space | Task | GPU quota? | agents.md |
|-------|------|-----------|-----------|
| [ERNIE-Image](https://huggingface.co/spaces/baidu/ERNIE-Image-Turbo) | Text → Image | ❌ Free | `curl https://huggingface.co/spaces/baidu/ERNIE-Image-Turbo/agents.md` |
| [FLUX Klein 9B](https://huggingface.co/spaces/black-forest-labs/flux-klein-9b-kv) | Text → Image | ✅ Uses quota | `curl https://huggingface.co/spaces/black-forest-labs/flux-klein-9b-kv/agents.md` |
| [TRELLIS.2](https://huggingface.co/spaces/microsoft/TRELLIS.2) | Image → 3D | ✅ Uses quota | `curl https://huggingface.co/spaces/microsoft/TRELLIS.2/agents.md` |
| [Qwen3 ASR](https://huggingface.co/spaces/Qwen/Qwen3-ASR) | Audio → Text | ✅ Uses quota | `curl https://huggingface.co/spaces/Qwen/Qwen3-ASR/agents.md` |
| [Z-Image Turbo](https://huggingface.co/spaces/mrfakename/Z-Image-Turbo) | Text → Image | ✅ Uses quota | `curl https://huggingface.co/spaces/mrfakename/Z-Image-Turbo/agents.md` |
"""


# ── Build the UI ─────────────────────────────────────────────────────

INTRO_MD = """
# 🤖 Spaces as Agent Tools

Every Gradio Space on Hugging Face exposes a plain-text **`agents.md`** that coding agents
(Claude Code, Codex, OpenCode, Pi, etc.) can call directly — **no client library, no SDK**.

A coding agent reads 4 lines (`agents.md`), learns the API schema, and makes standard HTTP calls. That's it.

---

**How to use this demo — pick a tab:**

| Tab | What it does | Token needed? |
|-----|-------------|---------------|
| 🔍 **Explore** | Look up any Space's agent interface & API schema | ❌ No |
| ⛓️ **Chain** | Generate an image, then turn it into a 3D model | ✅ Yes ([setup guide below](#chain-setup)) |
| 📋 **Reference** | Copy-paste curl commands for your coding agent | ❌ No |

> **Shortcut for any Space:** On any Space page on huggingface.co, click the **Agents** button in the header → curl command copied to clipboard.
"""

CHAIN_SETUP_MD = """
**Chain two Spaces with zero integration code:**
1. An **image generator** creates an image from your text prompt
2. **[TRELLIS.2](https://huggingface.co/spaces/microsoft/TRELLIS.2)** turns that image into a 3D GLB model

---

<details id="chain-setup">
<summary><b>⚙️ First-time setup (click to expand) — takes 2 minutes</b></summary>

This tab calls other Spaces on your behalf, so it needs a Hugging Face token.

**Step 1 — Get a token:**
1. Go to [huggingface.co/settings/tokens](https://huggingface.co/settings/tokens)
2. Click **Create new token** → type: **Read** → copy it

**Step 2 — Add it to this Space:**
1. Go to [⚙️ Space Settings → Variables & Secrets](https://huggingface.co/spaces/Crealink/spaces-as-agent-tools/settings)
2. Click **New secret** → Name: `HF_TOKEN` → Value: paste token → **Save**
3. Space restarts (~30 sec) → come back and click **Generate**

</details>

---

**▶ Choose an image source, type a prompt, and click Generate.**
"""

EXPLORE_MD = """
Type any Gradio Space ID and click **Explore** to see:
- Its **`agents.md`** — the 4-line text a coding agent reads
- A **curl command** you can paste into your agent
- The full **API schema** — every endpoint, parameter, type, and default

Try the examples below, or type any Space you know (e.g. your own).
"""

with gr.Blocks(
    title="Spaces as Agent Tools",
    theme=gr.themes.Soft(),
) as demo:
    gr.Markdown(INTRO_MD)

    with gr.Tabs():
        # ── Tab 1: Explorer ──
        with gr.TabItem("🔍 Explore a Space", id="explore"):
            gr.Markdown(EXPLORE_MD)
            with gr.Row():
                space_input = gr.Textbox(
                    label="Space ID",
                    placeholder="e.g. black-forest-labs/flux-klein-9b-kv",
                    value="black-forest-labs/flux-klein-9b-kv",
                    scale=3,
                )
                explore_btn = gr.Button("🔍 Explore", variant="primary", scale=1)
            explore_output = gr.Markdown(label="Result")
            explore_btn.click(explore_space, inputs=[space_input], outputs=[explore_output])

            gr.Examples(
                examples=[
                    ["black-forest-labs/flux-klein-9b-kv"],
                    ["microsoft/TRELLIS.2"],
                    ["baidu/ERNIE-Image-Turbo"],
                    ["Qwen/Qwen3-ASR"],
                    ["mrfakename/Z-Image-Turbo"],
                ],
                inputs=[space_input],
            )

        # ── Tab 2: Chain demo ──
        with gr.TabItem("⛓️ Chain: Text → Image → 3D", id="chain"):
            gr.Markdown(CHAIN_SETUP_MD)
            with gr.Row():
                chain_prompt = gr.Textbox(
                    label="Prompt",
                    placeholder="Describe what you want to create...",
                    value="a ceramic fox figurine, white background",
                    scale=3,
                )
            with gr.Row():
                image_source = gr.Dropdown(
                    label="Image generator",
                    choices=[
                        "ERNIE-Image (free — no GPU quota)",
                        "FLUX Klein 9B (higher quality — uses GPU quota)",
                    ],
                    value="ERNIE-Image (free — no GPU quota)",
                    scale=2,
                )
                chain_btn = gr.Button("🚀 Generate", variant="primary", scale=1)

            with gr.Row():
                chain_image = gr.Image(label="Step 1: Image from text", type="filepath")
                chain_preview = gr.Video(label="Step 2: TRELLIS.2 → 3D Preview")

            with gr.Row():
                chain_glb = gr.File(label="Step 3: Download GLB")

            chain_log = gr.Markdown(label="Pipeline Log")

            chain_btn.click(
                chain_prompt_to_3d,
                inputs=[chain_prompt, image_source],
                outputs=[chain_image, chain_preview, chain_glb, chain_log],
            )

        # ── Tab 3: Reference ──
        with gr.TabItem("📋 curl Reference", id="reference"):
            gr.Markdown(CURL_REFERENCE)

if __name__ == "__main__":
    demo.launch()