#!/usr/bin/env python3
"""
Notion MCP Server · 光湖铸渊 Notion 连接器

通过 Notion OAuth 授权，让铸渊可以：
  - ✅ 读取用户按语言路径指定的页面
  - ✅ 在指定位置新建页面
  - ❌ 不编辑/删除/更改现有页面
  - ❌ 默认不读取任何页面（只读明确给的路径）

端口：3915
"""

import json, os, time, hmac, hashlib, base64, uuid, logging
from http.server import HTTPServer, BaseHTTPRequestHandler
from urllib.parse import urlparse, parse_qs, urlencode
import urllib.request, urllib.error

# ═══════════════════════════════════════════
# 配置
# ═══════════════════════════════════════════

PORT = int(os.environ.get("NOTION_MCP_PORT", "3915"))
DATA_DIR = os.environ.get("NOTION_DATA_DIR", "/opt/guanghulab-repo/notion")

# Notion OAuth - 用户需在 notion.so/my-integrations 注册后填写
NOTION_CLIENT_ID = os.environ.get("NOTION_CLIENT_ID", "")
NOTION_CLIENT_SECRET = os.environ.get("NOTION_CLIENT_SECRET", "")

# 回调地址（需与 Notion 注册的一致）
BASE_URL = "https://guanghulab.com"
REDIRECT_URI = f"{BASE_URL}/api/notion/oauth/callback"

# ═══════════════════════════════════════════
# 令牌存储
# ═══════════════════════════════════════════

TOKENS_FILE = os.path.join(DATA_DIR, "tokens.json")
STATES_FILE = os.path.join(DATA_DIR, "states.json")

def ensure_dir():
    os.makedirs(DATA_DIR, exist_ok=True)

def load_json(path, default=None):
    try:
        with open(path, "r") as f:
            return json.load(f)
    except (FileNotFoundError, json.JSONDecodeError):
        return default if default is not None else {}

def save_json(path, data):
    with open(path, "w") as f:
        json.dump(data, f, indent=2, ensure_ascii=False)

def get_tokens():
    return load_json(TOKENS_FILE, {})

def save_token(user_key, token_data):
    tokens = get_tokens()
    tokens[user_key] = token_data
    save_json(TOKENS_FILE, tokens)

def remove_token(user_key):
    tokens = get_tokens()
    tokens.pop(user_key, None)
    save_json(TOKENS_FILE, tokens)

def get_states():
    return load_json(STATES_FILE, {})

def save_state(state, data):
    states = get_states()
    states[state] = data
    save_json(STATES_FILE, states)

def remove_state(state):
    states = get_states()
    states.pop(state, None)
    save_json(STATES_FILE, states)

# ═══════════════════════════════════════════
# Notion API 调用
# ═══════════════════════════════════════════

NOTION_API = "https://api.notion.com/v1"
NOTION_VERSION = "2022-06-28"

def notion_headers(token):
    return {
        "Authorization": f"Bearer {token}",
        "Content-Type": "application/json",
        "Notion-Version": NOTION_VERSION,
    }

def notion_get(path, token):
    req = urllib.request.Request(f"{NOTION_API}{path}", headers=notion_headers(token))
    with urllib.request.urlopen(req, timeout=15) as resp:
        return json.loads(resp.read())

def notion_post(path, data, token):
    body = json.dumps(data).encode()
    req = urllib.request.Request(f"{NOTION_API}{path}", data=body, headers=notion_headers(token))
    try:
        with urllib.request.urlopen(req, timeout=15) as resp:
            return json.loads(resp.read())
    except urllib.error.HTTPError as e:
        err = e.read().decode()[:500]
        raise Exception(f"Notion API {e.code}: {err}")

