app.py

"""Command Center chat bridge API for Hugging Face Spaces."""
from __future__ import annotations

import hmac
import json
import logging
import os
import threading
import time
import uuid
from pathlib import Path
from typing import Any, Literal, Optional
from urllib.parse import parse_qs, urlencode, urlparse, urlunparse

import httpx
from fastapi import FastAPI, Header, HTTPException, Query, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from pydantic import BaseModel, Field, field_validator

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger("chat-bridge")

ALLOWED_ORIGINS = [
    "https://simzy420.github.io",
    "http://localhost:5173",
    "http://127.0.0.1:5173",
    "*",
]


def _resolve_sessions_dir() -> Path:
    preferred = Path("/data/sessions")
    fallback = Path("/tmp/sessions")
    try:
        Path("/data").mkdir(parents=True, exist_ok=True)
        preferred.mkdir(parents=True, exist_ok=True)
        probe = preferred / ".write_test"
        probe.write_text("ok", encoding="utf-8")
        probe.unlink(missing_ok=True)
        log.info("Using sessions dir: %s", preferred)
        return preferred
    except OSError as exc:
        log.warning("/data not writable (%s); falling back to %s", exc, fallback)
        fallback.mkdir(parents=True, exist_ok=True)
        return fallback


SESSIONS_DIR = _resolve_sessions_dir()
_lock = threading.RLock()

app = FastAPI(title="Command Center Chat Bridge", version="1.0.0")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=False,
    allow_methods=["*"],
    allow_headers=["*"],
)


class ChatRequest(BaseModel):
    sessionId: str
    clientMsgId: str
    botId: Optional[str] = None
    botName: Optional[str] = None
    text: str
    history: Optional[Any] = None

    @field_validator("sessionId", "clientMsgId")
    @classmethod
    def non_empty(cls, v: str) -> str:
        if not isinstance(v, str) or not v.strip():
            raise ValueError("must be a non-empty string")
        return v

    @field_validator("text")
    @classmethod
    def text_ok(cls, v: str) -> str:
        if not isinstance(v, str):
            raise ValueError("text must be a string")
        if len(v) > 8000:
            raise ValueError("text max 8000 chars")
        return v


class ReplyRequest(BaseModel):
    sessionId: str
    clientMsgId: str
    botId: Optional[str] = None
    text: str
    status: Literal["partial", "final"] = "final"

    @field_validator("sessionId", "clientMsgId")
    @classmethod
    def non_empty(cls, v: str) -> str:
        if not isinstance(v, str) or not v.strip():
            raise ValueError("must be a non-empty string")
        return v


def _session_path(session_id: str) -> Path:
    # Keep filename safe
    safe = "".join(c if c.isalnum() or c in "-_." else "_" for c in session_id)[:200]
    if not safe:
        raise HTTPException(status_code=400, detail="invalid sessionId")
    return SESSIONS_DIR / f"{safe}.json"


def _load_session(session_id: str) -> dict[str, Any]:
    path = _session_path(session_id)
    if not path.exists():
        return {"sessionId": session_id, "messages": []}
    try:
        with path.open("r", encoding="utf-8") as f:
            data = json.load(f)
        if not isinstance(data, dict):
            return {"sessionId": session_id, "messages": []}
        data.setdefault("sessionId", session_id)
        data.setdefault("messages", [])
        return data
    except (json.JSONDecodeError, OSError) as exc:
        log.warning("Failed to load session %s: %s", session_id, exc)
        return {"sessionId": session_id, "messages": []}


def _save_session(session_id: str, data: dict[str, Any]) -> None:
    path = _session_path(session_id)
    tmp = path.with_suffix(".tmp")
    with tmp.open("w", encoding="utf-8") as f:
        json.dump(data, f, ensure_ascii=False, indent=2)
    tmp.replace(path)


def _now_iso() -> str:
    return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())


def _append_key_query(url: str, key: str) -> str:
    parsed = urlparse(url)
    qs = parse_qs(parsed.query, keep_blank_values=True)
    if "key" not in qs:
        qs["key"] = [key]
    new_query = urlencode(qs, doseq=True)
    return urlunparse(parsed._replace(query=new_query))




def _require_bridge_secret(authorization: Optional[str]) -> None:
    secret = (os.environ.get("CHAT_BRIDGE_SECRET") or "").strip()
    if not secret:
        raise HTTPException(status_code=503, detail="CHAT_BRIDGE_SECRET not configured")
    if not authorization or not authorization.startswith("Bearer "):
        raise HTTPException(status_code=401, detail="missing bearer token")
    token = authorization[7:].strip()
    if not hmac.compare_digest(token, secret):
        raise HTTPException(status_code=401, detail="invalid bearer token")


@app.get("/health")
def health() -> dict[str, bool]:
    return {"ok": True}


@app.get("/")
def root() -> dict[str, Any]:
    return {
        "ok": True,
        "service": "command-center-chat-bridge",
        "health": "/health",
        "chat": "/api/chat",
        "reply": "/api/chat/reply",
        "pending": "/api/chat/pending",
    }


@app.options("/api/chat")
@app.options("/api/chat/reply")
@app.options("/api/chat/pending")
async def options_ok() -> JSONResponse:
    return JSONResponse(
        content={"ok": True},
        headers={
            "Access-Control-Allow-Origin": "*",
            "Access-Control-Allow-Methods": "GET, POST, OPTIONS",
            "Access-Control-Allow-Headers": "*",
        },
    )


