#!/usr/bin/env python3
"""
统一记忆入口 (memory_recall.py)

同时查询织忆、TencentDB、Soulful 三套记忆系统，合并去重后输出统一上下文。
"""

import argparse
import json
import os
import subprocess
import threading
import time
import urllib.error
import urllib.request
from pathlib import Path

# ── 常量 ──────────────────────────────────────────────────────────────────────

ZHIYI_URL = "http://localhost:7821/api/v1/recall"
ZHIYI_HEADERS = {
    "X-API-Key": os.environ.get("ZHIYI_API_KEY", "zhiyi-dev-key-2026"),
    "Content-Type": "application/json",
}
TENCENTDB_URL = "http://localhost:8420"
SOULFUL_DIR = Path.home() / ".hermes" / "soulful"

DEFAULT_TOP_K = 3
TIMEOUT_SEC = 15

# ── 工具函数 ──────────────────────────────────────────────────────────────────

def http_post(url: str, payload: dict, headers: dict | None = None, timeout: int = TIMEOUT_SEC) -> dict | None:
    """POST JSON，返回解析后的 dict；失败自动重试最多3次，返回 None 并打印 warning。"""
    delays = [1, 2, 4]
    for attempt in range(4):  # 首次 + 3次重试
        try:
            req = urllib.request.Request(
                url,
                data=json.dumps(payload).encode("utf-8"),
                headers=headers or {},
                method="POST",
            )
            with urllib.request.urlopen(req, timeout=timeout) as resp:
                return json.loads(resp.read().decode("utf-8"))
        except Exception as e:
            if attempt < 3:
                print(f"  [warning] {url} -> {e} (重试 {attempt+1}/3, 等待 {delays[attempt]}s)", flush=True)
                time.sleep(delays[attempt])
            else:
                print(f"  [warning] {url} -> {e} (已重试3次，放弃)", flush=True)
                return None
    return None


def read_json_file(path: Path) -> dict | list | None:
    """读取 JSON 文件，失败返回 None 并打印 warning。"""
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except Exception as e:
        print(f"  [warning] read {path} -> {e}", flush=True)
        return None


def read_jsonl(path: Path):
    """逐行读取 JSONL，yield dict；失败打印 warning 并跳过。"""
    try:
        for line in path.read_text(encoding="utf-8").splitlines():
            line = line.strip()
            if not line:
                continue
            try:
                yield json.loads(line)
            except json.JSONDecodeError:
                pass
    except Exception as e:
        print(f"  [warning] read {path} -> {e}", flush=True)


# ── bge-embed keepalive ────────────────────────────────────────────────────────

BGE_HEALTH_URL = "http://localhost:8000/health"
_KEEPALIVE_RUNNING = True


def add_keepalive():
    """后台线程：每10秒ping bge-embed /health，防止服务因空闲超时断开。"""
    global _KEEPALIVE_RUNNING
    while _KEEPALIVE_RUNNING:
        try:
            req = urllib.request.Request(BGE_HEALTH_URL, method="GET")
            with urllib.request.urlopen(req, timeout=5) as resp:
                resp.read()  # 消费响应，无需解析
        except Exception:
            pass  # bge-embed 暂时不可用也静默，ping 下次继续
        time.sleep(10)


def start_keepalive():
    """启动 keepalive 守护线程（在 main 入口调用）。"""
    t = threading.Thread(target=add_keepalive, daemon=True)
    t.start()


def stop_keepalive():
    """停止 keepalive 线程。"""
    global _KEEPALIVE_RUNNING
    _KEEPALIVE_RUNNING = False


# ── 查询函数（供线程调用） ──────────────────────────────────────────────────────

def query_zhiyi(query: str, top_k: int) -> dict:
    result = http_post(ZHIYI_URL, {"query": query, "top_k": top_k}, ZHIYI_HEADERS)
    if result is None:
        # 降级：织忆不可用时从 llm_context.json 读取缓存上下文
        ctx = load_llm_context()
        if ctx is not None:
            print(f"  [degrade] 织忆不可用，从 llm_context.json 读取缓存上下文", flush=True)
            return {
                "error": False,
                "degraded": True,
                "results": [],
                "llm_context": ctx,
            }
        return {"error": True, "results": []}
    return {"error": False, "results": result.get("results", [])}


def query_tencentdb_l1(query: str, top_k: int) -> dict:
    result = http_post(
        f"{TENCENTDB_URL}/search/memories",
        {"query": query, "top_k": top_k},
    )
    if result is None:
        return {"error": True, "results": []}
    return {"error": False, "results": result.get("results", [])}


def query_tencentdb_l0(query: str, top_k: int) -> dict:
    result = http_post(
        f"{TENCENTDB_URL}/search/conversations",
        {"query": query, "top_k": top_k},
    )
    if result is None:
        return {"error": True, "results": []}
    return {"error": False, "results": result.get("results", [])}