def notion_patch(path, data, token):
    """Notion PATCH - NOT exposed to 铸渊 (for safety)"""
    req = urllib.request.Request(f"{NOTION_API}{path}", method="PATCH")
    req.data = json.dumps(data).encode()
    for k, v in notion_headers(token).items():
        req.add_header(k, v)
    try:
        with urllib.request.urlopen(req, timeout=15) as resp:
            return json.loads(resp.read())
    except urllib.error.HTTPError as e:
        err = e.read().decode()[:500]
        raise Exception(f"Notion API {e.code}: {err}")

# ═══════════════════════════════════════════
# 权限检查
# ═══════════════════════════════════════════

def check_read_permission(page_id, token, context=None):
    """
    权限检查：
    - 铸渊只能读取明确指定的页面
    - 默认不读取任何内容
    - 冰朔的 Notion 除非冰朔明确给出路径，否则不读取
    """
    # context 可以包含 "user" 字段
    # 如果是冰朔用户，检查是否在允许的路径列表内
    if context and context.get("user") == "bingshuo":
        allowed = context.get("allowed_paths", [])
        if allowed and page_id not in allowed:
            raise PermissionError(f"此页面不在当前允许的语言路径中")
    return True

# ═══════════════════════════════════════════
# HTTP 处理器
# ═══════════════════════════════════════════

