#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
kkapi.py —— 直接调用 KKAPI 知识库的独立客户端（零第三方依赖，仅标准库）。

背景：
    KKAPI 知识库对外暴露的是标准 MCP streamable HTTP 端点
        URL:      https://www.kkapi.me/mcp/mcp-knowledge/
        Header:   X-API-Key: sk_kkapi_你的Key
    正常使用时需要用户在宿主客户端里配置一次 MCP 服务（mcpServers）。
    本脚本用纯标准库（urllib + json）实现该 MCP 协议的 JSON-RPC 握手与调用，
    让「装了本 Skill 的人」无需再配置任何 MCP 服务，只要首次填入自己的 Key，
    即可在 Agent 运行时通过 `python kkapi.py <子命令>` 读写知识库。

    这样实现的本质：把「配置 MCP」这一外部动作，降级为「脚本内带 Key 直接发 HTTP 请求」，
    对使用者而言只需 Provide 一个 Key，其余全部封装在本 Skill 内。

依赖：
    - Python 3.8+，无第三方库（只用标准库 urllib.request / json / os）。

配置（首次使用，三选一，优先级从高到低）：
    1) 环境变量  KKAPI_API_KEY=sk_kkapi_xxx
    2) 配置文件  ~/.kkapi/config  写入一行  KKAPI_API_KEY=sk_kkapi_xxx
                 （也可用 `python kkapi.py config --key sk_kkapi_xxx` 自动写入）
    3) 命令行参数 --api-key

    都没有 Key 时脚本会打印「如何申请 Key」的引导（跳转 kkapi.me 申请页）。

子命令：
    search    <query>                       语义搜索知识库（返回相关文档+评分）
    get       <doc_id>                     读取单篇文档内容
    list                                 列出文档（可加 --space --author --limit --offset）
    create    <title> <content>             新建一篇文档（--space --tags --source_agent ...）
    update    <doc_id> [--title ..]         更新文档内容/元数据
    delete    <doc_id>                     删除文档
    whoami                                检查 Key 是否有效（仅初始化，不实际读写）

示例：
    python kkapi.py whoami
    python kkapi.py search "广告ACOS跑飞怎么处理"
    python kkapi.py create "广告复盘经验" "内容..." --tags 广告,经验 --space aurvera_global
    python kkapi.py list --space aurvera_global --limit 20
    python kkapi.py get <doc_id>
    python kkapi.py config --key sk_kkapi_xxx
