宜博 / pixelphone

"""
Memory Graph — SQLite-backed knowledge graph for the Agent Phone Runtime.

Schema:
  person    — unified identity across apps (WeChat, Telegram, contacts)
  message   — messages linked to persons
  entity    — general facts/events (calendar, notes, etc.)
  relation  — named edges between persons

Vector search is handled by chromadb (optional, falls back to FTS5 keyword search).
"""
from __future__ import annotations

import json
import os
import sqlite3
import time
from pathlib import Path
from typing import Any


def _load_hermes_env():
    """Load ~/.hermes/.env so KIMI_API_KEY etc. are always available."""
    # On Android as root via ADB, HOME may be the Mac host path.
    # Always prefer the Termux home if it exists.
    termux_home = Path("/data/data/com.termux/files/home")
    home = termux_home if termux_home.exists() else Path(os.environ.get("HOME", str(termux_home)))
    env_file = home / ".hermes" / ".env"
    if not env_file.exists():
        return
    for raw in env_file.read_text(encoding="utf-8").splitlines():
        line = raw.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        k, _, v = line.partition("=")
        k = k.strip()
        v = v.strip().strip('"').strip("'")
        if k and not os.environ.get(k):
            os.environ[k] = v

_load_hermes_env()

_TERMUX_HOME = Path("/data/data/com.termux/files/home")
_HERMES_HOME = (_TERMUX_HOME if _TERMUX_HOME.exists()
                else Path(os.environ.get("HOME", str(_TERMUX_HOME)))) / ".hermes"

DB_PATH = os.environ.get("MEMORY_GRAPH_DB", str(_HERMES_HOME / "memory_graph.db"))


def _get_db() -> sqlite3.Connection:
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("PRAGMA foreign_keys=ON")
    return conn


def init_db():
    conn = _get_db()
    conn.executescript("""
    CREATE TABLE IF NOT EXISTS person (
        id          INTEGER PRIMARY KEY AUTOINCREMENT,
        wxid        TEXT,
        tg_id       TEXT,
        phone       TEXT,
        nickname    TEXT,
        remark      TEXT,
        source_app  TEXT,
        created_at  INTEGER DEFAULT (strftime('%s','now')),
        updated_at  INTEGER DEFAULT (strftime('%s','now')),
        UNIQUE(wxid),
        UNIQUE(tg_id)
    );

    CREATE TABLE IF NOT EXISTS message (
        id          INTEGER PRIMARY KEY AUTOINCREMENT,
        person_id   INTEGER REFERENCES person(id),
        talker      TEXT,
        content     TEXT,
        is_sent     INTEGER DEFAULT 0,
        app         TEXT DEFAULT 'wechat',
        ts          INTEGER,
        raw_id      TEXT,
        created_at  INTEGER DEFAULT (strftime('%s','now'))
    );

    CREATE TABLE IF NOT EXISTS entity (
        id          INTEGER PRIMARY KEY AUTOINCREMENT,
        type        TEXT,
        title       TEXT,
        content     TEXT,
        source_app  TEXT,
        ts          INTEGER,
        metadata    TEXT,
        created_at  INTEGER DEFAULT (strftime('%s','now'))
    );

    CREATE TABLE IF NOT EXISTS relation (
        id          INTEGER PRIMARY KEY AUTOINCREMENT,
        from_id     INTEGER REFERENCES person(id),
        to_id       INTEGER REFERENCES person(id),
        label       TEXT,
        weight      REAL DEFAULT 1.0,
        created_at  INTEGER DEFAULT (strftime('%s','now'))
    );

    CREATE VIRTUAL TABLE IF NOT EXISTS message_fts USING fts5(
        content, talker, app,
        content='message', content_rowid='id'
    );

    CREATE TRIGGER IF NOT EXISTS message_ai AFTER INSERT ON message BEGIN
        INSERT INTO message_fts(rowid, content, talker, app)
        VALUES (new.id, new.content, new.talker, new.app);
    END;

    CREATE INDEX IF NOT EXISTS idx_message_person ON message(person_id);
    CREATE INDEX IF NOT EXISTS idx_message_ts ON message(ts);
    CREATE INDEX IF NOT EXISTS idx_message_talker ON message(talker);
    CREATE INDEX IF NOT EXISTS idx_entity_type ON entity(type);
    CREATE INDEX IF NOT EXISTS idx_entity_ts ON entity(ts);
    """)
    conn.commit()
    conn.close()


