#!/usr/bin/env python3
"""
Mem0-Sidecar für Mission Control 2.0 — das auto-lernende, semantische Gedächtnis.

WARUM ein eigener Dienst? Mem0 + chromadb laufen nur unter Python 3.12 (`~/.mem0/venv`),
das MC2-Backend aber unter Python 3.14 (kann mem0 nicht importieren). Darum kapselt dieser
schlanke FastAPI-Dienst die Mem0-Memory-Instanz und exponiert sie auf localhost. MC2
(`backend/services/memory.py`) spricht ihn per HTTP an — die `/api/memory`-API-Form bleibt
nach außen unverändert (UI + MCP-Server kompatibel).

Aufbau:
- Vektor-Store : Chroma embedded (kein Docker auf der Box) — Pfad MEM0_CHROMA_PATH.
- Embeddings   : llama.cpp `embed`-Rolle über llama-swap (/v1/embeddings, 1024 Dim).
- LLM-Extraktion: lokales `fast`-Hirn (Qwen3.6) über llama-swap (/v1/chat/completions).
                  Thinking wird abgeschaltet (NoThinkLLM), sonst bricht die JSON-Extraktion ab.

Läuft als systemd-User-Dienst (deploy/mem0-service.service) im `~/.mem0/venv`.
Bind: 127.0.0.1 (nur lokal; MC2 proxyt nach außen).
"""

import logging
import os
from contextlib import asynccontextmanager

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel

from mem0 import Memory
from mem0.configs.llms.openai import OpenAIConfig
from mem0.llms.openai import OpenAILLM

log = logging.getLogger("mem0_service")

# --- Konfiguration (alles über Env überschreibbar; Defaults = Box-Stand) ----------
USER_ID = os.environ.get("MEM0_USER_ID", "mission-control")  # ein geteiltes Gehirn
EMBED_URL = os.environ.get("MEM0_EMBED_URL", "http://127.0.0.1:8080/v1")
LLM_URL = os.environ.get("MEM0_LLM_URL", "http://127.0.0.1:8080/v1")
EMBED_MODEL = os.environ.get("MEM0_EMBED_MODEL", "embed")
LLM_MODEL = os.environ.get("MEM0_LLM_MODEL", "fast")
EMBED_DIMS = int(os.environ.get("MEM0_EMBED_DIMS", "1024"))
CHROMA_PATH = os.environ.get("MEM0_CHROMA_PATH", "/srv/models/mem0/chroma")
HISTORY_DB = os.environ.get("MEM0_HISTORY_DB", "/srv/models/mem0/history.db")
COLLECTION = os.environ.get("MEM0_COLLECTION", "mc2")
API_KEY = os.environ.get("MEM0_API_KEY", "sk-local")  # llama.cpp ignoriert den Key
LLM_MAX_TOKENS = int(os.environ.get("MEM0_LLM_MAX_TOKENS", "2048"))
SEARCH_TOP_K = int(os.environ.get("MEM0_SEARCH_TOP_K", "50"))

# Kategorien = fundierte Memory-Taxonomie (semantisch/prozedural/episodisch), bewusst knapp (4).
# Mem0-OSS kann nicht nativ klassifizieren (Cloud-Feature) → wir lassen das Hirn beim Lernen
# einordnen (classify_facts). Die Beschreibungen steuern die Treffsicherheit der Auto-Zuordnung.
CATEGORIES = ("identity", "knowledge", "rules", "events")
DEFAULT_CATEGORY = "knowledge"
CATEGORY_DESCRIPTIONS = {
    "identity": "Identität & Vorlieben des Nutzers: wer er ist (Name, Rolle), wie er angesprochen werden "
                "will, persönliche Vorlieben und seine Arbeitsweise.",
    "knowledge": "Wissen & Fakten: stabile Fakten über Stack, Projekte, Tools, Infrastruktur und Umgebung.",
    "rules": "Regeln & Konventionen: verbindliche Anweisungen und Workflows — wie etwas gemacht werden "
             "soll, Do's und Don'ts.",
    "events": "Ereignisse & Entscheidungen: was passiert ist, getroffene Entscheidungen, zeitgebundene Vorgänge.",
}


