#!/usr/bin/env python3
"""Daytrace desktop sync. Python 3 standard library, existing Codex MCP login only."""
import argparse
import base64
import gzip
import hashlib
import hmac
import json
import os
import queue
import re
import secrets
import shutil
import signal
import subprocess
import tempfile
import threading
import time
import webbrowser
from datetime import date, timedelta
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from urllib.error import HTTPError, URLError
from urllib.parse import urlsplit
from urllib.request import HTTPRedirectHandler, Request, build_opener

GATEWAY = "https://beluga-server-cn.heytapmobi.com/beluga-mcp/gateway"
TOOLS = ("recall_memory", "recall_note")
CONFIG_PATH = Path.home() / ".config/daytrace-sync/config.json"
CHILDREN = set()
CHILD_LOCK = threading.Lock()


def server_url(value):
    parts = urlsplit(value.strip())
    local = parts.hostname in ("127.0.0.1", "localhost", "::1")
    if (parts.scheme != "https" and not (parts.scheme == "http" and local)) or not parts.hostname:
        raise ValueError("网站地址需要使用 HTTPS，本机预览可使用 HTTP。")
    if parts.username or parts.password or parts.query or parts.fragment or parts.path not in ("", "/"):
        raise ValueError("请填写网站根地址，例如 https://www.funnycafe.top。")
    return parts.scheme + "://" + parts.netloc


class ApiError(ValueError):
    def __init__(self, status, message):
        super().__init__(message)
        self.status = status