# ── Soulful 读取 ───────────────────────────────────────────────────────────────

def load_soulful() -> dict:
    """加载 Soulful 三件套，返回聚合 dict；各文件缺失不报错。"""
    profile = read_json_file(SOULFUL_DIR / "user-profile.json") or {}
    cares_data = read_json_file(SOULFUL_DIR / "cares-queue.json") or {}
    heart_traces = list(read_jsonl(SOULFUL_DIR / "heart-traces.jsonl"))

    return {
        "error": False,
        "profile": profile,
        "cares": cares_data.get("cares", []) if isinstance(cares_data, dict) else [],
        "heart_traces": heart_traces,
    }


# ── 统一读取入口（优先 llm_context.json）────────────────────────────────────────

LLM_CONTEXT_FILE = Path.home() / ".hermes" / "llm_context.json"


def load_llm_context() -> dict | None:
    """读取 llm_context.json（daemon 统一写入的合并上下文）。失败返回 None。"""
    try:
        if not LLM_CONTEXT_FILE.exists():
            return None
        return json.loads(LLM_CONTEXT_FILE.read_text(encoding="utf-8"))
    except Exception:
        return None


def query_llm_context() -> dict:
    """主查询：优先 llm_context.json，再补充织忆/L1/L0/soulful 详细数据。

    返回格式同 parallel_query：(zhiyi, l1, l0, soulful, unified_context)
    unified_context 包含 user_profile / active_scenes / short_term / cares / recent_moments
    """
    ctx = load_llm_context()
    soulful_data = load_soulful()

    unified = {
        "error": False,
        "ctx": ctx,  # llm_context v2 dict or None
        "user_profile": (ctx or {}).get("user_profile", {}),
        "active_scenes": (ctx or {}).get("active_scenes", []),
        "short_term": (ctx or {}).get("short_term", {}),
        "cares": soulful_data.get("cares", []),  # 优先 soulful 原始
        "heart_traces": soulful_data.get("heart_traces", []),
        "distill_status": (ctx or {}).get("distill_status", "no_context"),
    }

    # 如果 ctx 有完整数据，织忆语义检索由主查询补充
    # （此处不重复检索，llm_context.json 已含 daemon 蒸馏结果）
    zhiyi = {"error": False, "results": []}
    tddb_l1 = {"error": False, "results": []}
    tddb_l0 = {"error": False, "results": []}

    return zhiyi, tddb_l1, tddb_l0, soulful_data, unified


# ── 并行查询（fallback 模式）────────────────────────────────────────────────────

def parallel_query(query: str, top_k: int):
    """用三个线程同时查询三个后端 + Soulful 文件，返回 (zhiyi, l1, l0, soulful)。"""

    def safe_update(result_dict, fn, *args, **kwargs):
        """安全执行查询，异常时标记 error"""
        try:
            result_dict.update(fn(*args, **kwargs))
        except Exception as e:
            result_dict["error"] = True
            result_dict["results"] = []

    threads = []
    zhiyi_result = {}
    l1_result = {}
    l0_result = {}
    soulful_result = {}

    t1 = threading.Thread(target=lambda: safe_update(zhiyi_result, query_zhiyi, query, top_k))
    t2 = threading.Thread(target=lambda: safe_update(l1_result, query_tencentdb_l1, query, top_k))
    t3 = threading.Thread(target=lambda: safe_update(l0_result, query_tencentdb_l0, query, top_k))
    t4 = threading.Thread(target=lambda: safe_update(soulful_result, load_soulful))

    for t in (t1, t2, t3, t4):
        t.start()
    for t in (t1, t2, t3, t4):
        t.join()

    return zhiyi_result, l1_result, l0_result, soulful_result


# ── 格式化输出 ────────────────────────────────────────────────────────────────

import re

def _fmt_val(v):
    """格式化 profile 中的 value：dict/list 走 JSON.dumps，其余直接转 str。"""
    if isinstance(v, (dict, list)):
        return json.dumps(v, ensure_ascii=False)
    return str(v)

