# -*- coding: utf-8 -*-
# memory_store.py — 结构化长期记忆存储（SQLite + FTS5，零外部依赖）
#
# 事实库: facts(uid, content, importance, keywords, created_at, archived, source)
# 召回:   FTS5 trigram（≥3 字子串）+ LIKE 兜底（短查询/中文双字）
# 迁移:   首次接触某 uid 时把 memory/<uid>.txt 的旧文本记忆导入 facts（meta 表幂等）
import json
import os
import re
import sqlite3
import threading
import time

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
DB_PATH = os.environ.get("CHAT_MEMORY_DB", os.path.join(BASE_DIR, "memory.db"))
LEGACY_DIR = os.environ.get("CHAT_MEMORY_LEGACY_DIR", os.path.join(BASE_DIR, "memory"))

_lock = threading.RLock()


def _connect():
    con = sqlite3.connect(DB_PATH, timeout=30)
    con.execute("PRAGMA journal_mode=WAL")
    return con


def init_db():
    with _lock, _connect() as con:
        con.executescript("""
        CREATE TABLE IF NOT EXISTS facts(
            id INTEGER PRIMARY KEY,
            uid INTEGER NOT NULL,
            content TEXT NOT NULL,
            importance INTEGER NOT NULL DEFAULT 3,
            keywords TEXT NOT NULL DEFAULT '',
            created_at REAL NOT NULL,
            archived INTEGER NOT NULL DEFAULT 0,
            source TEXT NOT NULL DEFAULT ''
        );
        CREATE VIRTUAL TABLE IF NOT EXISTS facts_fts USING fts5(
            content, keywords, content='facts', content_rowid='id', tokenize='trigram'
        );
        CREATE TRIGGER IF NOT EXISTS facts_ai AFTER INSERT ON facts BEGIN
            INSERT INTO facts_fts(rowid, content, keywords) VALUES (new.id, new.content, new.keywords);
        END;
        CREATE TRIGGER IF NOT EXISTS facts_ad AFTER DELETE ON facts BEGIN
            INSERT INTO facts_fts(facts_fts, rowid, content, keywords)
            VALUES ('delete', old.id, old.content, old.keywords);
        END;
        CREATE TRIGGER IF NOT EXISTS facts_au AFTER UPDATE ON facts BEGIN
            INSERT INTO facts_fts(facts_fts, rowid, content, keywords)
            VALUES ('delete', old.id, old.content, old.keywords);
            INSERT INTO facts_fts(rowid, content, keywords) VALUES (new.id, new.content, new.keywords);
        END;
        CREATE TABLE IF NOT EXISTS meta(key TEXT PRIMARY KEY, value TEXT);
        CREATE TABLE IF NOT EXISTS triples(
            id INTEGER PRIMARY KEY,
            uid INTEGER NOT NULL,
            subject TEXT NOT NULL,
            predicate TEXT NOT NULL,
            object TEXT NOT NULL,
            importance INTEGER NOT NULL DEFAULT 3,
            valid_from REAL NOT NULL,
            valid_until REAL,
            created_at REAL NOT NULL,
            source TEXT NOT NULL DEFAULT 'model'
        );
        CREATE INDEX IF NOT EXISTS triples_uid ON triples(uid);
        CREATE INDEX IF NOT EXISTS triples_subj ON triples(uid, subject);
        CREATE TABLE IF NOT EXISTS relations(
            uid INTEGER PRIMARY KEY,
            first_met REAL,
            msg_count INTEGER NOT NULL DEFAULT 0,
            affinity INTEGER NOT NULL DEFAULT 50,
            last_seen REAL,
            notes TEXT NOT NULL DEFAULT '[]'
        );
        """)


def _now():
    return time.time()


def _norm_content(s):
    return re.sub(r"\s+", " ", str(s or "")).strip()


def _norm_keywords(kw):
    if isinstance(kw, str):
        kw = [k.strip() for k in re.split(r"[,，、/ ]+", kw) if k.strip()]
    elif isinstance(kw, (list, tuple)):
        kw = [str(k).strip() for k in kw if str(k).strip()]
    else:
        kw = []
    return ",".join(kw[:8])


def add_facts(uid, facts, source="model"):
    """写入若干条事实；同 uid 完全相同（content 一致）的活跃条目去重跳过。
    返回新增的 id 列表。"""
    init_db()
    uid = int(uid)
    added = []
    with _lock, _connect() as con:
        for f in facts or []:
            if isinstance(f, str):
                f = {"content": f}
            if not isinstance(f, dict):
                continue
            content = _norm_content(f.get("content"))
            if not content or len(content) < 2:
                continue
            importance = f.get("importance", 3)
            try:
                importance = int(importance)
            except (TypeError, ValueError):
                importance = 3
            importance = max(1, min(5, importance))
            keywords = _norm_keywords(f.get("keywords"))
            if con.execute(
                    "SELECT 1 FROM facts WHERE uid=? AND content=? AND archived=0",
                    (uid, content)).fetchone():
                continue
            cur = con.execute(
                "INSERT INTO facts(uid, content, importance, keywords, created_at, archived, source)"
                " VALUES (?,?,?,?,?,0,?)",
                (uid, content, importance, keywords, _now(), source))
            added.append(cur.lastrowid)
    return added