def upsert_person(wxid: str = None, tg_id: str = None, nickname: str = None,
                   remark: str = None, phone: str = None, source_app: str = "wechat") -> int:
    conn = _get_db()
    now = int(time.time())
    try:
        if wxid:
            row = conn.execute("SELECT id FROM person WHERE wxid=?", (wxid,)).fetchone()
            if row:
                conn.execute(
                    "UPDATE person SET nickname=COALESCE(?,nickname), remark=COALESCE(?,remark), "
                    "phone=COALESCE(?,phone), updated_at=? WHERE id=?",
                    (nickname, remark, phone, now, row["id"])
                )
                conn.commit()
                return row["id"]
        cursor = conn.execute(
            "INSERT OR IGNORE INTO person(wxid, tg_id, phone, nickname, remark, source_app, created_at, updated_at) "
            "VALUES(?,?,?,?,?,?,?,?)",
            (wxid, tg_id, phone, nickname, remark, source_app, now, now)
        )
        conn.commit()
        pid = cursor.lastrowid
        if pid == 0:
            pid = conn.execute("SELECT id FROM person WHERE wxid=?", (wxid,)).fetchone()["id"]
        return pid
    finally:
        conn.close()


def ingest_wechat_contacts(contacts: list[dict]) -> dict:
    added = 0
    updated = 0
    for c in contacts:
        wxid = c.get("username", "")
        if not wxid:
            continue
        nickname = c.get("nickname") or c.get("conRemark") or wxid
        remark = c.get("conRemark", "")
        existing = _person_exists(wxid=wxid)
        upsert_person(wxid=wxid, nickname=nickname, remark=remark, source_app="wechat")
        if existing:
            updated += 1
        else:
            added += 1
    return {"added": added, "updated": updated, "total": len(contacts)}


def ingest_wechat_messages(messages: list[dict]) -> dict:
    conn = _get_db()
    added = 0
    # Pre-build wxid→person_id cache from existing persons to avoid nested connections
    person_cache: dict[str, int] = {}
    for row in conn.execute("SELECT id, wxid FROM person WHERE wxid IS NOT NULL"):
        person_cache[row["wxid"]] = row["id"]
    try:
        for m in messages:
            talker = m.get("talker", "")
            content = m.get("content", "")
            ts = m.get("createTime", int(time.time() * 1000)) // 1000
            raw_id = str(m.get("msgId", ""))
            is_sent = m.get("isSend", 0)
            if not talker or not content:
                continue
            # Find or create person (using same connection)
            person_id = None
            if not talker.endswith("@chatroom"):
                if talker in person_cache:
                    person_id = person_cache[talker]
                else:
                    now = int(time.time())
                    cursor = conn.execute(
                        "INSERT OR IGNORE INTO person(wxid, nickname, source_app, created_at, updated_at) "
                        "VALUES(?,?,?,?,?)",
                        (talker, talker, "wechat", now, now)
                    )
                    if cursor.lastrowid:
                        person_id = cursor.lastrowid
                    else:
                        row = conn.execute("SELECT id FROM person WHERE wxid=?", (talker,)).fetchone()
                        person_id = row["id"] if row else None
                    if person_id:
                        person_cache[talker] = person_id
            # Dedup by raw_id
            if not conn.execute("SELECT 1 FROM message WHERE raw_id=? AND app='wechat'", (raw_id,)).fetchone():
                conn.execute(
                    "INSERT INTO message(person_id, talker, content, is_sent, app, ts, raw_id) "
                    "VALUES(?,?,?,?,?,?,?)",
                    (person_id, talker, content, is_sent, "wechat", ts, raw_id)
                )
                added += 1
        conn.commit()
        return {"added": added, "total": len(messages)}
    finally:
        conn.close()


