#!/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
import re
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"))

# --- Rerank-Zweitstufe (07.07.2026): Die Vektor-Suche findet ÄHNLICHES, der Reranker ordnet
# nach ECHTER Relevanz (Qwen3-Reranker via llama-swap /v1/rerank). Es ändert sich NUR die
# Reihenfolge: `score` bleibt der Vektor-Score (RECALL_MIN_SCORE im Hermes-Plugin bleibt
# kalibriert), Konsumenten, die vorne abschneiden (RECALL_LIMIT), bekommen automatisch die
# relevantesten Fakten. Jeder Fehler → lautlos Vektor-Reihenfolge (Gedächtnis bricht NIE).
RERANK_ENABLED = os.environ.get("MC_RERANK_ENABLED", "1") != "0"
RERANK_URL = os.environ.get("MC_RERANK_URL", "http://127.0.0.1:8080/v1/rerank")
RERANK_MODEL = os.environ.get("MC_RERANK_MODEL", "Qwen3-Reranker-0.6B")
RERANK_TIMEOUT = float(os.environ.get("MC_RERANK_TIMEOUT", "8"))
RERANK_CANDIDATES = int(os.environ.get("MC_RERANK_CANDIDATES", "20"))  # nur die Top-N neu ordnen (Latenz)

# 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,
        # Sprache zwingend Deutsch erzwingen — mem0s englischer Extraktions-Prompt zieht sonst
        # Richtung Englisch. Eigennamen/Fachbegriffe bleiben unangetastet.
        "custom_instructions": (
            "RELEVANZ (ZWINGEND, WICHTIGSTE REGEL): Extrahiere NUR dauerhaft nützliche Fakten ÜBER DEN "
            "NUTZER (Identität, Vorlieben, Arbeitsweise), seine Projekte/Entscheidungen/Stände, stabile "
            "Fakten über Stack/Tools/Infrastruktur sowie verbindliche Regeln und Workflows. "
            "NIEMALS FAKTEN ÜBER DEN ASSISTENTEN SELBST (ZWINGEND): Extrahiere KEINE Aussagen über den "
            "Assistenten/die KI (Hermes, Lucy), seine Fähigkeiten, seinen Zustand, seine Tools/MCP-Server, "
            "seine Cronjobs, Commit-/Versionsstände oder welche Modelle er nutzt. Das Subjekt eines Fakts "
            "ist IMMER der Nutzer oder sein Projekt — niemals 'der Assistent', 'Hermes' oder 'Lucy'. "
            "KEINE ZEITGEBUNDENEN ZUSTÄNDE: Ignoriere transiente Momentaufnahmen (Fehler-/Job-Läufe, "
            "Commit-Rückstand, Versions-Updates, 'gerade dabei'-/'hochgefahren'-/'aufgewärmt'-Status, "
            "Test- und Onboarding-Fragen). IGNORIERE VOLLSTÄNDIG und gib KEINE Fakten zurück bei: "
            "Tagesnachrichten und Welt-/Politik-/Börsen-/Promi-Ereignissen, Wetter, allgemeinem Weltwissen "
            "ohne Bezug zum Nutzer, vom Assistenten erzeugten Zusammenfassungen oder Recherche-Ergebnissen, "
            "Smalltalk, sowie Meta-Aussagen über den Gesprächsverlauf (z.B. 'der Nutzer fragt nach X', "
            "'der Nutzer möchte Y prüfen', 'der Assistent tat Z'). Speichere nur, was in Wochen noch "
            "nützlich ist; im Zweifel NICHT speichern. "
            "SPRACHE (ZWINGEND): Formuliere ALLE Fakten auf DEUTSCH — unabhängig von der "
            "Eingabesprache, niemals auf Englisch übersetzen. Eigennamen, Befehle, Pfade, Modell- und "
            "Produktnamen sowie etablierte Fachbegriffe (z.B. PowerShell-Push, Deploy, Qwen3.6, Neovim, "
            "Vulkan) bleiben unverändert. Schreibe jeden Fakt knapp und atomar — genau eine Aussage."
        ),
    }
    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}")


