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

# 5 Kategorien wie im alten System (UI-Kompatibilität). Mem0 selbst kennt keine
# Kategorien — wir führen sie als Metadaten mit.
CATEGORIES = ("user", "instruction", "stable", "versioned", "ephemeral")


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 = "stable"
    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 = "stable"


@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.post("/learn")
def learn(body: LearnIn) -> dict:
    """Auto-Lernen: Gesprächs-Turns/Text durchreichen → Mem0 EXTRAHIERT Fakten selbst
    (infer=True) und entscheidet ADD/UPDATE/NONE gegen das bestehende Gedächtnis."""
    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},
    )
    return {"results": [
        {"id": r.get("id"), "content": r.get("memory"), "event": r.get("event")}
        for r in res.get("results", [])
    ]}


if __name__ == "__main__":
    import uvicorn

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