class NoThinkLLM(OpenAILLM):
    """`fast` (Qwen3.6) ist ein Thinking-Modell. Unter response_format=json_object
    verbrennt das Reasoning Tokens und schneidet die JSON-Fakten ab (leeres Ergebnis).
    `enable_thinking=false` (chat_template_kwargs) → deterministische, schnelle Extraktion."""

    def generate_response(self, messages, response_format=None, tools=None,
                          tool_choice="auto", **kwargs):
        eb = dict(kwargs.pop("extra_body", {}) or {})
        eb.setdefault("chat_template_kwargs", {"enable_thinking": False})
        return super().generate_response(
            messages, response_format=response_format, tools=tools,
            tool_choice=tool_choice, extra_body=eb, **kwargs,
        )


def build_memory() -> Memory:
    os.environ.setdefault("OPENAI_API_KEY", API_KEY)
    os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
    os.makedirs(CHROMA_PATH, exist_ok=True)
    os.makedirs(os.path.dirname(HISTORY_DB), exist_ok=True)
    cfg = {
        "vector_store": {"provider": "chroma", "config": {
            "collection_name": COLLECTION, "path": CHROMA_PATH}},
        "embedder": {"provider": "openai", "config": {
            "model": EMBED_MODEL, "embedding_dims": EMBED_DIMS,
            "openai_base_url": EMBED_URL, "api_key": API_KEY}},
        "llm": {"provider": "openai", "config": {
            "model": LLM_MODEL, "openai_base_url": LLM_URL, "api_key": API_KEY,
            "temperature": 0.1, "max_tokens": LLM_MAX_TOKENS}},
        "history_db_path": HISTORY_DB,
        # Fakten in der Originalsprache halten (deutsche Eingaben bleiben deutsch).
        "custom_instructions": (
            "Bewahre die Originalsprache der Fakten — deutsche Eingaben bleiben deutsch. "
            "Schreibe jeden Fakt knapp und atomar."
        ),
    }
    m = Memory.from_config(cfg)
    # LLM gegen die Thinking-freie Variante tauschen (gleiche Verbindung).
    m.llm = NoThinkLLM(OpenAIConfig(
        model=LLM_MODEL, openai_base_url=LLM_URL, api_key=API_KEY,
        temperature=0.1, max_tokens=LLM_MAX_TOKENS,
    ))
    return m


_mem: Memory | None = None


def mem() -> Memory:
    global _mem
    if _mem is None:
        _mem = build_memory()
    return _mem


@asynccontextmanager
async def lifespan(_app: FastAPI):
    logging.basicConfig(level=logging.INFO)
    try:
        mem()  # Memory beim Start aufbauen (Chroma öffnen, Modelle anpingen)
        log.info("Mem0 bereit (chroma=%s, embed=%s, llm=%s)", CHROMA_PATH, EMBED_MODEL, LLM_MODEL)
    except Exception:
        log.exception("Mem0-Init fehlgeschlagen (Dienst läuft, /health meldet down)")
    yield


app = FastAPI(title="MC2 Mem0 Sidecar", lifespan=lifespan)


# --- Mapping Mem0 <-> MC2-API-Form -----------------------------------------------
def _to_item(r: dict, score=None) -> dict:
    meta = r.get("metadata") or {}
    out = {
        "id": r.get("id"),
        "content": r.get("memory", ""),
        "category": meta.get("category") or "stable",
        "source": meta.get("source") or "auto",
        "created_at": r.get("created_at") or "",
        "updated_at": r.get("updated_at") or r.get("created_at") or "",
    }
    if score is not None:
        try:
            out["score"] = round(float(score), 4)
        except (TypeError, ValueError):
            pass
    return out


class MemIn(BaseModel):
    content: str
    category: str = DEFAULT_CATEGORY
    source: str = "manual"


class MemUp(BaseModel):
    content: str | None = None
    category: str | None = None


