backend/server.py

"""FastAPI routes attached to gr.Server.

CRITICAL FIXES vs previous build:
  • Strips <think>...</think> blocks (and other prefixes) from raw model output
    BEFORE JSON parsing. The fine-tuned MiniCPM-V emits a `<think>\n\n</think>`
    preamble — that was breaking parse and triggering the
    "raw JSON toast on top" + repeated `Inference failed` errors.
  • Accepts BOTH formats:
        (A) bare ElysiumResponse        — what the model actually emits
        (B) {user_msg, elysium_response} envelope — historical schema
    Auto-wraps (A) into (B) so downstream code is unchanged.
  • Fallback never dumps raw JSON into `direct_answer`; instead returns a
    short, human-readable error string so the UI cannot accidentally toast
    a 2 KB blob of JSON.
  • Always attaches a complete `_runtime` block (metrics, audio, tool_results,
    per_agent_audio, civilization_map) so the frontend never hits null/undefined.
  • Adds /api/civilization_map endpoint returning a compact spatial snapshot
    for the minimap (id, type, color, x_norm, y_norm, edges).
"""
import io
import json
import re
import time
import traceback
from typing import List, Optional
from PIL import Image
from fastapi import UploadFile, File, Form, HTTPException
from fastapi.responses import JSONResponse
from fastapi.staticfiles import StaticFiles

import spaces

from .config import (AUDIO_CACHE, MAX_UPLOAD_FILES, MAX_UPLOAD_BYTES,
                     ALLOWED_MIME_TYPES)
from .model_loader import make_llm
from .grammar import load_grammar
from .prompt_builder import build_messages, new_session_meta
from .schema import ElysiumEnvelope, ElysiumResponse
from .hypergraph import persistence
from .hypergraph.engine import Hypergraph
from .tools.dispatcher import execute_all
from .tts.debate_sequencer import build_debate, build_per_agent_audio


# ─── Singletons ─────────────────────────────────────────────────────────────
HG: Hypergraph = persistence.load()
GRAMMAR = load_grammar()
CIVILIZATION_START_TS = time.time()

# Pre-compiled regex to strip <think>...</think> blocks (including newlines)
_THINK_RE = re.compile(r"<think\b[^>]*>.*?</think\s*>", re.DOTALL | re.IGNORECASE)
# Match the first balanced JSON object in arbitrary text
_JSON_OBJECT_RE = re.compile(r"\{.*\}", re.DOTALL)


# ─── GPU-bound inference ────────────────────────────────────────────────────
@spaces.GPU(duration=120)
def _gpu_infer(messages: list, max_tokens: int = 4096) -> str:
    llm = make_llm()
    out = llm.create_chat_completion(
        messages=messages,
        max_tokens=max_tokens,
        temperature=0.7,
        grammar=GRAMMAR,
    )
    return out["choices"][0]["message"]["content"]


# ────────────────────────────────────────────────────────────────────────────
# Robust raw-output → ElysiumEnvelope parser
# ────────────────────────────────────────────────────────────────────────────
def _clean_raw_model_output(raw: str) -> str:
    """Strip thinking blocks, code fences, and surrounding whitespace.

    Handles all of the following patterns observed in real model outputs:
      <think>...</think>\n\n{json}
      ```json\n{json}\n```
      Some text {json} trailing text
    """
    if not raw:
        return ""
    s = raw.strip()
    # Strip <think>...</think>
    s = _THINK_RE.sub("", s).strip()
    # Strip ```json ... ``` fences
    if s.startswith("```"):
        s = re.sub(r"^```[a-zA-Z]*\s*", "", s)
        s = re.sub(r"\s*```\s*$", "", s)
        s = s.strip()
    return s


def _extract_json_blob(cleaned: str) -> Optional[dict]:
    """Try direct json.loads, then fall back to first-balanced-object extract."""
    if not cleaned:
        return None
    try:
        return json.loads(cleaned)
    except Exception:
        pass
    m = _JSON_OBJECT_RE.search(cleaned)
    if not m:
        return None
    try:
        return json.loads(m.group(0))
    except Exception:
        # Try a more aggressive cleanup: collapse stray newlines INSIDE strings
        # (the model sometimes wraps long string values across lines)
        candidate = m.group(0)
        # Replace literal newlines inside likely string contexts. This is a
        # best-effort: only safe for the model's typical output shape.
        try:
            candidate2 = re.sub(r'(?<!\\)\n', ' ', candidate)
            return json.loads(candidate2)
        except Exception:
            return None


