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

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

import argparse
import json
import threading
import urllib.error
import urllib.request
from datetime import datetime
from pathlib import Path

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

ZHIYI_URL = "http://localhost:7821/api/v1/recall"
ZHIYI_HEADERS = {
    "X-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 = 5

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

def http_post(url: str, payload: dict, headers: dict | None = None, timeout: int = TIMEOUT_SEC) -> dict | None:
    """POST JSON，返回解析后的 dict；失败返回 None 并打印 warning。"""
    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:
        print(f"  [warning] {url} -> {e}", flush=True)
        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，yeild 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)


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

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:
        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,
    }


# ── 并行查询 ───────────────────────────────────────────────────────────────────

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

    def run(query: str, top_k: int, fn):
        return fn(query, top_k)

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

    t1 = threading.Thread(target=lambda: zhiyi_result.update(run(query, top_k, query_zhiyi)))
    t2 = threading.Thread(target=lambda: l1_result.update(run(query, top_k, query_tencentdb_l1)))
    t3 = threading.Thread(target=lambda: l0_result.update(run(query, top_k, query_tencentdb_l0)))
    t4 = threading.Thread(target=lambda: soulful_result.update(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 main():
    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"], default="pretty", help="输出格式 (default: pretty)")
    args = parser.parse_args()

    # 实际并行查询
    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()