rime-auto-deploy/custom/rime_ice_wanxiang.py

#!/usr/bin/env python3
# /// script
# requires-python = ">=3.9"
# dependencies = [
#     "ruamel.yaml",
# ]
# ///
# -*- coding: utf-8 -*-

from __future__ import annotations

import argparse
import os
import shutil
import stat
import sys
import uuid
from datetime import datetime
from io import StringIO
from pathlib import Path
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen

try:
    from ruamel.yaml import YAML
    from ruamel.yaml.comments import CommentedMap
    from ruamel.yaml.error import YAMLError
except ModuleNotFoundError:  # 在真正执行操作时再给出友好提示
    YAML = None
    CommentedMap = None
    YAMLError = Exception

DEFAULT_RIME_DIR = Path.home() / "Library" / "Rime"
MODEL_FILENAME = "wanxiang-lts-zh-hans.gram"
CUSTOM_YAML_FILENAME = "rime_ice.custom.yaml"
MODEL_NAME = "wanxiang-lts-zh-hans"

# CNB 优先，GitHub 备用。
CNB_URL = (
    "https://cnb.cool/Mintimate/rime/oh-my-rime/-/releases/download/latest/"
    "wanxiang-lts-zh-hans.gram"
)
GITHUB_URL = (
    "https://github.com/amzxyz/rime-wanxiang-data/releases/download/LTS/"
    "wanxiang-lts-zh-hans.gram"
)

DOWNLOAD_TIMEOUT = 60  # 秒
CHUNK_SIZE = 1024 * 1024

# 集中维护万象模型参数；只补充缺失项，不覆盖用户已有调参。
PATCH_ITEMS = (
    ("grammar/language", MODEL_NAME),
    ("grammar/collocation_max_length", 6),
    ("grammar/collocation_min_length", 2),
    ("grammar/collocation_penalty", -14),
    ("grammar/non_collocation_penalty", -6),
    ("grammar/weak_collocation_penalty", -100),
    ("grammar/rear_penalty", -18),
    ("translator/contextual_suggestions", False),
    ("translator/max_homophones", 8),
    ("translator/max_homographs", 7),
)


class RimeMateError(Exception):
    """脚本可直接向用户报告的错误。"""


class ConfigError(RimeMateError):
    """YAML 配置无效或无法安全修改。"""


class DownloadError(RimeMateError):
    """语言模型下载失败。"""


def _require_ruamel() -> None:
    if YAML is None:
        raise RimeMateError(
            "缺少第三方依赖 ruamel.yaml。\n"
            "请确认当前使用的是 pyenv 安装的 Python 3.13:\n"
            "    which python3\n"
            "    python3 --version\n"
            "    pyenv global 3.13\n"
            "然后在当前 Python 环境安装依赖:\n"
            "    python3 -m pip install ruamel.yaml\n"
            "详细步骤见 README.python.md。"
        )


def _new_yaml():
    yaml = YAML(typ="rt")
    yaml.preserve_quotes = True
    yaml.width = 4096
    yaml.allow_unicode = True
    yaml.indent(mapping=2, sequence=4, offset=2)
    return yaml


def _ensure_macos() -> None:
    if sys.platform != "darwin":
        raise RimeMateError("本脚本只支持 macOS。")


def _ensure_rime_dir(rime_dir: Path) -> None:
    if not rime_dir.is_dir():
        raise RimeMateError(f"Rime 配置目录不存在: {rime_dir}")


def _load_config(path: Path):
    """读取配置并验证根节点、patch 节点，返回 round-trip 映射。"""
    if not path.exists():
        return CommentedMap()

    try:
        text = path.read_text(encoding="utf-8")
    except (OSError, UnicodeError) as exc:
        raise ConfigError(f"无法读取配置文件 {path}: {exc}") from exc

    # 空文件视为尚未写入任何配置。
    if not text.strip():
        return CommentedMap()

    yaml = _new_yaml()
    try:
        data = yaml.load(text)
    except YAMLError as exc:
        raise ConfigError(f"配置文件不是合法的 YAML: {path}\n{exc}") from exc

    if data is None:
        return CommentedMap()
    if not isinstance(data, CommentedMap):
        raise ConfigError(
            f"配置文件根节点必须是映射（键值对），但 {path} 的根节点是 "
            f"{type(data).__name__}。为避免破坏配置，已停止。"
        )

    patch = data.get("patch")
    if patch is not None and not isinstance(patch, CommentedMap):
        raise ConfigError(
            f"配置文件中的 patch 必须是映射（键值对），但实际是 "
            f"{type(patch).__name__}。为避免破坏配置，已停止。"
        )
    return data