def _coerce_to_envelope(raw: str, user_text: str) -> Optional[ElysiumEnvelope]:
    """Parse arbitrary model output into a validated ElysiumEnvelope.

    Returns None only when nothing JSON-shaped can be extracted at all.
    """
    cleaned = _clean_raw_model_output(raw)
    blob = _extract_json_blob(cleaned)
    if blob is None:
        return None

    # Case A: bare ElysiumResponse → wrap it
    if "elysium_response" not in blob and ("schema_version" in blob or
                                            "interaction_type" in blob or
                                            "hypergraph_delta" in blob):
        try:
            resp = ElysiumResponse.model_validate(blob)
            return ElysiumEnvelope(user_msg=user_text, elysium_response=resp)
        except Exception as e:
            print(f"[parse] bare-response validate failed: {e}")
            # Fall through to envelope attempt

    # Case B: full envelope
    try:
        return ElysiumEnvelope.model_validate(blob)
    except Exception as e:
        print(f"[parse] envelope validate failed: {e}")

    # Last resort: synthesize a minimal envelope, keeping any direct_answer
    meta = new_session_meta()
    direct = ""
    if isinstance(blob, dict):
        direct = str(blob.get("direct_answer", "") or "")[:300]
    return ElysiumEnvelope(
        user_msg=user_text,
        elysium_response=ElysiumResponse(
            session_id=meta["session_id"],
            timestamp_utc=meta["timestamp_utc"],
            interaction_type="SIMPLE_REPLY",
            direct_answer=direct or "(model response could not be fully parsed)",
        ),
    )


def _fallback_envelope(user_text: str, err: str) -> dict:
    """Last-resort fallback that NEVER contains raw JSON in direct_answer."""
    meta = new_session_meta()
    # Keep direct_answer short & human readable. Never echo raw JSON.
    safe_err = (err or "unknown error").splitlines()[0][:200]
    resp = ElysiumResponse(
        session_id=meta["session_id"],
        timestamp_utc=meta["timestamp_utc"],
        interaction_type="SIMPLE_REPLY",
        direct_answer=f"The civilization could not process that turn ({safe_err}). Try again.",
    )
    return {
        "user_msg": user_text,
        "elysium_response": resp.model_dump(),
        "_runtime": _empty_runtime(),
    }


def _empty_runtime() -> dict:
    """Always-present runtime envelope; never None on any field."""
    return {
        "tool_results":        [],
        "audio_url":           None,
        "per_agent_audio":     [],
        "metrics":             _baseline_metrics(),
        "attachment_errors":   [],
        "attachments_processed": [],
        "civilization_map":    _civilization_map_snapshot(),
    }


def _baseline_metrics() -> dict:
    nodes = HG.node_count()
    edges = HG.edge_count()
    density = 0.0 if nodes == 0 else min(1.0, edges / max(1, nodes * 1.4))
    age_min = int((time.time() - CIVILIZATION_START_TS) / 60)
    return {
        "mycelium_density_pct": round(density * 100),
        "council_active":       0,
        "knowledge_growth":     0,
        "coherence_pct":        70,
        "civilization_age_min": age_min,
        "nodes":                nodes,
        "edges":                edges,
        "alert_level":          "CALM",
    }


def _civilization_metrics(resp: ElysiumResponse) -> dict:
    nodes = HG.node_count()
    edges = HG.edge_count()
    density = 0.0 if nodes == 0 else min(1.0, edges / max(1, nodes * 1.4))
    council_active = len(resp.council_deliberation.agent_outputs or [])
    knowledge_growth = len(resp.hypergraph_delta.nodes_added or [])
    coherence = 1.0 - float(resp.strain_metadata.cognitive_strain or 0.3)
    age_min = int((time.time() - CIVILIZATION_START_TS) / 60)
    return {
        "mycelium_density_pct": round(density * 100),
        "council_active":       council_active,
        "knowledge_growth":     knowledge_growth,
        "coherence_pct":        round(max(0.0, min(1.0, coherence)) * 100),
        "civilization_age_min": age_min,
        "nodes":                nodes,
        "edges":                edges,
        "alert_level":          resp.ui_directives.alert_level or "CALM",
    }