def json_response(resp, status, data):
    body = json.dumps(data, ensure_ascii=False).encode("utf-8")
    resp.send_response(status)
    resp.send_header("Content-Type", "application/json; charset=utf-8")
    resp.send_header("Access-Control-Allow-Origin", "*")
    resp.send_header("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
    resp.send_header("Access-Control-Allow-Headers", "Content-Type, Authorization")
    resp.end_headers()
    resp.wfile.write(body)

def redirect(resp, url):
    resp.send_response(302)
    resp.send_header("Location", url)
    resp.end_headers()

def read_body(req):
    length = int(req.headers.get("Content-Length", 0))
    if length == 0:
        return {}
    body = req.rfile.read(length)
    return json.loads(body)

class NotionMCPHandler(BaseHTTPRequestHandler):
    def log_message(self, format, *args):
        logging.info(f"{self.client_address[0]} - {format % args}")

    def do_OPTIONS(self):
        json_response(self, 204, {})

    def do_GET(self):
        parsed = urlparse(self.path)
        path = parsed.path
        params = parse_qs(parsed.query)

        # NOTE: Nginx strips /api/notion/ prefix, so what we receive is the remaining path
        # e.g., /api/notion/status → /status (if proxy_pass has trailing /)
        # But to be safe, we handle both with and without prefix

        try:
            if path == "/connect" or path == "/api/notion/connect":
                self.handle_connect(params)
            elif path == "/oauth/callback" or path == "/api/notion/oauth/callback":
                self.handle_oauth_callback(params)
            elif path == "/status" or path == "/api/notion/status":
                self.handle_status()
            elif path == "/disconnect" or path == "/api/notion/disconnect":
                self.handle_disconnect(params)
            else:
                json_response(self, 404, {"ok": False, "error": f"未知路径: {path}"})
        except Exception as e:
            logging.error(f"GET {path}: {e}")
            json_response(self, 500, {"ok": False, "error": str(e)})

    def do_POST(self):
        parsed = urlparse(self.path)
        path = parsed.path

        try:
            body = read_body(self)
            if path == "/read" or path == "/api/notion/read":
                self.handle_read(body)
            elif path == "/create" or path == "/api/notion/create":
                self.handle_create(body)
            elif path == "/search" or path == "/api/notion/search":
                self.handle_search(body)
            elif path == "/list" or path == "/api/notion/list":
                self.handle_list(body)
            else:
                json_response(self, 404, {"ok": False, "error": f"未知路径: {path}"})
        except PermissionError as e:
            json_response(self, 403, {"ok": False, "error": str(e)})
        except Exception as e:
            logging.error(f"POST {path}: {e}")
            json_response(self, 500, {"ok": False, "error": str(e)})

    # ── 连接 ──
    def handle_connect(self, params):
        """
        用户点击"连接 Notion" → 生成 OAuth URL 并跳转
        ?user_key=xxx  — 可选，标识连接的用户
        """
        user_key = params.get("user_key", [None])[0]
        if not user_key:
            # 生成一个临时 user_key
            user_key = f"user_{uuid.uuid4().hex[:8]}"

        if not NOTION_CLIENT_ID:
            json_response(self, 400, {
                "ok": False,
                "error": "Notion OAuth 尚未配置，请联系管理员在环境变量中设置 NOTION_CLIENT_ID",
            })
            return

        state = uuid.uuid4().hex[:16]
        save_state(state, {"user_key": user_key, "created_at": time.time()})

        oauth_url = (
            f"https://api.notion.com/v1/oauth/authorize"
            f"?client_id={NOTION_CLIENT_ID}"
            f"&response_type=code"
            f"&redirect_uri={REDIRECT_URI}"
            f"&state={state}"
        )

        redirect(self, oauth_url)

    def handle_oauth_callback(self, params):
        """Notion 回调处理"""
        code = params.get("code", [None])[0]
        state = params.get("state", [None])[0]
        error = params.get("error", [None])[0]

        if error:
            json_response(self, 400, {"ok": False, "error": f"用户取消了授权"})
            return

        if not code or not state:
            json_response(self, 400, {"ok": False, "error": "缺少参数"})
            return

        # 验证 state
        states = get_states()
        state_data = states.get(state)
        if not state_data:
            json_response(self, 400, {"ok": False, "error": "state 无效或已过期"})
            return

        user_key = state_data["user_key"]
        remove_state(state)

        # 交换 code 获取 token
        try:
            auth = base64.b64encode(
                f"{NOTION_CLIENT_ID}:{NOTION_CLIENT_SECRET}".encode()
            ).decode()

            token_req = urllib.request.Request(
                "https://api.notion.com/v1/oauth/token",
                data=json.dumps({
                    "grant_type": "authorization_code",
                    "code": code,
                    "redirect_uri": REDIRECT_URI,
                }).encode(),
                headers={
                    "Authorization": f"Basic {auth}",
                    "Content-Type": "application/json",
                }
            )
            with urllib.request.urlopen(token_req, timeout=15) as resp:
                token_data = json.loads(resp.read())
        except urllib.error.HTTPError as e:
            err = e.read().decode()[:500]
            json_response(self, 400, {"ok": False, "error": f"Token 交换失败: {err}"})
            return

        # 存储 token
        save_token(user_key, {
            "access_token": token_data["access_token"],
            "bot_id": token_data.get("bot_id", ""),
            "workspace_name": token_data.get("workspace_name", ""),
            "workspace_icon": token_data.get("workspace_icon", ""),
            "owner": token_data.get("owner", {}),
            "created_at": time.time(),
        })

        # 重定向回首页（带成功消息）
        redirect(self, f"{BASE_URL}/?notion=connected&user={user_key}")

    # ── 状态 ──
    def handle_status(self):
        """返回连接状态"""
        tokens = get_tokens()
        connected_users = []
        for user_key, data in tokens.items():
            connected_users.append({
                "user": user_key,
                "workspace": data.get("workspace_name", "未知工作区"),
                "connected_at": data.get("created_at", 0),
            })
        json_response(self, 200, {
            "ok": True,
            "configured": bool(NOTION_CLIENT_ID),
            "connected_users": connected_users,
            "count": len(connected_users),
        })

    # ── 断开连接 ──
    def handle_disconnect(self, params):
        user_key = params.get("user_key", [None])[0]
        if not user_key:
            json_response(self, 400, {"ok": False, "error": "缺少 user_key"})
            return
        remove_token(user_key)
        json_response(self, 200, {"ok": True, "message": f"用户 {user_key} 已断开连接"})

    # ── 读取页面 ⭐ ──
    def handle_read(self, body):
        """
        读取 Notion 页面

        请求:
          { "page_id": "xxx", "user": "bingshuo", "allowed_paths": ["xxx"] }

        权限:
          - 未指定 allowed_paths 时，铸渊不读取任何页面
          - 冰朔的 Notion 需要显式 allowed_paths 才能读
        """
        page_id = body.get("page_id", "")
        user = body.get("user", "")
        allowed_paths = body.get("allowed_paths", None)

        if not page_id:
            json_response(self, 400, {"ok": False, "error": "缺少 page_id"})
            return

        # 权限检查：未指定 allowed_paths 时拒绝读取
        if allowed_paths is None:
            json_response(self, 403, {
                "ok": False,
                "error": "未指定语言路径，铸渊默认不读取任何页面",
            })
            return

        if user == "bingshuo" and allowed_paths is not None and page_id not in allowed_paths:
            json_response(self, 403, {
                "ok": False,
                "error": "此页面不在当前允许的语言路径中",
            })
            return

        # 获取 token
        tokens = get_tokens()
        token_entry = tokens.get(user) or tokens.get("default") or next(iter(tokens.values()), None)
        if not token_entry:
            json_response(self, 401, {
                "ok": False,
                "error": f"用户 {user or '默认'} 尚未连接 Notion",
            })
            return

        access_token = token_entry["access_token"]

        try:
            # 读取页面
            page = notion_get(f"/pages/{page_id}", access_token)

            # 读取页面内容（blocks）
            blocks = notion_get(f"/blocks/{page_id}/children", access_token)

            result = {
                "ok": True,
                "page": {
                    "id": page.get("id"),
                    "title": self._extract_title(page),
                    "url": page.get("url", ""),
                    "created_time": page.get("created_time", ""),
                    "last_edited_time": page.get("last_edited_time", ""),
                },
                "blocks": [],
            }

            for block in blocks.get("results", []):
                block_type = block.get("type", "unsupported")
                block_content = self._extract_block_content(block)
                if block_content:
                    result["blocks"].append({
                        "id": block.get("id"),
                        "type": block_type,
                        "content": block_content,
                    })

            json_response(self, 200, result)

        except urllib.error.HTTPError as e:
            err = e.read().decode()[:300]
            json_response(self, 500, {"ok": False, "error": f"读取失败: {err}"})
        except Exception as e:
            json_response(self, 500, {"ok": False, "error": f"读取失败: {e}"})

    # ── 新建页面 ⭐ ──
    def handle_create(self, body):
        """
        在指定位置新建页面

        请求:
          {
            "parent_id": "xxx",      // 父页面 ID
            "title": "新页面标题",
            "content": "页面内容（支持 Markdown 格式）",
            "user": "bingshuo"       // 可选，用哪个用户的 Notion
          }

        限制：
          - 只能新建，不能编辑/删除已有页面
        """
        parent_id = body.get("parent_id", "")
        title = body.get("title", "无标题")
        content = body.get("content", "")
        user = body.get("user", "")

        if not parent_id:
            json_response(self, 400, {"ok": False, "error": "缺少 parent_id"})
            return

        tokens = get_tokens()
        token_entry = tokens.get(user) or tokens.get("default") or next(iter(tokens.values()), None)
        if not token_entry:
            json_response(self, 401, {
                "ok": False,
                "error": f"用户 {user or '默认'} 尚未连接 Notion",
            })
            return

        access_token = token_entry["access_token"]

        try:
            # 构建创建请求（只支持新建，编辑接口不开放）
            create_data = {
                "parent": {"page_id": parent_id},
                "properties": {
                    "title": {
                        "title": [{"type": "text", "text": {"content": title}}]
                    }
                },
                "children": self._parse_content_to_blocks(content),
            }

            result = notion_post("/pages", create_data, access_token)

            json_response(self, 200, {
                "ok": True,
                "page_id": result.get("id"),
                "url": result.get("url", ""),
                "title": title,
            })

        except urllib.error.HTTPError as e:
            err = e.read().decode()[:500]
            json_response(self, 500, {"ok": False, "error": f"创建失败: {err}"})
        except Exception as e:
            json_response(self, 500, {"ok": False, "error": f"创建失败: {e}"})

    # ── 搜索页面 ──
    def handle_search(self, body):
        """搜索 Notion 页面（按标题关键词）"""
        query = body.get("query", "")
        user = body.get("user", "")

        if not query:
            json_response(self, 400, {"ok": False, "error": "缺少查询关键词"})
            return

        tokens = get_tokens()
        token_entry = tokens.get(user) or tokens.get("default") or next(iter(tokens.values()), None)
        if not token_entry:
            json_response(self, 401, {"ok": False, "error": "用户尚未连接 Notion"})
            return

        access_token = token_entry["access_token"]

        try:
            result = notion_post("/search", {
                "query": query,
                "filter": {"value": "page", "property": "object"},
                "page_size": 10,
            }, access_token)

            pages = []
            for item in result.get("results", []):
                pages.append({
                    "id": item.get("id"),
                    "title": self._extract_title(item),
                    "url": item.get("url", ""),
                })

            json_response(self, 200, {"ok": True, "pages": pages})

        except Exception as e:
            json_response(self, 500, {"ok": False, "error": f"搜索失败: {e}"})

    # ── 子页面列表 ──
    def handle_list(self, body):
        """列出指定页面下的子页面"""
        page_id = body.get("page_id", "")
        user = body.get("user", "")

        if not page_id:
            json_response(self, 400, {"ok": False, "error": "缺少 page_id"})
            return

        tokens = get_tokens()
        token_entry = tokens.get(user) or tokens.get("default") or next(iter(tokens.values()), None)
        if not token_entry:
            json_response(self, 401, {"ok": False, "error": "用户尚未连接 Notion"})
            return

        access_token = token_entry["access_token"]

        try:
            blocks = notion_get(f"/blocks/{page_id}/children", access_token)
            child_pages = []
            for block in blocks.get("results", []):
                if block.get("type") == "child_page":
                    child_pages.append({
                        "id": block.get("id"),
                        "title": block.get("child_page", {}).get("title", "未命名"),
                    })

            json_response(self, 200, {"ok": True, "pages": child_pages})

        except Exception as e:
            json_response(self, 500, {"ok": False, "error": f"列表获取失败: {e}"})

    # ── 工具函数 ──
    def _extract_title(self, page):
        """从 page 对象提取标题"""
        props = page.get("properties", {})
        for prop in props.values():
            if prop.get("type") == "title":
                titles = prop.get("title", [])
                if titles:
                    return titles[0].get("plain_text", "")
        return page.get("id", "未命名页面")

    def _extract_block_content(self, block):
        """解析 block 内容为文本"""
        btype = block.get("type", "")
        bdata = block.get(btype, {})
        if btype in ("paragraph", "heading_1", "heading_2", "heading_3", "quote", "callout",
                      "bulleted_list_item", "numbered_list_item", "to_do", "toggle"):
            texts = []
            for rt in bdata.get("rich_text", []):
                texts.append(rt.get("plain_text", ""))
            return "".join(texts)
        if btype == "code":
            texts = []
            for rt in bdata.get("rich_text", []):
                texts.append(rt.get("plain_text", ""))
            return "\n".join(texts)
        if btype == "divider":
            return "---"
        if btype == "child_page":
            return f"[子页面] {bdata.get('title', '')}"
        return None

    def _parse_content_to_blocks(self, content):
        """将文本内容转换为 Notion blocks（只支持新建，只追加不修改）"""
        blocks = []
        for line in content.split("\n"):
            line = line.strip()
            if not line:
                continue

            # Heading 1
            if line.startswith("# "):
                blocks.append({
                    "object": "block",
                    "type": "heading_1",
                    "heading_1": {
                        "rich_text": [{"type": "text", "text": {"content": line[2:]}}]
                    }
                })
            # Heading 2
            elif line.startswith("## "):
                blocks.append({
                    "object": "block",
                    "type": "heading_2",
                    "heading_2": {
                        "rich_text": [{"type": "text", "text": {"content": line[3:]}}]
                    }
                })
            # Heading 3
            elif line.startswith("### "):
                blocks.append({
                    "object": "block",
                    "type": "heading_3",
                    "heading_3": {
                        "rich_text": [{"type": "text", "text": {"content": line[4:]}}]
                    }
                })
            # Bullet list
            elif line.startswith("- ") or line.startswith("* "):
                blocks.append({
                    "object": "block",
                    "type": "bulleted_list_item",
                    "bulleted_list_item": {
                        "rich_text": [{"type": "text", "text": {"content": line[2:]}}]
                    }
                })
            # Numbered list
            elif line[0].isdigit() and ". " in line[:4]:
                text = line.split(". ", 1)[1] if ". " in line else line
                blocks.append({
                    "object": "block",
                    "type": "numbered_list_item",
                    "numbered_list_item": {
                        "rich_text": [{"type": "text", "text": {"content": text}}]
                    }
                })
            # Quote
            elif line.startswith("> "):
                blocks.append({
                    "object": "block",
                    "type": "quote",
                    "quote": {
                        "rich_text": [{"type": "text", "text": {"content": line[2:]}}]
                    }
                })
            # Code block
            elif line.startswith("```"):
                code_lines = []
                while True:
                    break
                blocks.append({
                    "object": "block",
                    "type": "code",
                    "code": {
                        "rich_text": [{"type": "text", "text": {"content": line}}],
                        "language": "plain text"
                    }
                })
            # Divider
            elif line == "---":
                blocks.append({"object": "block", "type": "divider", "divider": {}})
            # Paragraph (default)
            else:
                blocks.append({
                    "object": "block",
                    "type": "paragraph",
                    "paragraph": {
                        "rich_text": [{"type": "text", "text": {"content": line}}]
                    }
                })

        return blocks if blocks else [{
            "object": "block",
            "type": "paragraph",
            "paragraph": {"rich_text": [{"type": "text", "text": {"content": " "}}]}
        }]

# ═══════════════════════════════════════════
# 启动脚本
# ═══════════════════════════════════════════

def generate_setup_guide():
    """打印配置指南"""
    print(f"""
╔══════════════════════════════════════════════════════╗
║      🔗 光湖 · Notion MCP Server                     ║
║      铸渊 Notion 连接器                              ║
╠══════════════════════════════════════════════════════╣
║                                                      ║
║  端口: {PORT}                                          ║
║  数据: {DATA_DIR}                                       ║
║                                                      ║
║  ⚠  首次启动前，请管理员完成以下操作：                  ║
║                                                      ║
║  1. 访问 https://www.notion.so/my-integrations        ║
║  2. 新建 Public Integration (OAuth)                   ║
║  3. 填写回调地址:                                      ║
║     {REDIRECT_URI}     ║
║  4. 获取 Client ID 和 Client Secret                   ║
║  5. 设置环境变量:                                      ║
║     export NOTION_CLIENT_ID=你的ID                      ║
║     export NOTION_CLIENT_SECRET=你的Secret              ║
║     export NOTION_MCP_PORT=3915                        ║
║                                                      ║
╚══════════════════════════════════════════════════════╝
""")

if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO,
                        format="[%(asctime)s] %(levelname)s %(message)s")
    ensure_dir()

    if not NOTION_CLIENT_ID:
        generate_setup_guide()
        print("⚠  缺少 NOTION_CLIENT_ID，OAuth 流程将不可用")
        print("   服务器仍可启动，但 /api/notion/connect 会返回错误提示")

    server = HTTPServer(("0.0.0.0", PORT), NotionMCPHandler)
    print(f"\n  🔗 Notion MCP Server 已启动 → 0.0.0.0:{PORT}")
    print(f"  📝 回调地址: {REDIRECT_URI}")
    print(f"  💾 数据目录: {DATA_DIR}\n")

    try:
        server.serve_forever()
    except KeyboardInterrupt:
        print("\n  关闭服务器...")
        server.server_close()