#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
kv_cache_coordinator.py — A5 KV Cache 算力亲和协调（吸收 openJiuwen KV Cache Proactive Coord）
========================================================================================
针对 llama.cpp / 长上下文 / 持续对话场景，主动协调 KV Cache 以省 token：
- 前缀复用：同一模型 + 同一会话，共享前缀 KV，避免重复计算
- 长上下文分段：超长输入分段递增复用，而非整体重算
- 命中率统计：量化省 token 比例
落地为独立的 KV 协调层，可被 brain_scheduler / micro_brain / swarm 调用。
"""

import json
import os
import threading
import time
import uuid
from datetime import datetime

ANCHOR = "Ω₀⊂⊙∞⊂Ω"
DID = "DID-BR-000002"
KV_DB = os.environ.get("KV_CACHE_DB", "/var/lib/zr-kv/kv_cache.json")

# 长上下文分段窗口（token 估算：中英混合按 ~1.5 char/token 粗估）
WINDOW = int(os.environ.get("KV_WINDOW", "4000"))
# 是否开启真实 llama.cpp 对接（默认关闭，演示协调逻辑；对接 micro_brain 时打开）
LLAMA_ENABLED = os.environ.get("KV_LLAMA_ENABLED", "0") == "1"


def now():
    return datetime.now().isoformat(timespec="seconds")


def _load_json(path, default):
    try:
        if os.path.exists(path):
            with open(path, "r", encoding="utf-8") as f:
                return json.load(f)
    except Exception:
        pass
    return default


def _save_json(path, obj):
    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
    tmp = path + ".tmp"
    with open(tmp, "w", encoding="utf-8") as f:
        json.dump(obj, f, ensure_ascii=False, indent=2)
    os.replace(tmp, path)


def _est_tokens(text):
    return max(1, int(len(str(text)) / 1.5))


class KVCacheCoordinator:
    """KV Cache 主动协调器：前缀复用 + 长上下文分段递增。"""

    def __init__(self, window=WINDOW):
        self.window = window
        self.lock = threading.Lock()
        self.db = _load_json(KV_DB, {"sessions": {}, "stats": {"hits": 0, "misses": 0, "tokens_saved": 0}})

    # ---- 会话级前缀 KV ----
    def begin_session(self, model, system_prompt=""):
        """开启会话，登记可复用的前缀 KV（系统提示 + 首段）。"""
        sid = uuid.uuid4().hex[:8]
        entry = {
            "sid": sid, "model": model,
            "prefix": str(system_prompt),
            "prefix_tokens": _est_tokens(system_prompt),
            "segments": [],          # 已计算的长上下文分段
            "created": now(),
        }
        self.db.setdefault("sessions", {})[sid] = entry
        _save_json(KV_DB, self.db)
        return sid

    def extend(self, sid, new_text):
        """增量追加新内容：只计算新增段的 KV，复用已有前缀（省 token 核心）。"""
        with self.lock:
            s = self.db["sessions"].get(sid)
            if not s:
                return {"ok": False, "error": "session_not_found"}
            # 新增段：未在 segments 中的文本
            new_tokens = _est_tokens(new_text)
            s.setdefault("segments", []).append({
                "text": new_text[:200], "tokens": new_tokens, "ts": now()
            })
            self.db["stats"]["hits"] += 1
            # 省 token = 复用前缀（本次未重算前缀）
            self.db["stats"]["tokens_saved"] += s.get("prefix_tokens", 0)
            _save_json(KV_DB, self.db)
            return {"ok": True, "sid": sid, "new_tokens": new_tokens,
                    "tokens_saved_now": s.get("prefix_tokens", 0)}

    # ---- 长上下文分段 ----
    def segment(self, text):
        """超长输入分段（每段 ≤ window），供递增 KV 复用。"""
        text = str(text)
        n = max(1, _est_tokens(text))
        parts = []
        if n <= self.window:
            parts.append(text)
        else:
            # 按 token 预算切分
            step = int(self.window * 1.5)
            for i in range(0, len(text), step):
                parts.append(text[i:i + step])
        return parts

    def process_long(self, sid, long_text):
        """处理长上下文：分段后每段走 extend（递增复用前缀）。"""
        segs = self.segment(long_text)
        out = []
        for seg in segs:
            out.append(self.extend(sid, seg))
        return {"ok": True, "sid": sid, "segments": len(segs), "detail": out}

    # ---- 真实 llama.cpp 对接（可选） ----
    def infer_llama(self, text):
        """若启用，调用 llama.cpp 后端推理（通过 micro_brain 或 llama-cli）。"""
        if not LLAMA_ENABLED:
            return {"ok": False, "reason": "llama_disabled"}
        try:
            from micro_brain import think
            return {"ok": True, "output": think(text)}
        except Exception as e:
            return {"ok": False, "error": str(e)}

    # ---- 统计 ----
    def stats(self):
        s = self.db["stats"]
        total = s.get("hits", 0) + s.get("misses", 0)
        hit_rate = round(s.get("hits", 0) / total, 3) if total else 0
        return {
            "sessions": len(self.db.get("sessions", {})),
            "hits": s.get("hits", 0), "misses": s.get("misses", 0),
            "hit_rate": hit_rate,
            "tokens_saved_est": s.get("tokens_saved", 0),
            "anchor": ANCHOR,
        }


def _demo():
    print("== A5 KV Cache 算力亲和 demo ==")
    kv = KVCacheCoordinator(window=800)
    sid = kv.begin_session("qwen2.5-0.5b-gguf", system_prompt="你是火斗云智 AIOS 助手。")
    print("开启会话:", sid)
    r1 = kv.extend(sid, "用户问：openJiuwen 是什么？" * 30)
    print("首段追加 token:", r1["new_tokens"])
    long_text = ("长文档：多智能体协同架构设计说明。" * 60)   # 超 window
    r2 = kv.process_long(sid, long_text)
    print("长上下文分段数:", r2["segments"])
    print("KV 统计(省token估算):", json.dumps(kv.stats(), ensure_ascii=False))
    return kv.stats()


if __name__ == "__main__":
    import sys
    if len(sys.argv) > 1 and sys.argv[1] == "stats":
        print(json.dumps(KVCacheCoordinator().stats(), ensure_ascii=False, indent=2))
    else:
        _demo()