def _civilization_map_snapshot() -> dict:
    """Compact spatial summary used by the minimap on cold start
    or as a sync source. Frontend already maintains its own spatial layout;
    this endpoint provides the canonical ids/types/edges."""
    nodes_out = []
    for i in HG.g.node_indexes():
        d = HG.g[i]
        nodes_out.append({
            "node_id":  d["node_id"],
            "label":    d.get("label", d["node_id"]),
            "type":     d.get("node_type", "DOMAIN"),
        })
    edges_out = []
    for s, t in HG.g.edge_list():
        d = HG.g.get_edge_data(s, t)
        edges_out.append({
            "edge_id":  d.get("edge_id", f"e_{s}_{t}"),
            "source":   HG.g[s]["node_id"],
            "target":   HG.g[t]["node_id"],
            "type":     d.get("edge_type", "GENERIC"),
            "weight":   d.get("weight", 0.5),
        })
    return {"nodes": nodes_out, "edges": edges_out,
            "node_count": HG.node_count(), "edge_count": HG.edge_count()}


async def _load_attachment(uf: UploadFile) -> Optional[dict]:
    if uf is None or not uf.filename:
        return None
    raw = await uf.read()
    if not raw:
        return None
    if len(raw) > MAX_UPLOAD_BYTES:
        return {"kind": "error", "name": uf.filename,
                "error": f"file too large ({len(raw)} > {MAX_UPLOAD_BYTES} bytes)"}
    mime = (uf.content_type or "").lower()
    if mime not in ALLOWED_MIME_TYPES:
        low = uf.filename.lower()
        if low.endswith(".pdf"):
            mime = "application/pdf"
        elif low.endswith((".png", ".jpg", ".jpeg", ".webp", ".gif")):
            mime = "image/jpeg"
        else:
            return {"kind": "error", "name": uf.filename, "error": f"unsupported type: {mime}"}

    if mime == "application/pdf":
        return {"kind": "pdf", "bytes": raw, "name": uf.filename}
    try:
        img = Image.open(io.BytesIO(raw))
        img.load()
        return {"kind": "image", "image": img, "name": uf.filename, "bytes": raw}
    except Exception as e:
        return {"kind": "error", "name": uf.filename, "error": f"image decode failed: {e}"}