def format_pretty(zhiyi, l1, l0, soulful):
    """pretty 模式：分区块显示。"""
    lines = []

    # ── 语义记忆（织忆） ──
    if zhiyi.get("error"):
        lines.append("## 语义记忆（织忆）")
        lines.append("暂不可用")
    else:
        lines.append("## 语义记忆（织忆）")
        for item in zhiyi.get("results", []):
            if isinstance(item, dict):
                content = item.get("content", "")
                score = item.get("score", 0.0)
                cat = item.get("category", "")
                ts = item.get("timestamp", "")
            else:
                content = str(item)
                score, cat, ts = 0.0, "", ""
            cat_str = f" [{cat}]" if cat else ""
            ts_str = f" ({ts})" if ts else ""
            lines.append(f"[{content}]{cat_str}{ts_str} (score: {score:.2f})")
        lines.append("")

    # ── 人格记忆（TencentDB L1） ──
    if l1.get("error"):
        lines.append("## 人格记忆（TencentDB）")
        lines.append("暂不可用\n")
    else:
        lines.append("## 人格记忆（TencentDB）")
        results = l1.get("results", [])
        # TencentDB 返回的 results 是 markdown 字符串，不是 JSON 数组
        if isinstance(results, str):
            # 整段输出，不拆分
            lines.append(results)
        else:
            for item in results:
                if isinstance(item, dict):
                    content = item.get("content", "")
                    score = item.get("score", 0.0)
                    ptype = item.get("type", "instruction")
                    scene = item.get("scene", "")
                else:
                    content = str(item)
                    score, ptype, scene = 0.0, "instruction", ""
                scene_str = f" [{scene}]" if scene else ""
                lines.append(f"[type: {ptype}]{scene_str}")
                lines.append(f"[{content}] (score: {score:.2f})")
        lines.append("")

    # ── 对话记忆（TencentDB L0） ──
    if l0.get("error"):
        lines.append("## 对话记忆（TencentDB L0）")
        lines.append("暂不可用\n")
    else:
        lines.append("## 对话记忆（TencentDB L0）")
        results = l0.get("results", [])
        # TencentDB 返回的 results 是 markdown 字符串，不是 JSON 数组
        if isinstance(results, str):
            lines.append(results)
        else:
            for item in results:
                if isinstance(item, dict):
                    msg = item.get("message", "")
                    score = item.get("score", 0.0)
                    ts = item.get("timestamp", "")
                    session = item.get("session", "")
                else:
                    msg = str(item)
                    score, ts, session = 0.0, "", ""
                session_str = f"[{session}]" if session else ""
                ts_str = f" {ts}" if ts else ""
                lines.append(f"{session_str}{ts_str}")
                lines.append(f"[{msg}] (score: {score:.2f})")
        lines.append("")

    # ── 关系感知（Soulful） ──
    lines.append("## 关系感知（Soulful）")

    if soulful.get("error"):
        lines.append("暂不可用\n")
    else:
        # 画像
        lines.append("### 画像")
        profile = soulful.get("profile", {})
        if profile:
            for key, value in profile.items():
                lines.append(f"{key}: {_fmt_val(value)}")
        else:
            lines.append("(无)")

        # 牵挂
        lines.append("\n### 待关注（牵挂）")
        cares = soulful.get("cares", [])
        if cares:
            for item in cares:
                content = item.get("content", "")
                due = item.get("follow_up_date") or item.get("due") or "null"
                lines.append(f"- [{content}] (due: {due})")
        else:
            lines.append("(无)")

        # 心迹
        lines.append("\n### 心迹")
        heart_traces = soulful.get("heart_traces", [])
        if heart_traces:
            for item in heart_traces:
                ts = item.get("timestamp", "")
                content = item.get("content", "")
                lines.append(f"[{ts}] [{content}]")
        else:
            lines.append("(无)")

    return "\n".join(lines)


def format_compact(zhiyi, l1, l0, soulful):
    """compact 模式：单行动态摘要。"""
    parts = []

    # 织忆
    if not zhiyi.get("error"):
        results = zhiyi.get("results", [])
        if results:
            top = results[0]
            parts.append(f"织忆/{top.get('content', '')[:40]}(s={top.get('score',0):.2f})")

    # L1 - TencentDB returns string for markdown, list for parsed
    if not l1.get("error"):
        results = l1.get("results", [])
        if isinstance(results, str):
            # Markdown string - extract content from first meaningful line
            # L1 format: "- **[persona/instruction]** ..." then blank then "  content"
            #            or "  content" appears on 2nd or 3rd non-empty line
            lines = [l for l in results.split("\n") if l.strip() and not l.startswith("Found") and not l.startswith("---")]
            if len(lines) >= 2:
                # Second line has the actual content text (after blank line)
                content = lines[1].strip().lstrip("-* []").strip()
                parts.append(f"人格:{content[:35]}...")
            elif lines:
                content = lines[0].strip().lstrip("-* []").strip()
                parts.append(f"人格:{content[:35]}")
        elif results:
            top = results[0]
            ptype = top.get("type", "?")
            parts.append(f"人格/{ptype}:{top.get('content', '')[:30]}(s={top.get('score',0):.2f})")

    # L0 - TencentDB returns string for markdown, list for parsed
    if not l0.get("error"):
        results = l0.get("results", [])
        if isinstance(results, str):
            # L0 format: "**[user]** Session: ... (score: X)\n\nMESSAGE CONTENT\n\n---..."
            # Match header line + blank line + first non-empty content line
            m = re.search(r'\*\*\[(user|assistant)\]\*\*.*?\n\n(.+?)(?:\n\n|---\Z)', results, re.DOTALL)
            if m:
                msg = m.group(2).strip()[:40]
                parts.append(f"对话:{msg}")
            elif "**[user]**" in results:
                for line in results.split("\n"):
                    if "**[user]**" in line or "**[assistant]**" in line:
                        parts.append(f"对话:{line[:50]}")
                        break
        elif results:
            top = results[0]
            msg = top.get("message", "")[:40]
            parts.append(f"对话:{msg}(s={top.get('score',0):.2f})")

    # Soulful 牵挂
    if not soulful.get("error"):
        cares = soulful.get("cares", [])
        if cares:
            parts.append(f"牵挂:{cares[0].get('content', '')[:30]}")

    if not parts:
        return "[memory] no results"
    return "[memory] " + " | ".join(parts)