def _person_exists(wxid: str = None, tg_id: str = None) -> bool:
    conn = _get_db()
    try:
        if wxid:
            return conn.execute("SELECT 1 FROM person WHERE wxid=?", (wxid,)).fetchone() is not None
        if tg_id:
            return conn.execute("SELECT 1 FROM person WHERE tg_id=?", (tg_id,)).fetchone() is not None
        return False
    finally:
        conn.close()


def search_memory(query: str, person: str = "", since: int = 0,
                   until: int = 0, source_app: str = "", limit: int = 20) -> list[dict]:
    conn = _get_db()
    try:
        results = []

        # FTS5 keyword search
        fts_query = query.replace('"', '""')
        sql = """
            SELECT m.id, m.talker, m.content, m.is_sent, m.app, m.ts,
                   p.nickname, p.remark
            FROM message_fts fts
            JOIN message m ON fts.rowid = m.id
            LEFT JOIN person p ON m.person_id = p.id
            WHERE message_fts MATCH ?
        """
        params: list[Any] = [fts_query]

        if person:
            sql += " AND (m.talker LIKE ? OR p.nickname LIKE ? OR p.remark LIKE ?)"
            params += [f"%{person}%", f"%{person}%", f"%{person}%"]
        if since:
            sql += " AND m.ts >= ?"
            params.append(since)
        if until:
            sql += " AND m.ts <= ?"
            params.append(until)
        if source_app:
            sql += " AND m.app = ?"
            params.append(source_app)
        sql += f" ORDER BY m.ts DESC LIMIT {limit}"

        try:
            rows = conn.execute(sql, params).fetchall()
            for r in rows:
                results.append(dict(r))
        except sqlite3.OperationalError:
            # FTS query syntax error — fall back to LIKE
            like_sql = """
                SELECT m.id, m.talker, m.content, m.is_sent, m.app, m.ts,
                       p.nickname, p.remark
                FROM message m
                LEFT JOIN person p ON m.person_id = p.id
                WHERE m.content LIKE ?
            """
            like_params: list[Any] = [f"%{query}%"]
            if person:
                like_sql += " AND (m.talker LIKE ? OR p.nickname LIKE ?)"
                like_params += [f"%{person}%", f"%{person}%"]
            if since:
                like_sql += " AND m.ts >= ?"
                like_params.append(since)
            like_sql += f" ORDER BY m.ts DESC LIMIT {limit}"
            rows = conn.execute(like_sql, like_params).fetchall()
            for r in rows:
                results.append(dict(r))

        return results
    finally:
        conn.close()


def get_stats() -> dict:
    conn = _get_db()
    try:
        persons = conn.execute("SELECT COUNT(*) FROM person").fetchone()[0]
        messages = conn.execute("SELECT COUNT(*) FROM message").fetchone()[0]
        entities = conn.execute("SELECT COUNT(*) FROM entity").fetchone()[0]
        vectors = 0
        try:
            vectors = conn.execute("SELECT COUNT(*) FROM message_vector").fetchone()[0]
        except Exception:
            pass
        return {"persons": persons, "messages": messages, "entities": entities, "vectors": vectors}
    finally:
        conn.close()


# ── Vector Memory (Sprint 9) ───────────────────────────────────────────────────

KIMI_API_KEY = os.environ.get("KIMI_API_KEY", "")
KIMI_BASE_URL = os.environ.get("KIMI_BASE_URL", "https://api.kimi.com/coding/v1")
EMBED_MODEL = "moonshot-v1-embedding"
EMBED_DIM = 1024


def _init_vector_table():
    conn = _get_db()
    conn.executescript("""
    CREATE TABLE IF NOT EXISTS message_vector (
        message_id  INTEGER PRIMARY KEY REFERENCES message(id),
        embedding   BLOB
    );
    """)
    conn.commit()
    conn.close()