def attach(app):
    """Register all /api routes on the gr.Server FastAPI app."""

    app.mount("/audio", StaticFiles(directory=str(AUDIO_CACHE)), name="audio")

    @app.get("/api/health")
    async def health():
        return {"status": "ok",
                "nodes": HG.node_count(),
                "edges": HG.edge_count(),
                "grammar": GRAMMAR is not None,
                "max_upload_files": MAX_UPLOAD_FILES}

    @app.get("/api/hypergraph")
    async def hypergraph():
        nodes, edges = [], []
        for i in HG.g.node_indexes():
            d = HG.g[i]
            nodes.append({
                "node_id":   d["node_id"],
                "label":     d.get("label", d["node_id"]),
                "node_type": d.get("node_type", "DOMAIN"),
                "payload":   d.get("payload", {}),
                "embedding_hint": d.get("embedding_hint", ""),
            })
        for s, t in HG.g.edge_list():
            d = HG.g.get_edge_data(s, t)
            edges.append({
                "edge_id":         d.get("edge_id", f"e_{s}_{t}"),
                "source_node_id":  HG.g[s]["node_id"],
                "target_node_id":  HG.g[t]["node_id"],
                "edge_type":       d.get("edge_type", "GENERIC"),
                "weight":          d.get("weight", 0.5),
                "payload":         d.get("payload", {}),
            })
        return {
            "nodes": nodes, "edges": edges,
            "node_count": HG.node_count(), "edge_count": HG.edge_count(),
        }

    @app.get("/api/civilization_map")
    async def civilization_map():
        return _civilization_map_snapshot()

    @app.get("/api/node/{node_id}")
    async def node_detail(node_id: str):
        if node_id not in HG._idx:
            raise HTTPException(404, "node not found")
        idx = HG._idx[node_id]
        n   = HG.g[idx]

        incoming, outgoing = [], []
        for s, t in HG.g.edge_list():
            d = HG.g.get_edge_data(s, t)
            if t == idx:
                incoming.append({
                    "from": HG.g[s]["node_id"],
                    "from_label": HG.g[s].get("label", ""),
                    "edge_type": d.get("edge_type", ""),
                    "weight": d.get("weight", 0.5),
                })
            if s == idx:
                outgoing.append({
                    "to": HG.g[t]["node_id"],
                    "to_label": HG.g[t].get("label", ""),
                    "edge_type": d.get("edge_type", ""),
                    "weight": d.get("weight", 0.5),
                })
        return {
            "node_id":  n["node_id"],
            "label":    n.get("label", n["node_id"]),
            "node_type": n.get("node_type", ""),
            "payload":  n.get("payload", {}),
            "embedding_hint": n.get("embedding_hint", ""),
            "incoming": incoming,
            "outgoing": outgoing,
            "degree":   len(incoming) + len(outgoing),
        }

    @app.post("/api/turn")
    async def turn(user_text: str = Form(""),
                   attachments: List[UploadFile] = File(default=[])):
        try:
            # 1. Process up to MAX_UPLOAD_FILES attachments
            atts = []
            for uf in (attachments or [])[:MAX_UPLOAD_FILES]:
                a = await _load_attachment(uf)
                if a:
                    atts.append(a)

            errors = [a for a in atts if a.get("kind") == "error"]
            valid  = [a for a in atts if a.get("kind") in ("image", "pdf")]

            # 2. Build messages with hypergraph context
            messages = build_messages(user_text, valid, HG.context_summary())

            # 3. GPU inference (returns strict JSON, possibly wrapped in <think>)
            raw = _gpu_infer(messages)
            if not raw:
                return JSONResponse(_fallback_envelope(user_text, "empty model output"))

            # 4. Robust parse: strips <think>, extracts JSON, accepts bare ElysiumResponse
            envelope = _coerce_to_envelope(raw, user_text)
            if envelope is None:
                return JSONResponse(_fallback_envelope(user_text, "no JSON in model output"))

            resp = envelope.elysium_response

            # 5. Apply hypergraph delta
            try:
                HG.apply_delta(resp.hypergraph_delta)
                persistence.save(HG)
            except Exception as e:
                print(f"[hypergraph] apply_delta failed: {e}")
                traceback.print_exc()

            # 6. Execute tools (best-effort)
            tool_results = []
            try:
                tool_results = execute_all(resp.tool_calls) if resp.tool_calls else []
            except Exception as e:
                print(f"[tools] execute_all failed: {e}")

            # 7. Audio drama (combined + per-agent), best-effort
            audio_url = None
            per_agent_audio = []
            if resp.council_deliberation.debate_mode in ("AUDIO_DRAMA", "SILENT") \
               and resp.council_deliberation.agent_outputs:
                try:
                    agents_dump = [a.model_dump() for a in resp.council_deliberation.agent_outputs]
                    if resp.council_deliberation.debate_mode == "AUDIO_DRAMA":
                        audio_url = build_debate(agents_dump)
                    per_agent_audio = build_per_agent_audio(agents_dump)
                except Exception as e:
                    print(f"[tts] debate failed: {e}")
                    traceback.print_exc()

            payload = envelope.model_dump()
            payload["_runtime"] = {
                "tool_results":          tool_results,
                "audio_url":             audio_url,
                "per_agent_audio":       per_agent_audio,
                "metrics":               _civilization_metrics(resp),
                "attachment_errors":     [{"name": e["name"], "error": e["error"]} for e in errors],
                "attachments_processed": [{"kind": a["kind"], "name": a["name"]} for a in valid],
                "civilization_map":      _civilization_map_snapshot(),
            }
            return JSONResponse(payload)

        except Exception as e:
            traceback.print_exc()
            return JSONResponse(
                _fallback_envelope(user_text, str(e)), status_code=200)

    @app.post("/api/reset")
    async def reset():
        global HG, CIVILIZATION_START_TS
        HG = Hypergraph()
        persistence.save(HG)
        CIVILIZATION_START_TS = time.time()
        return {"status": "reset"}