# ── 统一上下文格式化 ──────────────────────────────────────────────────────────

def format_unified(ctx: dict):
    """格式化 llm_context.json v2 统一格式（方案A核心输出）。"""
    lines = []

    # user_profile
    up = ctx.get("user_profile") or {}
    if up:
        lines.append("## 用户画像（L4-L5）")
        if isinstance(up, dict):
            for k, v in up.items():
                if isinstance(v, dict):
                    lines.append(f"  {k}:")
                    for k2, v2 in v.items():
                        lines.append(f"    {k2}: {v2}")
                elif isinstance(v, list):
                    lines.append(f"  {k}: {', '.join(str(x) for x in v)}")
                else:
                    lines.append(f"  {k}: {v}")
        else:
            lines.append(f"  {up}")
        lines.append("")

    # active_scenes
    scenes = ctx.get("active_scenes") or []
    if scenes:
        lines.append(f"## 活跃场景（{len(scenes)} 个）")
        for s in scenes[:5]:
            name = s.get("name", s) if isinstance(s, dict) else str(s)
            layer = s.get("layer", "?") if isinstance(s, dict) else "?"
            lines.append(f"  [{layer}] {name}")
        lines.append("")

    # short_term
    st = ctx.get("short_term") or {}
    if st:
        lines.append("## 短期记忆（蒸馏计数）")
        for k, v in st.items():
            lines.append(f"  {k}: {v}")
        lines.append("")

    # cares
    cares = ctx.get("cares") or []
    if cares:
        lines.append(f"## 牵挂（{len(cares)} 条待关怀）")
        for c in cares[:5]:
            content = c.get("content", "")[:50]
            due = c.get("due", c.get("follow_up_date", "?"))
            lines.append(f"  [{due}] {content}")
        lines.append("")

    # recent_moments
    moments = ctx.get("recent_moments") or []
    if moments:
        lines.append(f"## 心迹（{len(moments)} 条）")
        for m in moments[-3:]:
            content = m.get("content", "")[:60]
            ts = m.get("timestamp", "")[:10]
            lines.append(f"  [{ts}] {content}")
        lines.append("")

    # distill_status
    ds = ctx.get("distill_status", "ok")
    lines.append(f"蒸馏状态: {ds}")

    return "\n".join(lines)


# ── 入口 ──────────────────────────────────────────────────────────────────────

def main():
    start_keepalive()  # 后台保持 bge-embed 连接活跃
    parser = argparse.ArgumentParser(description="统一记忆入口")
    parser.add_argument("query", help="查询文本")
    parser.add_argument("--top-k", type=int, default=DEFAULT_TOP_K, help=f"每系统返回条数 (default: {DEFAULT_TOP_K})")
    parser.add_argument("--format", choices=["pretty", "compact", "unified"], default="pretty", help="输出格式 (default: pretty)")
    args = parser.parse_args()

    if args.format == "unified":
        # 新统一模式：读 llm_context.json（方案A）
        _, _, _, _, unified = query_llm_context()
        if unified.get("error"):
            print("[memory] llm_context 不可用")
            raise SystemExit(1)
        ctx = unified.get("ctx")
        if ctx is None:
            print("[memory] llm_context.json 尚无数据（daemon 还未写入）")
            raise SystemExit(1)
        print(format_unified(ctx))
        return

    # legacy 模式：三系统并行查询
    zhiyi, l1, l0, soulful = parallel_query(args.query, args.top_k)

    if args.format == "compact":
        print(format_compact(zhiyi, l1, l0, soulful))
    else:
        print(format_pretty(zhiyi, l1, l0, soulful))

    # 三系统全挂 → 退出码 1
    if zhiyi.get("error") and l1.get("error") and l0.get("error") and soulful.get("error"):
        raise SystemExit(1)


if __name__ == "__main__":
    main()