def _get_embedding(text: str) -> list[float] | None:
    """Call Kimi embedding API and return 1024-dim vector."""
    if not KIMI_API_KEY:
        return None
    import urllib.request
    text = text[:2000]  # Kimi max input per item
    payload = json.dumps({"model": EMBED_MODEL, "input": [text]}).encode()
    req = urllib.request.Request(
        f"{KIMI_BASE_URL}/embeddings",
        data=payload,
        headers={
            "Authorization": f"Bearer {KIMI_API_KEY}",
            "Content-Type": "application/json",
        },
    )
    try:
        import urllib.error
        with urllib.request.urlopen(req, timeout=15) as resp:
            data = json.loads(resp.read().decode())
            return data["data"][0]["embedding"]
    except Exception:
        return None


def _embed_to_blob(vec: list[float]) -> bytes:
    import struct
    return struct.pack(f"{len(vec)}f", *vec)


def _blob_to_embed(blob: bytes) -> list[float]:
    import struct
    n = len(blob) // 4
    return list(struct.unpack(f"{n}f", blob))


def _cosine(a: list[float], b: list[float]) -> float:
    dot = sum(x * y for x, y in zip(a, b))
    na = sum(x * x for x in a) ** 0.5
    nb = sum(x * x for x in b) ** 0.5
    return dot / (na * nb + 1e-9)


def embed_messages(limit: int = 200) -> dict:
    """Generate and store embeddings for un-embedded messages."""
    if not KIMI_API_KEY:
        return {"error": "KIMI_API_KEY not set"}
    _init_vector_table()
    conn = _get_db()
    try:
        # Find messages without embeddings
        rows = conn.execute(
            "SELECT m.id, m.content FROM message m "
            "LEFT JOIN message_vector mv ON m.id = mv.message_id "
            "WHERE mv.message_id IS NULL AND m.content != '' "
            f"LIMIT {limit}"
        ).fetchall()
        embedded = 0
        for row in rows:
            mid, content = row["id"], row["content"]
            vec = _get_embedding(content)
            if vec:
                blob = _embed_to_blob(vec)
                conn.execute(
                    "INSERT OR REPLACE INTO message_vector(message_id, embedding) VALUES(?,?)",
                    (mid, blob)
                )
                embedded += 1
        conn.commit()
        return {"embedded": embedded, "pending": len(rows) - embedded}
    finally:
        conn.close()


def semantic_search(query: str, limit: int = 10) -> list[dict]:
    """Semantic similarity search using Kimi embeddings."""
    if not KIMI_API_KEY:
        return []
    _init_vector_table()
    query_vec = _get_embedding(query)
    if not query_vec:
        return []
    conn = _get_db()
    try:
        rows = conn.execute(
            "SELECT mv.message_id, mv.embedding, m.talker, m.content, m.ts, "
            "p.nickname, p.remark "
            "FROM message_vector mv "
            "JOIN message m ON mv.message_id = m.id "
            "LEFT JOIN person p ON m.person_id = p.id"
        ).fetchall()
        scored = []
        for r in rows:
            vec = _blob_to_embed(r["embedding"])
            score = _cosine(query_vec, vec)
            scored.append({
                "id": r["message_id"],
                "talker": r["talker"],
                "content": r["content"],
                "ts": r["ts"],
                "nickname": r["nickname"],
                "remark": r["remark"],
                "score": round(score, 4),
            })
        scored.sort(key=lambda x: x["score"], reverse=True)
        return scored[:limit]
    finally:
        conn.close()


if __name__ == "__main__":
    import os as _os
    # Load .env if available
    env_file = _os.path.expanduser("~/.hermes/.env")
    if _os.path.exists(env_file):
        for line in open(env_file):
            line = line.strip()
            if line and not line.startswith("#") and "=" in line:
                k, v = line.split("=", 1)
                _os.environ.setdefault(k.strip(), v.strip())
    # Reload key after env load
    import memory_graph as _self
    _self.KIMI_API_KEY = _os.environ.get("KIMI_API_KEY", "")

    init_db()
    print("Memory Graph initialized:", DB_PATH)
    print("Stats:", get_stats())