def list_facts(uid, include_archived=False, limit=50):
    init_db()
    q = ("SELECT id, content, importance, keywords, created_at, archived "
         "FROM facts WHERE uid=?")
    if not include_archived:
        q += " AND archived=0"
    q += " ORDER BY created_at DESC LIMIT ?"
    with _lock, _connect() as con:
        return con.execute(q, (int(uid), limit)).fetchall()


def forget_fact(uid, fact_id):
    """归档一条事实（不物理删除，可审计）。返回是否命中。"""
    init_db()
    with _lock, _connect() as con:
        cur = con.execute(
            "UPDATE facts SET archived=1 WHERE uid=? AND id=? AND archived=0",
            (int(uid), int(fact_id)))
        return cur.rowcount > 0


def count_facts(uid, active_only=True):
    init_db()
    q = "SELECT COUNT(*) FROM facts WHERE uid=?"
    if active_only:
        q += " AND archived=0"
    with _lock, _connect() as con:
        return con.execute(q, (int(uid),)).fetchone()[0]


def search_facts(uid, query, limit=5):
    """FTS5 trigram 召回（≥3 字）；短查询/无命中时 LIKE 兜底。"""
    init_db()
    uid = int(uid)
    query = _norm_content(query)
    if not query:
        return []
    rows = []
    with _lock, _connect() as con:
        if len(query) >= 3:
            try:
                rows = con.execute(
                    "SELECT f.id, f.content, f.importance, f.created_at "
                    "FROM facts_fts JOIN facts f ON f.id = facts_fts.rowid "
                    "WHERE facts_fts MATCH ? AND f.uid=? AND f.archived=0 "
                    "ORDER BY rank, f.created_at DESC LIMIT ?",
                    (query, uid, limit)).fetchall()
            except sqlite3.OperationalError:
                rows = []
        if not rows:
            like = f"%{query}%"
            rows = con.execute(
                "SELECT id, content, importance, created_at FROM facts "
                "WHERE uid=? AND archived=0 AND (content LIKE ? OR keywords LIKE ?) "
                "ORDER BY created_at DESC LIMIT ?",
                (uid, like, like, limit)).fetchall()
    return rows


def touch_relation(uid):
    """每轮对话调用：自动登记/更新关系统计（相识时间、对话次数、最近活跃）。
    每 50 轮顺手做一次程序式整理（归档低重要度老记忆，不花 token）。"""
    init_db()
    uid = int(uid)
    now = _now()
    with _lock, _connect() as con:
        row = con.execute("SELECT msg_count FROM relations WHERE uid=?", (uid,)).fetchone()
        if row:
            new_count = row[0] + 1
            con.execute("UPDATE relations SET msg_count=?, last_seen=? WHERE uid=?",
                        (new_count, now, uid))
        else:
            new_count = 1
            con.execute(
                "INSERT INTO relations(uid, first_met, msg_count, affinity, last_seen, notes)"
                " VALUES (?,?,1,50,?,'[]')", (uid, now, now))
    if new_count % 50 == 0:
        maybe_archive_stale(uid)


def get_relation(uid):
    init_db()
    with _lock, _connect() as con:
        row = con.execute(
            "SELECT first_met, msg_count, affinity, last_seen, notes FROM relations WHERE uid=?",
            (int(uid),)).fetchone()
    if not row:
        return None
    try:
        notes = json.loads(row[4] or "[]")
    except Exception:
        notes = []
    return {"first_met": row[0], "msg_count": row[1], "affinity": row[2],
            "last_seen": row[3], "notes": notes}


def add_relation_note(uid, note):
    """模型写的「关系上的重要变化/约定/里程碑」（一句话）。notes 保留最近 5 条。"""
    init_db()
    uid = int(uid)
    note = _norm_content(note)
    if not note:
        return
    with _lock, _connect() as con:
        row = con.execute("SELECT notes FROM relations WHERE uid=?", (uid,)).fetchone()
        if not row:
            touch_relation(uid)
            row = ("[]",)
        try:
            notes = json.loads(row[0] or "[]")
        except Exception:
            notes = []
        notes.append({"ts": _now(), "text": note[:200]})
        notes = notes[-5:]
        con.execute("UPDATE relations SET notes=? WHERE uid=?",
                    (json.dumps(notes, ensure_ascii=False), uid))