class NoRedirect(HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        return None


def api(server, token, path, data):
    request = Request(server + "/api/diary/memory/desktop" + path,
                      data=json.dumps(data, ensure_ascii=False).encode(), method="POST",
                      headers={"Authorization": "Bearer " + token, "Content-Type": "application/json", "Accept": "application/json"})
    try:
        with build_opener(NoRedirect()).open(request, timeout=45) as response:
            return json.load(response)
    except HTTPError as error:
        error.close()
        messages = {401: "电脑同步凭证已失效，请重新绑定日记。", 404: "网站尚未更新电脑同步接口。",
                    409: "本次同步请求已更新或取消。", 413: "同步内容超过网站大小限制。",
                    422: "同步内容未通过网站校验，原记录已保留。", 429: "网站请求频繁，稍后自动重试。"}
        raise ApiError(error.code, messages.get(error.code, "网站暂时无法接收同步，请稍后重试。")) from None
    except (URLError, TimeoutError, OSError, ValueError):
        raise ApiError(0, "暂时无法连接日迹网站，请检查网络与网站地址。") from None


def save_config(config, path=CONFIG_PATH):
    path.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
    os.chmod(path.parent, 0o700)
    fd, temporary = tempfile.mkstemp(prefix=".config-", dir=path.parent)
    try:
        with os.fdopen(fd, "w") as stream:
            os.chmod(temporary, 0o600)
            json.dump(config, stream)
            stream.flush()
            os.fsync(stream.fileno())
        os.replace(temporary, path)
    finally:
        if os.path.exists(temporary):
            os.unlink(temporary)


def load_config(path=CONFIG_PATH):
    if not path.exists():
        return {}
    os.chmod(path, 0o600)
    config = json.loads(path.read_text())
    config["server"] = server_url(config["server"])
    if not str(config.get("token", "")).startswith("dsc_"):
        raise ValueError("本机同步配置无效，请移除配置后重新绑定。")
    return config


def codex_command():
    binary = shutil.which("codex")
    candidates = [Path.home() / ".deskclaw/node/bin/codex", Path("/Applications/Codex.app/Contents/Resources/codex")]
    if not binary:
        binary = next((str(path) for path in candidates if path.is_file()), None)
    if not binary:
        raise ValueError("没有找到 Codex，请先安装 Codex 命令行并授权心力球。")
    sibling = Path(binary).parent / "node"
    if Path(binary).read_bytes()[:40].startswith(b"#!/usr/bin/env node") and sibling.is_file():
        return [str(sibling), binary]
    return [binary]


def toml(value):
    if isinstance(value, dict):
        return "{" + ", ".join(json.dumps(key) + " = " + toml(item) for key, item in value.items()) + "}"
    if isinstance(value, list):
        return "[" + ", ".join(toml(item) for item in value) + "]"
    return json.dumps(value)


def parse_page(result):
    if result.get("isError") or result.get("is_error"):
        raise ValueError("心力球接口返回错误，请检查本机授权。")
    structured = result.get("structured_content") or result.get("structuredContent") or {}
    raw = structured.get("result")
    if raw is None:
        raw = next((item.get("text") for item in result.get("content", []) if item.get("type") == "text"), "")
    envelope = json.loads(raw) if isinstance(raw, str) else raw
    if not isinstance(envelope, dict) or envelope.get("ok") is not True or envelope.get("body", {}).get("success") is not True:
        raise ValueError("心力球未返回可验证的数据，请检查本机授权。")
    page = envelope["body"].get("data")
    if not isinstance(page, dict) or not isinstance(page.get("content"), list):
        raise ValueError("心力球分页数据不完整。")
    return page


class Bridge:
    """Use MCP transport directly. Never starts a model turn or exports OAuth tokens."""
    def __enter__(self):
        args = codex_command() + ["app-server", "--listen", "stdio://"]
        config = {"approval_policy": "never", "sandbox_mode": "read-only", "project_doc_max_bytes": 0,
                  "web_search": "disabled", "agents.enabled": False,
                  "mcp_servers": {"xinliqiu": {"url": GATEWAY, "enabled": True, "required": True,
                      "enabled_tools": list(TOOLS), "tools": {tool: {"approval_mode": "approve"} for tool in TOOLS}}}}
        for feature in ("shell_tool", "multi_agent", "apps", "plugins", "hooks", "memories", "browser_use",
                        "computer_use", "image_generation", "skill_search", "tool_suggest"):
            config["features." + feature] = False
        for key, value in config.items():
            args += ["-c", key + "=" + toml(value)]
        self.tmp = tempfile.TemporaryDirectory(prefix="daytrace-read-")
        self.messages = queue.Queue()
        self.sequence = 0
        try:
            self.proc = subprocess.Popen(args, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL,
                                         text=True, cwd=self.tmp.name, start_new_session=True)
        except Exception:
            self.tmp.cleanup()
            raise
        with CHILD_LOCK:
            CHILDREN.add(self.proc)
        self.reader = threading.Thread(target=self.read, daemon=True)
        self.reader.start()
        try:
            self.call("initialize", {"clientInfo": {"name": "daytrace-desktop", "version": "1.0"}, "capabilities": {"experimentalApi": True}})
            self.send({"method": "initialized"})
            self.thread = self.call("thread/start", {"cwd": self.tmp.name, "ephemeral": True, "sandbox": "read-only", "approvalPolicy": "never"})["thread"]["id"]
            return self
        except Exception:
            self.__exit__()
            raise

    def read(self):
        try:
            for line in self.proc.stdout:
                try:
                    self.messages.put(json.loads(line))
                except ValueError:
                    pass
        finally:
            self.messages.put(None)

    def send(self, value):
        self.proc.stdin.write(json.dumps(value) + "\n")
        self.proc.stdin.flush()

    def call(self, method, params):
        self.sequence += 1
        wanted = self.sequence
        self.send({"id": wanted, "method": method, "params": params})
        deadline = time.monotonic() + 90
        while time.monotonic() < deadline:
            try:
                message = self.messages.get(timeout=max(.01, deadline - time.monotonic()))
            except queue.Empty:
                break
            if message is None:
                raise ValueError("本机连接已结束，请检查 Codex 与心力球登录。")
            if message.get("id") == wanted and "method" not in message:
                if "error" in message:
                    raise ValueError("心力球读取失败，请检查本机授权与 Codex 版本。")
                return message["result"]
            if "method" in message and "id" in message:
                self.send({"id": message["id"], "error": {"code": -32601, "message": "Interactive actions are disabled"}})
        raise ValueError("心力球读取超时，稍后自动重试。")

    def recall(self, tool, arguments):
        if tool not in TOOLS:
            raise ValueError("只允许读取全时记与随口记。")
        result = self.call("mcpServer/tool/call", {"threadId": self.thread, "server": "xinliqiu", "tool": tool, "arguments": arguments})
        return parse_page(result)

    def __exit__(self, *unused):
        if self.proc.poll() is None:
            if os.name == "posix":
                try:
                    os.killpg(self.proc.pid, signal.SIGTERM)
                except ProcessLookupError:
                    pass
            else:
                self.proc.terminate()
            try:
                self.proc.wait(timeout=3)
            except subprocess.TimeoutExpired:
                self.proc.kill()
                self.proc.wait()
        self.proc.stdin.close()
        self.proc.stdout.close()
        self.reader.join(timeout=1)
        with CHILD_LOCK:
            CHILDREN.discard(self.proc)
        self.tmp.cleanup()


def redact(value):
    if isinstance(value, dict):
        return {key: redact(item) for key, item in value.items()}
    if isinstance(value, list):
        return [redact(item) for item in value]
    if isinstance(value, str):
        value = re.sub(r"\b(?:dj_|dsc_|sk-)[A-Za-z0-9_-]{12,}\b", "[已隐藏密钥]", value)
        value = re.sub(r"(?<!\d)1[3-9]\d{9}(?!\d)", "[已隐藏手机号]", value)
        value = re.sub(r"[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}", "[已隐藏邮箱]", value)
        return re.sub(r"((?:密码|验证码|api[ _-]?key|access[ _-]?token|refresh[ _-]?token|authorization)\s*[:：=]?\s*)[a-z0-9_./+=-]{4,}", r"\1[已隐藏]", value, flags=re.I)
    return value


def fetch_payload(task, progress, bridge_factory=Bridge):
    pages = {}
    with bridge_factory() as bridge:
        for tool in TOOLS:
            collected = []
            total = count = None
            records = 0
            for index in range(200):
                page = bridge.recall(tool, {"start_time": task["start_time"], "end_time": task["end_time"], "page": index, "size": 20})
                current_total, current_count = page.get("totalPages"), page.get("totalElements")
                if type(current_total) is not int or type(current_count) is not int or not 0 <= current_total <= 200 or current_count < 0:
                    raise ValueError("心力球分页总数无效，原记录已保留。")
                if index == 0:
                    total, count = current_total, current_count
                if (current_total != total or current_count != count or page.get("number") != index
                        or page.get("last") is not (index + 1 == max(1, total))):
                    raise ValueError("读取期间分页发生变化，稍后重新同步。")
                if not isinstance(page.get("content"), list) or any(not isinstance(item, dict) for item in page["content"]):
                    raise ValueError("记录格式不完整，原记录已保留。")
                records += len(page["content"])
                collected.append({key: page[key] for key in ("number", "totalPages", "totalElements", "last", "content")})
                progress(("全时记" if tool == "recall_memory" else "随口记") + "：已读取第 %s / %s 页" % (index + 1, max(1, total)))
                if page["last"]:
                    break
            if records != count or len(collected) != max(1, total):
                raise ValueError("分页条数不一致，原记录已保留。")
            pages[tool] = collected
    payload = {"schemaVersion": 1, "date": task["date"], "start_time": task["start_time"], "end_time": task["end_time"], "pages": redact(pages)}
    raw = json.dumps(payload, ensure_ascii=False, separators=(",", ":")).encode()
    compressed = base64.b64encode(gzip.compress(raw, mtime=0)).decode()
    if len(raw) > 12 * 1024 * 1024 or len(compressed) > 800000:
        raise ValueError("这一天的记录超过单次同步上限，原记录已保留。")
    return compressed


class Sync:
    def __init__(self, config_path=CONFIG_PATH):
        self.config_path = config_path
        self.config = load_config(config_path)
        self.lock = threading.Lock()
        self.wake = threading.Event()
        self.closed = threading.Event()
        self.message = "等待绑定日记" if not self.config else "等待连接网站"
        self.last_success = None
        self.server_today = None
        self.auto_at = 0

    def status(self):
        with self.lock:
            return {"paired": bool(self.config), "server": self.config.get("server", "https://www.funnycafe.top"),
                    "running": self.config.get("running", False), "name": self.config.get("name", "我的电脑"),
                    "message": self.message, "last_success": self.last_success}

    def pair(self, server, key, name):
        server = server_url(server)
        if not isinstance(key, str) or not key.startswith("dj_") or len(key) > 256:
            raise ValueError("请粘贴 dj_ 开头的日记密钥。")
        name = str(name).strip()[:80] or "我的电脑"
        result = api(server, key, "/register", {"name": name})
        token = result.get("connector_token", "")
        if not token.startswith("dsc_"):
            raise ValueError("网站未返回电脑同步凭证。")
        config = {"server": server, "token": token, "name": name, "running": True}
        save_config(config, self.config_path)
        with self.lock:
            self.config = config
            self.message = "已绑定，正在连接心力球"
            self.server_today = None
            self.auto_at = 0
        self.wake.set()

    def set_running(self, enabled):
        with self.lock:
            if not self.config:
                raise ValueError("请先绑定日记。")
            config = dict(self.config, running=enabled)
            save_config(config, self.config_path)
            self.config = config
            self.message = "等待同步" if enabled else "已暂停后续同步，当前任务完成后停止"
        self.wake.set()

    def tick(self):
        with self.lock:
            config = dict(self.config)
        if not config.get("running"):
            return False
        task = None
        try:
            dates = []
            if self.server_today and time.monotonic() >= self.auto_at:
                current = date.fromisoformat(self.server_today)
                dates = [current.isoformat(), (current - timedelta(days=1)).isoformat()]
            response = api(config["server"], config["token"], "/next", {"dates": dates})
            self.server_today = response["server_today"]
            if dates:
                self.auto_at = time.monotonic() + 300
            task = response.get("task")
            if not task:
                self.message = "已连接，等待网页请求；每 5 分钟同步今天与昨天"
                return bool(not dates and self.auto_at == 0)
            path = "/tasks/" + task["id"]
            progress_at = 0

            def progress(stage):
                nonlocal progress_at
                self.message = task["date"] + " · " + stage
                if time.monotonic() - progress_at >= 3:
                    api(config["server"], config["token"], path + "/progress", {"lease": task["lease"], "stage": stage})
                    progress_at = time.monotonic()

            progress("正在读取心力球")
            payload = fetch_payload(task, progress)
            progress("正在上传已读取的记录")
            for attempt in range(3):
                try:
                    api(config["server"], config["token"], path + "/upload", {"lease": task["lease"], "payload_b64": payload})
                    break
                except ApiError as error:
                    if error.status in (401, 404, 409, 413, 422) or attempt == 2:
                        raise
                    self.closed.wait(2 ** attempt)
            self.last_success = time.strftime("%Y-%m-%d %H:%M:%S")
            self.message = task["date"] + " 已同步"
            return True
        except Exception as error:
            # Only controlled messages are displayed. Never echo upstream bodies or credentials.
            self.message = str(error) if isinstance(error, (ValueError, ApiError)) and not isinstance(error, json.JSONDecodeError) else "本机读取未完成，请检查 Codex、心力球登录与网络。"
            if task:
                try:
                    api(config["server"], config["token"], "/tasks/" + task["id"] + "/progress", {"lease": task["lease"], "failed": True})
                except Exception:
                    pass
            if isinstance(error, ApiError) and error.status == 401:
                with self.lock:
                    if self.config.get("token") == config["token"]:
                        self.config = dict(self.config, running=False)
                        save_config(self.config, self.config_path)
            return False

    def run(self):
        while not self.closed.is_set():
            immediate = self.tick()
            self.wake.wait(1 if immediate else 20)
            self.wake.clear()


PAGE = """<!doctype html><html lang="zh-CN"><meta charset="utf-8"><meta name="viewport" content="width=device-width,initial-scale=1">
<title>日迹 · 电脑同步</title><style>
body{background:#101010;color:#f2efe7;font:16px/1.9 system-ui;margin:0}main{max-width:620px;margin:8vh auto;padding:24px}
h1{font:38px/1.4 Georgia,serif}label{display:block;margin-top:18px}input{box-sizing:border-box;width:100%;padding:12px;background:#181818;color:inherit;border:1px solid #555;font:inherit}
button{background:#e23a22;color:white;border:0;padding:12px 18px;font:inherit;cursor:pointer;margin:16px 8px 0 0}button:disabled{opacity:.5}
p,li{color:#aaa}code{color:#ddd;overflow-wrap:anywhere}#status{border-left:2px solid #e23a22;padding:14px;white-space:pre-wrap}small{color:#aaa}
</style><main><small>DAYTRACE / DESKTOP SYNC</small><h1>让电脑留下记录，<br>手机继续回看。</h1>
<p>使用电脑现有的心力球授权读取全时记与随口记，同步到你绑定的日迹。心力球授权保留在本机。点击开始后，每 5 分钟同步今天和昨天，也会处理网页请求的日期。</p>
<div id="status" role="status">正在读取本机状态……</div><form id="pair">
<label>日迹网站地址<input id="server" type="url" value="https://www.funnycafe.top" required></label>
<label>电脑名称<input id="name" value="我的电脑" maxlength="80" required></label>
<label>日记密钥<input id="key" type="password" placeholder="dj_ 开头，与手机上的日记相同" autocomplete="off" required></label>
<button type="submit">绑定并开始同步</button><button type="button" id="toggle" hidden>暂停后续同步</button></form>
<p>首次使用心力球时，在终端执行下面的授权命令，然后按提示完成登录。已有授权可以直接开始。</p>
<code>codex -c 'mcp_servers.xinliqiu.url="https://beluga-server-cn.heytapmobi.com/beluga-mcp/gateway"' mcp login xinliqiu</code>
<p>保持此程序运行。电脑休眠或网络中断时暂停更新，网页可继续使用已同步记录。关闭终端即退出程序，下次运行会继续同步。可以在网页“连接电脑端”里撤销同步凭证。</p>
<small>日记密钥仅用于本次绑定，程序只保存独立的电脑同步凭证。原文在内存中读取，上传前过滤常见密钥、密码、邮箱及手机号。同步文件保存在 ~/.config/daytrace-sync/。</small></main>
<script>
const csrf='__NONCE__';let running=false;const $=s=>document.querySelector(s);
async function call(path,data){const r=await fetch('/api/'+path,{method:data?'POST':'GET',headers:{'Content-Type':'application/json','X-CSRF':csrf},body:data?JSON.stringify(data):undefined});const v=await r.json();if(!r.ok)throw Error(v.message);return v}
async function status(){try{const s=await call('status');running=s.running;$('#status').textContent=s.message+(s.last_success?'\\n最近同步 '+s.last_success:'');$('#toggle').hidden=!s.paired;$('#toggle').textContent=s.running?'暂停后续同步':'继续同步';if(!document.activeElement.closest('form')){$('#server').value=s.server;$('#name').value=s.name}}catch(e){$('#status').textContent=e.message}}
$('#pair').onsubmit=async e=>{e.preventDefault();const b=$('button[type=submit]');b.disabled=true;const key=$('#key').value;$('#key').value='';try{await call('pair',{server:$('#server').value,key,name:$('#name').value});await status()}catch(e){$('#status').textContent=e.message}finally{b.disabled=false}};
$('#toggle').onclick=async()=>{try{await call('running',{enabled:!running});await status()}catch(e){$('#status').textContent=e.message}};status();setInterval(status,4000);
</script></html>"""


def handler(sync, port):
    nonce = secrets.token_urlsafe(32)
    origin = "http://127.0.0.1:" + str(port)
    setup_lock = threading.Lock()

    class Handler(BaseHTTPRequestHandler):
        def log_message(self, *unused):
            pass

        def send_data(self, status, data, content_type="application/json; charset=utf-8"):
            raw = data.encode() if isinstance(data, str) else json.dumps(data, ensure_ascii=False).encode()
            self.send_response(status)
            self.send_header("Content-Type", content_type)
            self.send_header("Content-Length", str(len(raw)))
            self.send_header("Cache-Control", "no-store")
            self.send_header("X-Content-Type-Options", "nosniff")
            self.send_header("Referrer-Policy", "no-referrer")
            self.send_header("Content-Security-Policy", "default-src 'self'; script-src 'unsafe-inline'; style-src 'unsafe-inline'; frame-ancestors 'none'; base-uri 'none'; form-action 'self'")
            self.end_headers()
            self.wfile.write(raw)

        def allowed(self, mutation=False):
            if self.headers.get("Host") != "127.0.0.1:" + str(port):
                return False
            if self.headers.get("Origin") not in (None, origin):
                return False
            if self.path.startswith("/api/") and not hmac.compare_digest(self.headers.get("X-CSRF", ""), nonce):
                return False
            return not mutation or (self.headers.get("Origin") == origin and self.headers.get_content_type() == "application/json")

        def do_GET(self):
            if not self.allowed():
                return self.send_data(403, {"message": "请从本机同步窗口访问。"})
            if self.path == "/":
                return self.send_data(200, PAGE.replace("__NONCE__", nonce), "text/html; charset=utf-8")
            if self.path == "/api/status":
                return self.send_data(200, sync.status())
            self.send_data(404, {"message": "页面不存在"})

        def do_POST(self):
            if not self.allowed(True):
                return self.send_data(403, {"message": "请从本机同步窗口操作。"})
            try:
                size = int(self.headers.get("Content-Length", "0"))
                if not 0 < size <= 16384:
                    raise ValueError("请求格式不正确。")
                data = json.loads(self.rfile.read(size))
                if not isinstance(data, dict):
                    raise ValueError("操作不存在。")
                with setup_lock:
                    if self.path == "/api/pair" and isinstance(data.get("server"), str):
                        sync.pair(data["server"], data["key"], data.get("name", "我的电脑"))
                    elif self.path == "/api/running" and type(data.get("enabled")) is bool:
                        sync.set_running(data["enabled"])
                    else:
                        raise ValueError("操作不存在。")
                self.send_data(200, {"ok": True})
            except (ValueError, KeyError, TypeError, OSError) as error:
                message = str(error) if isinstance(error, ValueError) and not isinstance(error, json.JSONDecodeError) else "操作未完成，请检查输入与本机配置权限。"
                self.send_data(400, {"message": message})

    return Handler


def main():
    parser = argparse.ArgumentParser(description="日迹电脑同步")
    parser.add_argument("--port", type=int, default=4185)
    parser.add_argument("--no-browser", action="store_true")
    arguments = parser.parse_args()
    sync = Sync()
    try:
        server = ThreadingHTTPServer(("127.0.0.1", arguments.port), handler(sync, arguments.port))
    except OSError:
        print("同步窗口端口已被占用，请打开现有窗口或使用 --port 指定其他端口。")
        return
    worker = threading.Thread(target=sync.run, daemon=True)
    worker.start()
    url = "http://127.0.0.1:" + str(arguments.port) + "/"
    print("日迹电脑同步已启动 " + url)
    print("关闭此终端或按 Ctrl+C 停止同步。")
    if not arguments.no_browser:
        webbrowser.open(url)
    def interrupted(signum, frame):
        raise KeyboardInterrupt
    signal.signal(signal.SIGTERM, interrupted)
    if hasattr(signal, "SIGHUP"):
        signal.signal(signal.SIGHUP, interrupted)
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        pass
    finally:
        sync.closed.set()
        sync.wake.set()
        server.server_close()
        with CHILD_LOCK:
            children = list(CHILDREN)
        for process in children:
            if process.poll() is None:
                try:
                    if os.name == "posix":
                        os.killpg(process.pid, signal.SIGTERM)
                    else:
                        process.terminate()
                except ProcessLookupError:
                    pass
        worker.join(timeout=4)


if __name__ == "__main__":
    main()