import gradio as gr
import requests, json, os, time
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
import uvicorn

OLLAMA = "http://localhost:11434"
MODEL = "FableForge-AI/shellwhisperer"

fastapi_app = FastAPI(docs_url=None, redoc_url=None)

@fastapi_app.post("/api/chat")
async def api_chat(request: Request):
    body = await request.json()
    try:
        r = requests.post(f"{OLLAMA}/api/chat", json=body, timeout=120)
        return JSONResponse(content=r.json())
    except Exception as e:
        return JSONResponse(status_code=500, content={"error": str(e)})

# ── helpers ──

def model_info():
    """Get model details and status."""
    info = {
        "ready": False,
        "model": MODEL,
        "size": "?",
        "speed": "?",
        "error": "",
    }
    try:
        r = requests.get(f"{OLLAMA}/api/tags", timeout=5)
        if r.status_code == 200:
            models = r.json().get("models", [])
            for m in models:
                if MODEL in m.get("name", ""):
                    info["ready"] = True
                    size_gb = m.get("size", 0) / 1e9
                    info["size"] = f"{size_gb:.1f} GB"
                    break
            if not info["ready"]:
                info["error"] = "Model not pulled yet"
        else:
            info["error"] = f"Ollama returned {r.status_code}"
    except requests.ConnectionError:
        info["error"] = "Ollama not running"
    except Exception as e:
        info["error"] = str(e)
    return info

def generate(prompt, temp):
    if not prompt.strip():
        yield "⚠️  Enter a prompt first."
        return

    info = model_info()
    if not info["ready"]:
        yield f"⏳ Model loading... ({info['error']})"
        return

    try:
        r = requests.post(f"{OLLAMA}/api/chat", json={
            "model": MODEL,
            "messages": [
                {"role": "system", "content": "You are ShellWhisperer-1.5B, a shell and CLI specialist. Output working code only, no explanations."},
                {"role": "user", "content": prompt}
            ],
            "stream": True,
            "options": {"temperature": temp, "num_ctx": 16384}
        }, stream=True, timeout=120)

        full = []
        for line in r.iter_lines():
            if not line:
                continue
            try:
                d = json.loads(line)
                chunk = d.get("message", {}).get("content", "")
                if chunk:
                    full.append(chunk)
                    yield "".join(full)
            except json.JSONDecodeError:
                pass

        if not full:
            yield "⚠️  Empty response from model. Try again."

    except requests.Timeout:
        yield "⏰  Request timed out after 120s. Try a shorter prompt."
    except Exception as e:
        yield f"❌  Error: {e}"

# ── UI ──

with gr.Blocks(
    title="ShellWhisperer-1.5B · API",
    theme=gr.themes.Soft(),
    fill_height=True,
    css="""footer { display: none !important; }
.status-ok { color: #22c55e; font-weight: 600; }
.status-loading { color: #f59e0b; font-weight: 600; }
.status-err { color: #ef4444; font-weight: 600; }
.api-box { background: #1f2937; color: #e5e7eb; padding: 1em; border-radius: 8px; font-family: monospace; font-size: 0.9em; overflow-x: auto; }
""",
) as demo:

    gr.Markdown("""# 🐚 ShellWhisperer-1.5B · API Demo
**CLI & Shell Code Specialist** · 1 GB · 29 tok/s · 16K context · Apache 2.0

Built by **FableForge AI** — part of the [Mythos model ecosystem](https://github.com/KingLabsA/mythos).
""")

    # ── status dashboard ──
    with gr.Row():
        status_badge = gr.Markdown("⏳ Checking...")
        model_size = gr.Markdown("")
        model_speed = gr.Markdown("")

    with gr.Tabs():
        with gr.TabItem("🧪 Try it"):
            with gr.Row():
                with gr.Column(scale=3):
                    inp = gr.Textbox(
                        label="Prompt",
                        placeholder="Write a bash script to...",
                        lines=4,
                    )
                    with gr.Row():
                        temp = gr.Slider(0.0, 1.0, value=0.3, step=0.05, label="Temperature")
                        btn = gr.Button("🚀 Generate", variant="primary", scale=1, size="lg")
                    out = gr.Textbox(label="Output", lines=16)

                    gr.Markdown("### 💡 Try these")
                    gr.Examples(
                        examples=[
                            ["Write a bash script to monitor a directory for new files and log them"],
                            ["Write a Python script to batch resize images to 800px wide"],
                            ["Write a Docker Compose file for a web app with PostgreSQL"],
                            ["Write a git pre-commit hook that runs tests"],
                        ],
                        inputs=inp,
                        label="",
                    )

        with gr.TabItem("📡 API"):
            gr.Markdown("""### REST API

This space exposes a standard Ollama-compatible chat endpoint.

```
POST /api/chat
Content-Type: application/json
```

**cURL:**
```bash
curl -X POST https://karma-devops-shellwhisperer-demo.hf.space/api/chat \\
  -H "Content-Type: application/json" \\
  -d '{
    "model": "FableForge-AI/shellwhisperer",
    "messages": [{"role": "user", "content": "Write a bash script"}],
    "stream": false,
    "options": {"temperature": 0.3, "num_ctx": 16384}
  }'
```

**Python:**
```python
import requests
r = requests.post("https://karma-devops-shellwhisperer-demo.hf.space/api/chat", json={
    "model": "FableForge-AI/shellwhisperer",
    "messages": [{"role": "user", "content": "Write a bash script"}],
    "stream": False,
    "options": {"temperature": 0.3, "num_ctx": 16384}
})
print(r.json()["message"]["content"])
```
""")

        with gr.TabItem("ℹ️ About"):
            gr.Markdown("""### Model

| Property | Value |
|---|---|
| Name | ShellWhisperer-1.5B |
| Author | FableForge AI ([KingLabsA](https://github.com/KingLabsA)) |
| Size | 1 GB |
| Speed | ~29 tok/s (T4 GPU) |
| Context | 16,384 tokens |
| License | Apache 2.0 |
| Base | Qwen2.5-Coder-1.5B-Instruct |

### Links

- [GitHub: KingLabsA/mythos](https://github.com/KingLabsA/mythos)
- [HuggingFace: King3Djbl](https://huggingface.co/King3Djbl)
- [HuggingFace: fableforge-ai](https://huggingface.co/fableforge-ai)
- [Ollama: FableForge-AI](https://ollama.com/FableForge-AI)
""")

    # ── events ──

    def refresh_status():
        info = model_info()
        if info["ready"]:
            badge = f'<span class="status-ok">✅ Model: Ready</span>'
            size = f'<span class="status-ok">📦 {info["size"]}</span>'
            speed = '<span class="status-ok">⚡ ~29 tok/s</span>'
        else:
            badge = f'<span class="status-loading">⏳ Model: {info["error"]}</span>'
            size = ""
            speed = ""
        return badge, size, speed

    demo.load(fn=refresh_status, outputs=[status_badge, model_size, model_speed])
    gr.Timer(30).tick(fn=refresh_status, outputs=[status_badge, model_size, model_speed])

    btn.click(
        fn=generate,
        inputs=[inp, temp],
        outputs=out,
    )

app = gr.mount_gradio_app(fastapi_app, demo, path="/")

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)