def maybe_archive_stale(uid, older_than_days=30, max_importance=2):
    """轻量整理：归档「老且不重要」的记忆（程序式，不花 token）。返回归档条数。"""
    init_db()
    cutoff = _now() - older_than_days * 86400
    with _lock, _connect() as con:
        cur = con.execute(
            "UPDATE facts SET archived=1 WHERE uid=? AND archived=0 AND importance<=? AND created_at<?",
            (int(uid), max_importance, cutoff))
        return cur.rowcount


def add_links(uid, links, source="model"):
    """写入 SPO 三元组。同 (uid, subject, predicate) 的旧边若 object 不同，
    置 valid_until（时效窗：旧事实保留但不再生效）。返回新增 id 列表。"""
    init_db()
    uid = int(uid)
    now = _now()
    added = []
    with _lock, _connect() as con:
        for lk in links or []:
            if not isinstance(lk, dict):
                continue
            s = _norm_content(lk.get("subject"))
            pr = _norm_content(lk.get("predicate"))
            o = _norm_content(lk.get("object"))
            if not s or not pr or not o:
                continue
            try:
                importance = max(1, min(5, int(lk.get("importance", 3))))
            except (TypeError, ValueError):
                importance = 3
            # 同主同谓的旧边：值相同则跳过；值不同则置时效窗
            old = con.execute(
                "SELECT id, object FROM triples WHERE uid=? AND subject=? AND predicate=?"
                " AND valid_until IS NULL", (uid, s, pr)).fetchall()
            if any(row[1] == o for row in old):
                continue
            for row in old:
                con.execute("UPDATE triples SET valid_until=? WHERE id=?", (now, row[0]))
            cur = con.execute(
                "INSERT INTO triples(uid, subject, predicate, object, importance,"
                " valid_from, valid_until, created_at, source) VALUES (?,?,?,?,?,?,NULL,?,?)",
                (uid, s, pr, o, importance, now, now, source))
            added.append(cur.lastrowid)
    return added


def entity_card(uid, name, limit=10):
    """某实体的卡片：以它为主语或宾语的当前有效三元组。"""
    init_db()
    uid = int(uid)
    name = _norm_content(name)
    if not name:
        return []
    with _lock, _connect() as con:
        return con.execute(
            "SELECT subject, predicate, object, importance FROM triples "
            "WHERE uid=? AND valid_until IS NULL AND (subject=? OR object=?) "
            "ORDER BY importance DESC, created_at DESC LIMIT ?",
            (uid, name, name, limit)).fetchall()


def list_entities(uid, limit=50):
    """全部实体（主语/宾语出现过的名字）及各自有效边数。"""
    init_db()
    with _lock, _connect() as con:
        rows = con.execute(
            "SELECT name, COUNT(*) FROM ("
            "  SELECT subject AS name FROM triples WHERE uid=? AND valid_until IS NULL"
            "  UNION ALL SELECT object FROM triples WHERE uid=? AND valid_until IS NULL"
            ") GROUP BY name ORDER BY COUNT(*) DESC LIMIT ?",
            (int(uid), int(uid), limit)).fetchall()
    return rows


def recall_facts(uid, query, top_k=5, pinned_k=3):
    """召回式读取：高重要度置顶 pinned_k 条 + 按当前输入 FTS5 召回 top_k 条，按 id 去重。
    返回 [(id, content, importance), ...]"""
    init_db()
    uid = int(uid)
    with _lock, _connect() as con:
        pinned = con.execute(
            "SELECT id, content, importance FROM facts WHERE uid=? AND archived=0 "
            "ORDER BY importance DESC, created_at DESC LIMIT ?",
            (uid, pinned_k)).fetchall()
    seen = {r[0] for r in pinned}
    out = list(pinned)
    for r in search_facts(uid, query, limit=top_k):
        if r[0] not in seen:
            out.append((r[0], r[1], r[2]))
            seen.add(r[0])
    return out


def migrate_legacy(uid, legacy_dir=None):
    """首次接触某 uid 时把 memory/<uid>.txt 的旧文本记忆导入 facts（meta 幂等）。"""
    init_db()
    uid = int(uid)
    legacy_dir = legacy_dir or LEGACY_DIR
    marker = f"imported:{uid}"
    with _lock, _connect() as con:
        if con.execute("SELECT 1 FROM meta WHERE key=?", (marker,)).fetchone():
            return 0
    path = os.path.join(legacy_dir, f"{uid}.txt")
    if not os.path.exists(path):
        return 0
    facts = []
    with open(path, encoding="utf-8") as f:
        for ln in f:
            ln = ln.strip()
            if not ln:
                continue
            m = re.match(r"^\[\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}\]\s*(.*)$", ln)
            facts.append({"content": m.group(1) if m else ln,
                          "importance": 3, "keywords": []})
    added = len(add_facts(uid, facts, source="legacy"))
    with _lock, _connect() as con:
        con.execute("INSERT OR REPLACE INTO meta(key, value) VALUES (?,?)",
                    (marker, str(_now())))
    return added