def _rerank(query: str, items: list[dict]) -> list[dict]:
    """Suchtreffer nach echter Relevanz umordnen; bei jedem Fehler unverändert zurück."""
    if not (RERANK_ENABLED and len(items) > 1):
        return items
    try:
        import httpx
        r = httpx.post(RERANK_URL, json={
            "model": RERANK_MODEL,
            "query": query,
            "documents": [i.get("content", "") for i in items],
            "top_n": len(items),
        }, timeout=RERANK_TIMEOUT)
        r.raise_for_status()
        results = r.json().get("results") or []
        ranked = []
        for s in sorted(results, key=lambda x: x.get("relevance_score", 0.0), reverse=True):
            idx = s.get("index")
            if isinstance(idx, int) and 0 <= idx < len(items):
                item = dict(items[idx])
                item["rerank_score"] = round(float(s.get("relevance_score", 0.0)), 4)
                ranked.append(item)
        # Nur übernehmen, wenn der Reranker ALLE Kandidaten bewertet hat (sonst Verlustgefahr).
        return ranked if len(ranked) == len(items) else items
    except Exception as exc:  # noqa: BLE001 — Gedächtnis darf am Reranker nie scheitern
        log.debug("rerank übersprungen: %s", exc)
        return items


@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", [])]
        # Zweitstufe: die vorderen Kandidaten nach echter Relevanz ordnen (Rest bleibt hinten dran).
        items = _rerank(q, items[:RERANK_CANDIDATES]) + items[RERANK_CANDIDATES:]
    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 {}


# Deterministischer Guard: mem0-OSS extrahiert trotz custom_instructions hartnäckig Fakten ÜBER
# den Assistenten selbst (Hermes/Lucy) und transiente Zustände (Commit-Rückstand, Cronjob-Läufe,
# Onboarding-/Test-Fragen). Diese Muster werden nach der Extraktion verworfen (der Fakt wird wieder
# gelöscht) — reine Sicherungsschicht, unabhängig vom Extraktions-LLM.
_JUNK_FACT_PATTERNS = [
    r"^\s*(?:der|die|das)\s+assistent",          # "Der Assistent ..." (Selbstbeschreibung)
    r"\bassistent(?:in)?\s+(?:hermes|lucy)\b",   # "... Assistent Hermes/Lucy ..."
    r"\bhermes[- ]agent\b",
    r"commits?\s+hinter",                        # Versions-Rückstand (transient)
    r"\bcronjob\b",
    r"beim\s+letzten\s+(?:lauf|mal)",
    r"onboarding[-\s]?frage|\btest-?frage\b",
    r"hat\s+den\s+assistenten\s+gebeten",
    r"auf\s+betriebstemperatur|hochgefahren",
    r"aktualisiert\s+sein\s+ged(?:ä|ae)chtnis",  # Selbst-Meta über das Gedächtnis
]
_JUNK_FACT_RE = re.compile("|".join(_JUNK_FACT_PATTERNS), re.IGNORECASE)


def _is_junk_fact(text: str) -> bool:
    """True, wenn der extrahierte Fakt Assistenten-Meta oder ein transienter Zustand ist."""
    return bool(_JUNK_FACT_RE.search(text or ""))


@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", [])

    # Guard: frisch extrahierte Assistenten-Meta/transiente Fakten sofort wieder verwerfen.
    kept = []
    for r in results:
        if r.get("event") in ("ADD", "UPDATE") and _is_junk_fact(r.get("memory", "")):
            try:
                mem().delete(r["id"])
            except Exception:
                log.debug("Junk-Fakt-Löschung fehlgeschlagen für %s", r.get("id"))
            log.info("mem0 guard: Junk-Fakt verworfen: %s", (r.get("memory") or "")[:80])
            continue
        kept.append(r)
    results = kept

    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")))