@app.post("/api/chat")
async def post_chat(body: ChatRequest) -> dict[str, Any]:
    msg = {
        "id": str(uuid.uuid4()),
        "role": "user",
        "botId": body.botId,
        "botName": body.botName,
        "text": body.text,
        "clientMsgId": body.clientMsgId,
        "createdAt": _now_iso(),
        "status": "final",
    }
    with _lock:
        data = _load_session(body.sessionId)
        data["messages"].append(msg)
        _save_session(body.sessionId, data)

    webhook_url = (os.environ.get("GROK_WEBHOOK_URL") or "").strip()
    sender_key = (os.environ.get("GROK_WEBHOOK_SENDER_KEY") or "").strip()

    if not webhook_url or not sender_key:
        log.warning("webhook not configured; stored message only")
        return {
            "ok": True,
            "clientMsgId": body.clientMsgId,
            "warning": "webhook not configured",
        }

    forward_url = _append_key_query(webhook_url, sender_key)
    payload = body.model_dump()
    headers = {
        "Authorization": f"Bearer {sender_key}",
        "X-Webhook-Key": sender_key,
        "Content-Type": "application/json",
    }
    try:
        async with httpx.AsyncClient(timeout=30.0) as client:
            resp = await client.post(forward_url, json=payload, headers=headers)
            if resp.status_code >= 400:
                log.error("webhook returned %s: %s", resp.status_code, resp.text[:500])
                # Still ok for client — message stored
                return {
                    "ok": True,
                    "clientMsgId": body.clientMsgId,
                    "warning": f"webhook status {resp.status_code}",
                }
    except Exception as exc:  # noqa: BLE001
        log.exception("webhook forward failed: %s", exc)
        return {
            "ok": True,
            "clientMsgId": body.clientMsgId,
            "warning": "webhook forward failed",
        }

    return {"ok": True, "clientMsgId": body.clientMsgId}


@app.get("/api/chat")
async def get_chat(sessionId: str = Query(..., min_length=1)) -> dict[str, Any]:
    with _lock:
        data = _load_session(sessionId)
    messages = data.get("messages") or []
    # Normalize shape for client
    out = []
    for m in messages:
        if not isinstance(m, dict):
            continue
        out.append(
            {
                "id": m.get("id"),
                "role": m.get("role"),
                "botId": m.get("botId"),
                "text": m.get("text"),
                "clientMsgId": m.get("clientMsgId"),
                "createdAt": m.get("createdAt"),
                "status": m.get("status"),
            }
        )
    return {"messages": out}


@app.post("/api/chat/reply")
async def post_reply(
    body: ReplyRequest,
    authorization: Optional[str] = Header(default=None),
) -> dict[str, bool]:
    _require_bridge_secret(authorization)

    with _lock:
        data = _load_session(body.sessionId)
        messages = data.get("messages") or []
        updated = False
        for m in messages:
            if (
                isinstance(m, dict)
                and m.get("clientMsgId") == body.clientMsgId
                and m.get("role") == "assistant"
            ):
                m["text"] = body.text
                m["status"] = body.status
                m["botId"] = body.botId if body.botId is not None else m.get("botId")
                m["updatedAt"] = _now_iso()
                updated = True
                break
        if not updated:
            messages.append(
                {
                    "id": str(uuid.uuid4()),
                    "role": "assistant",
                    "botId": body.botId,
                    "text": body.text,
                    "clientMsgId": body.clientMsgId,
                    "createdAt": _now_iso(),
                    "status": body.status,
                }
            )
        data["messages"] = messages
        _save_session(body.sessionId, data)

    return {"ok": True}



@app.get("/api/chat/pending")
async def get_pending(
    authorization: Optional[str] = Header(default=None),
) -> dict[str, Any]:
    """Return unanswered user messages (no matching assistant final for clientMsgId)."""
    _require_bridge_secret(authorization)

    pending: list[dict[str, Any]] = []
    with _lock:
        try:
            session_files = sorted(SESSIONS_DIR.glob("*.json"))
        except OSError as exc:
            log.warning("Failed to list sessions: %s", exc)
            session_files = []

        for path in session_files:
            if path.name.endswith(".tmp"):
                continue
            try:
                with path.open("r", encoding="utf-8") as f:
                    data = json.load(f)
            except (json.JSONDecodeError, OSError) as exc:
                log.warning("Failed to read session file %s: %s", path, exc)
                continue
            if not isinstance(data, dict):
                continue
            session_id = data.get("sessionId") or path.stem
            messages = data.get("messages") or []
            if not isinstance(messages, list):
                continue

            final_assistant_ids: set[str] = set()
            for m in messages:
                if not isinstance(m, dict):
                    continue
                if (
                    m.get("role") == "assistant"
                    and m.get("status") == "final"
                    and isinstance(m.get("clientMsgId"), str)
                ):
                    final_assistant_ids.add(m["clientMsgId"])

            for m in messages:
                if not isinstance(m, dict):
                    continue
                if m.get("role") != "user":
                    continue
                client_msg_id = m.get("clientMsgId")
                if not isinstance(client_msg_id, str) or not client_msg_id:
                    continue
                if client_msg_id in final_assistant_ids:
                    continue
                pending.append(
                    {
                        "sessionId": session_id,
                        "clientMsgId": client_msg_id,
                        "botId": m.get("botId"),
                        "botName": m.get("botName"),
                        "text": m.get("text"),
                        "createdAt": m.get("createdAt"),
                    }
                )

    def _sort_key(item: dict[str, Any]) -> str:
        created = item.get("createdAt")
        return created if isinstance(created, str) else ""

    pending.sort(key=_sort_key)
    return {"messages": pending[:50]}


@app.on_event("startup")
def on_startup() -> None:
    global SESSIONS_DIR
    SESSIONS_DIR = _resolve_sessions_dir()
    log.info("Chat bridge ready; sessions=%s", SESSIONS_DIR)