Timmy-time-dashboard/src/dashboard/routes/grok.py

"""Grok (xAI) dashboard routes — premium cloud augmentation controls.

Endpoints
---------
GET  /grok/status     — JSON status (API)
POST /grok/toggle     — Enable/disable Grok Mode (HTMX)
POST /grok/chat       — Direct Grok query (HTMX)
GET  /grok/stats      — Usage statistics (JSON)
"""

import logging

from fastapi import APIRouter, Form, Request
from fastapi.responses import HTMLResponse

from config import settings
from dashboard.templating import templates

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/grok", tags=["grok"])

# In-memory toggle state (persists per process lifetime)
_grok_mode_active: bool = False


@router.get("/status", response_class=HTMLResponse)
async def grok_status(request: Request):
    """Return Grok backend status as an HTML dashboard page."""
    from timmy.backends import grok_available

    status = {
        "enabled": settings.grok_enabled,
        "available": grok_available(),
        "active": _grok_mode_active,
        "model": settings.grok_default_model,
        "free_mode": settings.grok_free,
        "max_sats_per_query": settings.grok_max_sats_per_query,
        "api_key_set": bool(settings.xai_api_key),
    }

    # Include usage stats if backend exists
    stats = None
    try:
        from timmy.backends import get_grok_backend

        backend = get_grok_backend()
        stats = {
            "total_requests": backend.stats.total_requests,
            "total_prompt_tokens": backend.stats.total_prompt_tokens,
            "total_completion_tokens": backend.stats.total_completion_tokens,
            "estimated_cost_sats": backend.stats.estimated_cost_sats,
            "errors": backend.stats.errors,
        }
    except Exception as exc:
        logger.warning("Failed to load Grok stats: %s", exc)

    return templates.TemplateResponse(
        request,
        "grok_status.html",
        {
            "status": status,
            "stats": stats,
        },
    )


@router.post("/toggle")
async def toggle_grok_mode(request: Request):
    """Toggle Grok Mode on/off. Returns HTMX partial for the toggle card."""
    global _grok_mode_active

    from timmy.backends import grok_available

    if not grok_available():
        return HTMLResponse(
            '<div class="alert" style="color: var(--danger);">'
            "Grok unavailable — set GROK_ENABLED=true and XAI_API_KEY in .env"
            "</div>",
            status_code=200,
        )

    _grok_mode_active = not _grok_mode_active
    state = "ACTIVE" if _grok_mode_active else "STANDBY"

    logger.info("Grok Mode toggled: %s", state)

    # Log to Spark
    try:
        from spark.engine import spark_engine

        spark_engine.on_tool_executed(
            agent_id="default",
            tool_name="grok_mode_toggle",
            success=True,
        )
    except Exception as exc:
        logger.warning("Failed to log Grok toggle to Spark: %s", exc)

    return HTMLResponse(
        _render_toggle_card(_grok_mode_active),
        status_code=200,
    )


def _run_grok_query(message: str) -> dict:
    """Run a Grok query and return result dict.

    Returns:
        {"response": str | None, "error": str | None}
    """
    from timmy.backends import get_grok_backend, grok_available

    if not grok_available():
        return {
            "response": None,
            "error": "Grok is not available. Set GROK_ENABLED=true and XAI_API_KEY.",
        }

    backend = get_grok_backend()

    invoice_note = ""
    if not settings.grok_free:
        try:
            from lightning.factory import get_backend as get_ln_backend

            ln = get_ln_backend()
            sats = min(settings.grok_max_sats_per_query, settings.grok_sats_hard_cap)
            ln.create_invoice(sats, f"Grok: {message[:50]}")
            invoice_note = f" | {sats} sats"
        except Exception as exc:
            logger.warning("Lightning invoice creation failed: %s", exc)

    try:
        result = backend.run(message)
        return {"response": f"**[Grok]{invoice_note}:** {result.content}", "error": None}
    except Exception as exc:
        logger.exception("Grok query failed")
        return {"response": None, "error": f"Grok error: {exc}"}


@router.post("/chat", response_class=HTMLResponse)
async def grok_chat(request: Request, message: str = Form(...)):
    """Send a message directly to Grok and return HTMX chat partial."""
    from datetime import datetime

    from dashboard.store import message_log

    timestamp = datetime.now().strftime("%H:%M:%S")
    result = _run_grok_query(message)

    user_msg = f"[Ask Grok] {message}"
    message_log.append(role="user", content=user_msg, timestamp=timestamp, source="browser")

    if result["response"]:
        message_log.append(
            role="agent", content=result["response"], timestamp=timestamp, source="browser"
        )
    else:
        message_log.append(
            role="error", content=result["error"], timestamp=timestamp, source="browser"
        )

    return templates.TemplateResponse(
        request,
        "partials/chat_message.html",
        {
            "user_message": user_msg,
            "response": result["response"],
            "error": result["error"],
            "timestamp": timestamp,
        },
    )


@router.get("/stats")
async def grok_stats():
    """Return detailed Grok usage statistics."""
    try:
        from timmy.backends import get_grok_backend

        backend = get_grok_backend()
        return {
            "total_requests": backend.stats.total_requests,
            "total_prompt_tokens": backend.stats.total_prompt_tokens,
            "total_completion_tokens": backend.stats.total_completion_tokens,
            "total_latency_ms": round(backend.stats.total_latency_ms, 2),
            "avg_latency_ms": round(
                backend.stats.total_latency_ms / max(backend.stats.total_requests, 1),
                2,
            ),
            "estimated_cost_sats": backend.stats.estimated_cost_sats,
            "errors": backend.stats.errors,
            "model": settings.grok_default_model,
        }
    except Exception as exc:
        logger.exception("Failed to load Grok stats")
        return {"error": str(exc)}


def _render_toggle_card(active: bool) -> str:
    """Render the Grok Mode toggle card HTML."""
    import html

    color = "#00ff88" if active else "#666"
    state = "ACTIVE" if active else "STANDBY"
    glow = "0 0 20px rgba(0, 255, 136, 0.4)" if active else "none"
    model_name = html.escape(settings.grok_default_model)

    return f"""
    <div id="grok-toggle-card"
         style="border: 2px solid {color}; border-radius: 12px; padding: 16px;
                background: var(--bg-secondary); box-shadow: {glow};
                transition: all 0.3s ease;">
        <div style="display: flex; justify-content: space-between; align-items: center;">
            <div>
                <div style="font-weight: 700; font-size: 1.1rem; color: {color};">
                    GROK MODE: {state}
                </div>
                <div style="font-size: 0.8rem; color: var(--text-muted); margin-top: 4px;">
                    xAI frontier reasoning | {model_name}
                </div>
            </div>
            <button hx-post="/grok/toggle"
                    hx-target="#grok-toggle-card"
                    hx-swap="outerHTML"
                    style="background: {color}; color: #000; border: none;
                           border-radius: 8px; padding: 8px 20px; cursor: pointer;
                           font-weight: 700; font-family: inherit;">
                {"DEACTIVATE" if active else "ACTIVATE"}
            </button>
        </div>
    </div>
    """


def is_grok_mode_active() -> bool:
    """Check if Grok Mode is currently active (used by other modules)."""
    return _grok_mode_active