class LearnIn(BaseModel):
    text: str | None = None
    messages: list[dict] | None = None
    source: str = "auto"
    category: str = DEFAULT_CATEGORY
    classify: bool = True   # Hirn ordnet die gelernten Fakten automatisch ein


@app.get("/health")
def health() -> dict:
    try:
        mem().get_all(filters={"user_id": USER_ID}, top_k=1)
        return {"ok": True}
    except Exception as exc:  # noqa: BLE001
        raise HTTPException(503, f"mem0 nicht bereit: {exc}")


@app.get("/memory")
def list_memory(q: str = "", category: str = "") -> list[dict]:
    """q gesetzt → semantische Suche (mit Relevanz-Score). Sonst → alle Fakten."""
    if q:
        res = mem().search(q, filters={"user_id": USER_ID}, top_k=SEARCH_TOP_K)
        items = [_to_item(r, r.get("score")) for r in res.get("results", [])]
    else:
        res = mem().get_all(filters={"user_id": USER_ID}, top_k=1000)
        items = [_to_item(r) for r in res.get("results", [])]
        items.sort(key=lambda i: i.get("created_at", ""), reverse=True)
    if category:
        items = [i for i in items if i["category"] == category]
    return items


@app.post("/memory", status_code=201)
def add_memory(body: MemIn) -> dict:
    """Manueller/agentischer Einzel-Fakt → VERBATIM speichern (infer=False, keine
    LLM-Umformung). Auto-Lernen aus Gesprächen läuft über /learn (infer=True)."""
    content = body.content.strip()
    if not content:
        raise HTTPException(400, "content leer")
    res = mem().add(
        [{"role": "user", "content": content}],
        user_id=USER_ID, infer=False,
        metadata={"category": body.category, "source": body.source},
    )
    results = res.get("results", [])
    if not results:
        raise HTTPException(500, "Mem0 hat nichts gespeichert")
    new_id = results[0]["id"]
    full = mem().get(new_id) or {}
    return _to_item(full) if full else _to_item({"id": new_id, "memory": content,
            "metadata": {"category": body.category, "source": body.source}})


@app.put("/memory/{mid}")
def update_memory(mid: str, body: MemUp) -> dict:
    existing = mem().get(mid)
    if not existing:
        raise HTTPException(404, "Eintrag nicht gefunden")
    cur = _to_item(existing)
    new_content = body.content.strip() if body.content is not None else cur["content"]
    new_cat = body.category if body.category is not None else cur["category"]
    mem().update(mid, data=new_content,
                 metadata={"category": new_cat, "source": cur["source"]})
    full = mem().get(mid) or {}
    return _to_item(full)


@app.delete("/memory/{mid}")
def delete_memory(mid: str) -> dict:
    if not mem().get(mid):
        raise HTTPException(404, "Eintrag nicht gefunden")
    mem().delete(mid)
    return {"ok": True}


@app.get("/graph")
def graph(min_score: float = 0.45, top_k: int = 3) -> dict:
    """Fakten als Graph: Knoten = Fakten, Kanten = semantische Ähnlichkeit (Kosinus
    der gespeicherten Embeddings, je Knoten die top_k Nachbarn ≥ min_score). Für die
    Obsidian-artige Gedächtnis-Visualisierung im UI."""
    items = mem().get_all(filters={"user_id": USER_ID}, top_k=2000).get("results", [])
    nodes = [_to_item(r) for r in items]
    node_ids = [n["id"] for n in nodes]
    idx = {nid: i for i, nid in enumerate(node_ids)}

    # Embeddings direkt aus der Chroma-Collection ziehen (kein Re-Embedding).
    col = mem().vector_store.collection
    raw = col.get(include=["embeddings"])
    raw_ids = raw.get("ids") or []
    raw_embs = raw.get("embeddings")
    edges: list[dict] = []
    vecs = [None] * len(node_ids)
    have = 0
    if raw_embs is not None:
        for rid, emb in zip(raw_ids, raw_embs):
            if rid in idx and emb is not None:
                vecs[idx[rid]] = emb
                have += 1

    if have >= 2:
        import numpy as np
        present = [i for i, v in enumerate(vecs) if v is not None]
        M = np.array([vecs[i] for i in present], dtype=float)
        norms = np.linalg.norm(M, axis=1, keepdims=True)
        norms[norms == 0] = 1.0
        Mn = M / norms
        sim = Mn @ Mn.T
        seen: set[tuple[int, int]] = set()
        for a in range(len(present)):
            order = np.argsort(-sim[a])
            cnt = 0
            for b in order:
                if b == a:
                    continue
                s = float(sim[a][b])
                if s < min_score:
                    break
                i, j = present[a], present[b]
                key = (min(i, j), max(i, j))
                if key not in seen:
                    seen.add(key)
                    edges.append({"source": node_ids[i], "target": node_ids[j],
                                  "weight": round(s, 3)})
                cnt += 1
                if cnt >= top_k:
                    break
    return {"nodes": nodes, "edges": edges}