def _ensure_no_other_model(patch, action: str) -> None:
    """发现其他语言模型的明确设置时保守中止，避免覆盖或误删共享调参。"""
    if patch is None or "grammar/language" not in patch:
        return

    value = patch["grammar/language"]
    if value is None:
        return
    language = str(value).strip()
    if not language or language == MODEL_NAME:
        return

    raise ConfigError(
        "检测到 grammar/language 已明确设置为其他语言模型"
        f"（{language}）。为避免覆盖其他模型的设置，已停止{action}。"
    )


def _dump_config(root) -> bytes:
    if len(root) == 0:
        # 保留配置文件本身，不删除用户文件。
        return b"{}\n"

    yaml = _new_yaml()
    buffer = StringIO()
    yaml.dump(root, buffer)
    return buffer.getvalue().encode("utf-8")


def _backup_file(path: Path) -> Path:
    stamp = datetime.now().strftime("%Y%m%d-%H%M%S")
    backup = path.with_name(f"{path.name}.bak-{stamp}-{uuid.uuid4().hex[:8]}")
    try:
        shutil.copy2(path, backup)
    except OSError as exc:
        raise ConfigError(f"备份配置文件失败，未做任何修改: {exc}") from exc
    return backup


def _atomic_write(path: Path, data: bytes, mode: int) -> None:
    temp_path = path.with_name(f".{path.name}.tmp-{uuid.uuid4().hex}")
    try:
        fd = os.open(temp_path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
        with os.fdopen(fd, "wb") as handle:
            handle.write(data)
            handle.flush()
            os.fsync(handle.fileno())
        os.chmod(temp_path, mode)
        os.replace(temp_path, path)
    except OSError as exc:
        try:
            temp_path.unlink()
        except OSError:
            pass
        raise ConfigError(f"写入配置文件失败: {exc}") from exc


def _write_config(root, path: Path) -> None:
    """备份后原子写入配置，并保留原文件权限。"""
    mode = 0o644
    if path.exists():
        try:
            mode = stat.S_IMODE(path.stat().st_mode)
        except OSError as exc:
            raise ConfigError(f"无法读取配置文件权限: {exc}") from exc
        backup = _backup_file(path)
        print(f"已备份原配置: {backup}")

    _atomic_write(path, _dump_config(root), mode)


def _safe_unlink(path: Path) -> None:
    try:
        path.unlink()
    except FileNotFoundError:
        pass
    except OSError:
        pass


def _validate_payload(head: bytes) -> None:
    stripped = head.lstrip().lower()
    if stripped.startswith(b"<") or b"<html" in stripped[:256] or b"<!doctype" in stripped[:256]:
        raise DownloadError("下载内容疑似 HTML 错误页面，而不是语言模型文件")
    if stripped.startswith(b"version https://git-lfs"):
        raise DownloadError("下载内容是 Git LFS 指针文件，而不是语言模型文件")


def _download_to(url: str, temp_path: Path, timeout: float) -> None:
    try:
        _download_to_inner(url, temp_path, timeout)
    except BaseException:
        _safe_unlink(temp_path)
        raise


def _download_to_inner(url: str, temp_path: Path, timeout: float) -> None:
    request = Request(
        url,
        headers={"User-Agent": "rime-ice-wanxiang/1.0 (+https://github.com/hunter-ji/rime-mate)"},
    )

    try:
        response = urlopen(request, timeout=timeout)
    except HTTPError as exc:
        raise DownloadError(f"HTTP 错误 {exc.code}: {exc.reason}") from exc
    except (URLError, TimeoutError, OSError) as exc:
        raise DownloadError(f"网络请求失败: {exc}") from exc

    with response:
        status = getattr(response, "status", None)
        if status is None:
            status = response.getcode()
        if status != 200:
            raise DownloadError(f"服务器返回异常状态码: {status}")

        expected_length = None
        raw_length = response.headers.get("Content-Length")
        if raw_length:
            try:
                expected_length = int(raw_length)
            except (TypeError, ValueError):
                expected_length = None

        total = 0
        head = b""
        try:
            fd = os.open(temp_path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o644)
        except OSError as exc:
            raise DownloadError(f"无法创建临时文件: {exc}") from exc

        try:
            handle = os.fdopen(fd, "wb")
        except OSError as exc:
            os.close(fd)
            raise DownloadError(f"无法打开临时文件: {exc}") from exc

        with handle:
            while True:
                try:
                    chunk = response.read(CHUNK_SIZE)
                except (URLError, TimeoutError, OSError) as exc:
                    raise DownloadError(f"下载中断: {exc}") from exc
                if not chunk:
                    break
                if len(head) < 512:
                    head += chunk[: 512 - len(head)]
                try:
                    handle.write(chunk)
                except OSError as exc:
                    raise DownloadError(f"写入临时文件失败: {exc}") from exc
                total += len(chunk)
            try:
                handle.flush()
                os.fsync(handle.fileno())
            except OSError as exc:
                raise DownloadError(f"写入临时文件失败: {exc}") from exc

        if total == 0:
            raise DownloadError("下载内容为空")
        if expected_length is not None and total != expected_length:
            raise DownloadError(
                f"下载内容大小与 Content-Length 不一致（{total} != {expected_length}）"
            )
        _validate_payload(head)


def _download_model(rime_dir: Path, timeout: float) -> Path:
    """依次尝试 CNB 和 GitHub，成功后原子替换旧模型。"""
    model_path = rime_dir / MODEL_FILENAME
    errors = []

    for label, url in (("CNB", CNB_URL), ("GitHub", GITHUB_URL)):
        temp_path = rime_dir / f".{MODEL_FILENAME}.tmp-{uuid.uuid4().hex}"
        try:
            print(f"尝试从 {label} 下载语言模型...")
            _download_to(url, temp_path, timeout)
            os.replace(temp_path, model_path)
            return model_path
        except (RimeMateError, OSError) as exc:
            _safe_unlink(temp_path)
            errors.append(f"{label}: {exc}")
            print(f"{label} 下载失败: {exc}", file=sys.stderr)

    raise DownloadError(
        "所有下载源均失败，已保留原有语言模型。\n" + "\n".join(errors)
    )


def install(rime_dir: Path, timeout: float = DOWNLOAD_TIMEOUT) -> None:
    """安装或更新万象语言模型，只补充缺失的 patch 项。"""
    _require_ruamel()
    rime_dir = Path(rime_dir).expanduser()
    _ensure_rime_dir(rime_dir)

    print(f"Rime 配置目录: {rime_dir}")
    config_path = rime_dir / CUSTOM_YAML_FILENAME

    # 先解析并验证配置，再下载；避免配置有误时白白下载或破坏旧模型。
    root = _load_config(config_path)
    patch = root.get("patch")
    _ensure_no_other_model(patch, "安装")

    if patch is None:
        patch = CommentedMap()
        root["patch"] = patch

    added = []
    for key, value in PATCH_ITEMS:
        if key not in patch:
            if not added and len(patch) > 0:
                # 无论 patch 中已有何种配置，新增万象配置前始终留一行空白。
                patch.yaml_set_comment_before_after_key(key, before="\n")
            patch[key] = value
            added.append(key)

    model_path = _download_model(rime_dir, timeout)
    print(f"语言模型已保存: {model_path}")

    if added:
        _write_config(root, config_path)
        print("配置已更新，新增: " + "、".join(added))
    else:
        print("配置中已包含全部万象模型设置，无需修改。")

    _print_redeploy_hint()


def uninstall(rime_dir: Path) -> None:
    """清理万象 patch 项和目标模型文件。"""
    _require_ruamel()
    rime_dir = Path(rime_dir).expanduser()
    _ensure_rime_dir(rime_dir)

    print(f"Rime 配置目录: {rime_dir}")
    config_path = rime_dir / CUSTOM_YAML_FILENAME
    model_path = rime_dir / MODEL_FILENAME

    # 先解析配置，只有配置写入成功后才删除模型文件。
    # 卸载时按键名清理整组万象参数，不检查当前值；用户手动调整过的
    # 万象参数也必须一并删除。
    root = _load_config(config_path)
    patch = root.get("patch")

    removed = []
    if patch is not None:
        for key, _ in PATCH_ITEMS:
            if key in patch:
                del patch[key]
                removed.append(key)
        if removed and len(patch) == 0:
            del root["patch"]

    if removed:
        # 只有配置备份并成功写入后，才进入删除模型文件的步骤。
        _write_config(root, config_path)
        print("配置已清理，移除: " + "、".join(removed))
    else:
        print("配置中未发现万象模型相关设置，无需修改。")

    if model_path.exists() or model_path.is_symlink():
        try:
            model_path.unlink()
        except OSError as exc:
            raise RimeMateError(
                f"配置已清理，但删除语言模型文件失败: {exc}"
            ) from exc
        print(f"语言模型已删除: {model_path}")
    else:
        print("语言模型文件不存在，无需删除。")

    _print_redeploy_hint()


def _print_redeploy_hint() -> None:
    print()
    print("请手动重新部署鼠须管（菜单栏图标 -> 重新部署，或按 Control+Option+`）使配置生效。")
    print("本脚本不会自动触发部署。")


def _confirm(prompt: str) -> bool:
    try:
        answer = input(prompt).strip().lower()
    except EOFError:
        print()
        return False
    return answer in ("y", "yes")


def interactive_menu(rime_dir: Path) -> int:
    print("Rime Ice 万象语言模型管理")
    print(f"Rime 配置目录: {rime_dir}")

    while True:
        print()
        print("1. 安装/更新")
        print("2. 卸载")
        print("3. 退出")
        try:
            choice = input("请选择 [1-3]: ").strip()
        except EOFError:
            print()
            return 0

        if choice == "1":
            try:
                install(rime_dir)
            except RimeMateError as exc:
                print(f"错误: {exc}", file=sys.stderr)
        elif choice == "2":
            if _confirm("确认卸载万象语言模型并清理相关配置？[y/N] "):
                try:
                    uninstall(rime_dir)
                except RimeMateError as exc:
                    print(f"错误: {exc}", file=sys.stderr)
            else:
                print("已取消。")
        elif choice == "3":
            return 0
        else:
            print("请输入 1、2 或 3。")


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="在 macOS 上安装、更新或卸载 Rime Ice 万象语言模型。",
        epilog="不指定操作时显示中文菜单。仅使用 ruamel.yaml 一个第三方依赖。",
    )
    parser.add_argument(
        "action",
        nargs="?",
        choices=("install", "update", "uninstall"),
        help="install/update 安装或更新；uninstall 卸载；省略则显示菜单",
    )
    parser.add_argument(
        "--rime-dir",
        type=Path,
        default=DEFAULT_RIME_DIR,
        help=f"Rime 配置目录（默认: {DEFAULT_RIME_DIR}，主要用于测试）",
    )
    return parser


def main(argv=None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)

    try:
        _ensure_macos()
        _require_ruamel()
        rime_dir = args.rime_dir.expanduser()
        _ensure_rime_dir(rime_dir)

        if args.action is None:
            return interactive_menu(rime_dir)
        if args.action in ("install", "update"):
            install(rime_dir)
        else:
            uninstall(rime_dir)
        return 0
    except RimeMateError as exc:
        print(f"错误: {exc}", file=sys.stderr)
        return 1
    except KeyboardInterrupt:
        print("\n已取消。", file=sys.stderr)
        return 130


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