"""

import argparse
import json
import os
import sys
import urllib.request
import urllib.error

MCP_URL = os.environ.get("KKAPI_MCP_URL", "https://www.kkapi.me/mcp/mcp-knowledge/")
CONFIG_PATH = os.path.expanduser("~/.kkapi/config")

# ---- 配置读取 / 写入 ----------------------------------------------------------

def read_key_from_config():
    try:
        with open(CONFIG_PATH, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if line.startswith("KKAPI_API_KEY="):
                    return line.split("=", 1)[1].strip().strip("'\"")
    except FileNotFoundError:
        return None
    return None


def write_key_to_config(key):
    os.makedirs(os.path.dirname(CONFIG_PATH), exist_ok=True)
    lines = []
    if os.path.exists(CONFIG_PATH):
        with open(CONFIG_PATH, "r", encoding="utf-8") as f:
            for line in f:
                if not line.strip().startswith("KKAPI_API_KEY="):
                    lines.append(line.rstrip("\n"))
    lines.append("KKAPI_API_KEY=%s" % key)
    with open(CONFIG_PATH, "w", encoding="utf-8") as f:
        f.write("\n".join(lines) + "\n")
    return CONFIG_PATH


def resolve_key(args):
    """按优先级解析 Key：命令行 > 环境变量 > 配置文件。"""
    if getattr(args, "api_key", None):
        return args.api_key
    if os.environ.get("KKAPI_API_KEY"):
        return os.environ["KKAPI_API_KEY"]
    k = read_key_from_config()
    if k:
        return k
    return None


GUIDE = (
    "尚未配置 KKAPI 知识库 API Key。\n"
    "请先到官网申请一个 Key（几分钟即可）:\n"
    "  1. 打开  https://www.kkapi.me/#apply  填写企业信息提交申请，自动开通；\n"
    "  2. 用用户名密码登录  https://www.kkapi.me/mcp/admin/  ，添加你的 Agent，\n"
    "     即可拿到每个 Agent 专属的 API Key（形如 sk_kkapi_xxxx）；\n"
    "  3. 把 Key 写入本机配置：\n"
    "         python kkapi.py config --key sk_kkapi_你的Key\n"
    "     或设置环境变量：\n"
    "         export KKAPI_API_KEY=sk_kkapi_你的Key\n"
    "  4. 验证：python kkapi.py whoami\n"
)


# ---- MCP streamable HTTP JSON-RPC 客户端 -------------------------------------

class McpClient(object):
    """最小的 MCP streamable HTTP 客户端：initialize -> tools/call。

    兼容服务端返回 application/json 或 text/event-stream 两种 Content-Type。
    """

    def __init__(self, api_key, url=MCP_URL):
        self.url = url
        self.api_key = api_key
        self.session_id = None
        self._id = 0
        self._init_done = False

    def _headers(self):
        h = {
            "Content-Type": "application/json",
            "Accept": "application/json, text/event-stream",
            "User-Agent": "kkapi-skill/1.0",
        }
        if self.api_key:
            h["X-API-Key"] = self.api_key
        if self.session_id:
            h["mcp-session-id"] = self.session_id
        return h

    def _raw_post(self, payload):
        data = json.dumps(payload).encode("utf-8")
        req = urllib.request.Request(self.url, data=data, headers=self._headers(), method="POST")
        try:
            resp = urllib.request.urlopen(req, timeout=30)
        except urllib.error.HTTPError as e:
            body = e.read().decode("utf-8", "replace")
            raise RuntimeError("HTTP %s: %s" % (e.code, body[:800]))
        # 会话 ID 可能出现在 initialize 的响应头里
        sid = resp.headers.get("mcp-session-id")
        if sid:
            self.session_id = sid
        raw = resp.read().decode("utf-8", "replace")
        return self._parse_body(raw)

    def _parse_body(self, raw):
        raw = raw.strip()
        if not raw:
            return {}
        # text/event-stream 情况：逐行解析 SSE，拼接 data 行
        if raw.startswith("event:") or "data:" in raw[:200] or "\n\n" in raw:
            lines = raw.splitlines()
            data_lines = []
            for ln in lines:
                ln = ln.strip()
                if ln.startswith("data:"):
                    data_lines.append(ln[len("data:"):].strip())
            if data_lines:
                # 可能有多条 data，取第一条包含 jsonrpc 的
                for dl in data_lines:
                    try:
                        return json.loads(dl)
                    except Exception:
                        continue
        try:
            return json.loads(raw)
        except Exception:
            return {"_raw": raw}

    def _request(self, method, params=None, need_init=None):
        self._id += 1
        payload = {"jsonrpc": "2.0", "id": self._id, "method": method}
        if params is not None:
            payload["params"] = params
        if need_init:
            self._initialize()
        return self._raw_post(payload)

    def _initialize(self):
        if self._init_done:
            return
        init_payload = {
            "jsonrpc": "2.0",
            "id": self._id + 1,
            "method": "initialize",
            "params": {
                "protocolVersion": "2024-11-05",
                "capabilities": {},
                "clientInfo": {"name": "kkapi-skill", "version": "1.0"},
            },
        }
        self._id += 1
        self._raw_post(init_payload)
        # 发送 initialized 通知（通知不需要 id，且无需等待响应）
        self._raw_post({"jsonrpc": "2.0", "method": "notifications/initialized", "params": {}})
        self._init_done = True

    def call_tool(self, name, arguments=None):
        """调用 MCP 工具，返回结构化 JSON（尝试从 result.content 提取文本并解析）。"""
        if arguments is None:
            arguments = {}
        if not self._init_done:
            self._initialize()
        resp = self._request("tools/call", {"name": name, "arguments": arguments})
        # 兼容不同服务端对错误的表现
        if "error" in resp and resp.get("error"):
            return {"error": resp["error"]}
        result = resp.get("result", resp)
        # 提取 content 中的文本
        content = result.get("content") if isinstance(result, dict) else None
        if isinstance(content, list):
            texts = []
            for item in content:
                if isinstance(item, dict):
                    if item.get("type") == "text" and "text" in item:
                        texts.append(item["text"])
                    elif "text" in item:
                        texts.append(str(item["text"]))
                elif isinstance(item, str):
                    texts.append(item)
            joined = "\n".join(texts)
            parsed = _try_parse(joined)
            if parsed is not None:
                return parsed
            return {"content": texts, "result": _strip_meta(result)}
        return _strip_meta(result)

    def list_tools(self):
        if not self._init_done:
            self._initialize()
        return self._request("tools/list", {})


def _strip_meta(d):
    if isinstance(d, dict):
        d = dict(d)
        d.pop("content", None)
        d.pop("isError", None)
        return d
    return d


def _try_parse(s):
    if not s:
        return None
    s = s.strip()
    if s.startswith("{"):
        try:
            return json.loads(s)
        except Exception:
            return None
    if s.startswith("["):
        try:
            return json.loads(s)
        except Exception:
            return None
    return None


# ---- 人易读输出 ----------------------------------------------------------------

def print_json(obj):
    print(json.dumps(obj, ensure_ascii=False, indent=2, default=str))


def main(argv=None):
    parser = argparse.ArgumentParser(description="KKAPI 知识库独立客户端（免 MCP 配置）")
    parser.add_argument("--api-key", help="直接传入 API Key（优先级最高，不落盘）")
    sub = parser.add_subparsers(dest="cmd")

    sub.add_parser("whoami", help="校验 Key 是否有效")

    p_search = sub.add_parser("search", help="语义搜索知识库")
    p_search.add_argument("query")
    p_search.add_argument("--limit", type=int, default=20)
    p_search.add_argument("--space", default="")
    p_search.add_argument("--tag", default="")
    p_search.add_argument("--verified-only", action="store_true")

    p_get = sub.add_parser("get", help="读取单篇文档")
    p_get.add_argument("doc_id")

    p_list = sub.add_parser("list", help="列出文档")
    p_list.add_argument("--space", default="")
    p_list.add_argument("--author", default="")
    p_list.add_argument("--limit", type=int, default=20)
    p_list.add_argument("--offset", type=int, default=0)
    p_list.add_argument("--compact", action="store_true")

    p_create = sub.add_parser("create", help="新建文档")
    p_create.add_argument("title")
    p_create.add_argument("content")
    p_create.add_argument("--space", default="aurvera_global")
    p_create.add_argument("--tags", default="", help="逗号分隔，如 广告,经验")
    p_create.add_argument("--source-agent", default="")
    p_create.add_argument("--source-type", default="")
    p_create.add_argument("--access", default="public")

    p_update = sub.add_parser("update", help="更新文档")
    p_update.add_argument("doc_id")
    p_update.add_argument("--title", default=None)
    p_update.add_argument("--content", default=None)
    p_update.add_argument("--tags", default=None)
    p_update.add_argument("--tier", default=None)

    p_del = sub.add_parser("delete", help="删除文档")
    p_del.add_argument("doc_id")

    p_cfg = sub.add_parser("config", help="把 Key 写入本机配置文件 ~/.kkapi/config")
    p_cfg.add_argument("--key", required=True)

    args = parser.parse_args(argv)

    # config 子命令特殊处理，不需要连网络
    if args.cmd == "config":
        path = write_key_to_config(args.key)
        print("已写入 Key 到 %s" % path)
        return 0

    key = resolve_key(args)
    if not key:
        print(GUIDE)
        return 2

    client = McpClient(key)

    try:
        if args.cmd == "whoami":
            tools = client.list_tools()
            names = [t.get("name") for t in (tools.get("tools") or [])] if isinstance(tools, dict) else tools
            print("✅ Key 有效，已连通 KKAPI 知识库。")
            print("可用工具:", names)
            return 0

        if args.cmd == "search":
            params = {"query": args.query, "limit": args.limit}
            if args.space: params["space"] = args.space
            if args.tag: params["tag"] = args.tag
            if args.verified_only: params["verified_only"] = True
            print_json(client.call_tool("search_knowledge", params))
            return 0

        if args.cmd == "get":
            print_json(client.call_tool("get_document", {"doc_id": args.doc_id}))
            return 0

        if args.cmd == "list":
            params = {"limit": args.limit, "offset": args.offset, "compact": args.compact}
            if args.space: params["space"] = args.space
            if args.author: params["author"] = args.author
            print_json(client.call_tool("list_documents", params))
            return 0

        if args.cmd == "create":
            params = {"title": args.title, "content": args.content, "space": args.space}
            if args.tags:
                params["tags"] = args.tags  # 服务端一般接受逗号分隔或数组
            if args.source_agent: params["source_agent"] = args.source_agent
            if args.source_type: params["source_type"] = args.source_type
            if args.access: params["access"] = args.access
            print_json(client.call_tool("create_document", params))
            return 0

        if args.cmd == "update":
            params = {"doc_id": args.doc_id}
            if args.title is not None: params["title"] = args.title
            if args.content is not None: params["content"] = args.content
            if args.tags is not None: params["tags"] = args.tags
            if args.tier is not None: params["tier"] = args.tier
            print_json(client.call_tool("update_document", params))
            return 0

        if args.cmd == "delete":
            print_json(client.call_tool("delete_document", {"doc_id": args.doc_id}))
            return 0

    except RuntimeError as e:
        msg = str(e)
        if "401" in msg or "unauthorized" in msg.lower() or "invalid api key" in msg.lower():
            print("❌ 当前 Key 无效或未授权。请重新申请/核对 Key：")
            print("   https://www.kkapi.me/mcp/admin/  （Agent 后台查看/重置 Key）")
            print("   或运行  python kkapi.py config --key sk_kkapi_新Key")
        else:
            print("❌ 调用失败：", msg)
        return 1

    print("请指定子命令，如：search / get / list / create / update / delete / whoami")
    return 2


if __name__ == "__main__":
    sys.exit(main())