def classify_facts(facts: list[tuple[str, str]]) -> dict:
    """Ordnet jeden Fakt per LLM GENAU einer Kategorie zu → {id: category}.
    Nutzt dasselbe (Thinking-freie) Hirn wie die Extraktion. Fehler = leeres Mapping (Fallback)."""
    if not facts:
        return {}
    import json as _json
    import re as _re
    cat_lines = "\n".join(f"- {k}: {v}" for k, v in CATEGORY_DESCRIPTIONS.items())
    system = (
        "Du bist ein präziser Klassifikator für ein Langzeitgedächtnis. Ordne jeden Fakt GENAU EINER "
        "Kategorie zu. Verfügbare Kategorien (nur diese Slugs verwenden):\n" + cat_lines +
        "\nAntworte NUR als JSON-Objekt der Form {\"<id>\": \"<slug>\", …}."
    )
    body = "\n".join(f"{fid}: {txt}" for fid, txt in facts)
    try:
        resp = mem().llm.generate_response(
            messages=[{"role": "system", "content": system},
                      {"role": "user", "content": "Fakten:\n" + body}],
            response_format={"type": "json_object"},
        )
        m = _re.search(r"\{.*\}", resp or "", _re.S)
        data = _json.loads(m.group(0)) if m else {}
        return {str(k): v for k, v in data.items() if v in CATEGORIES}
    except Exception:
        log.exception("Auto-Klassifikation fehlgeschlagen")
        return {}


@app.post("/learn")
def learn(body: LearnIn) -> dict:
    """Auto-Lernen: Gesprächs-Turns/Text durchreichen → Mem0 EXTRAHIERT Fakten selbst
    (infer=True). Anschließend ordnet das Hirn jeden neuen Fakt automatisch einer Kategorie zu."""
    msgs = body.messages
    if not msgs and body.text:
        msgs = [{"role": "user", "content": body.text}]
    if not msgs:
        raise HTTPException(400, "text oder messages erforderlich")
    res = mem().add(
        msgs, user_id=USER_ID, infer=True,
        metadata={"category": body.category, "source": body.source},
    )
    results = res.get("results", [])

    cats: dict = {}
    if body.classify:
        to_classify = [(r["id"], r.get("memory", "")) for r in results
                       if r.get("id") and r.get("memory") and r.get("event") in ("ADD", "UPDATE")]
        cats = classify_facts(to_classify)
        for r in results:
            cat = cats.get(r.get("id"))
            if cat:
                try:
                    mem().update(r["id"], data=r.get("memory", ""),
                                 metadata={"category": cat, "source": body.source})
                except Exception:
                    log.debug("Kategorie-Update fehlgeschlagen für %s", r.get("id"))

    return {"results": [
        {"id": r.get("id"), "content": r.get("memory"), "event": r.get("event"),
         "category": cats.get(r.get("id"), body.category)}
        for r in results
    ]}


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="127.0.0.1", port=int(os.environ.get("MEM0_PORT", "8765")))