# -*- coding: utf-8 -*-
"""StaSh 一键安装 + 优化脚本（Pythonista 专用）。

功能：
  1. 读取云端配置 download_config.json，选择「加速 / 不加速」下载源
  2. 检查 update.json，提示是否有新版脚本
  3. 执行 getstash.py 安装 StaSh（走 gitproxy 代理）
  4. 优化：补装依赖（pyte/wcwidth）+ 修复启动闪退 + 修复 pip

补丁是否已打，统一读取 stash 目录下的 patch_state.json 判断。
配置只从云端/本地 download_config.json 读取，不内置任何 JSON。
"""

import io
import json
import os
import re
import shutil
import sys
import tempfile
import zipfile

import requests

# ───────────────────────── 常量 ─────────────────────────

SCRIPT_VERSION = 4

# 本地配置缓存文件名（与云端 download_config.json 对应）
CONFIG_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'download_config.json')

# 云端配置与更新信息（EdgeOne Pages 项目 stash-config）
REMOTE_CONFIG_URL = 'https://stash-config.edgeone.dev/download_config.json'
REMOTE_UPDATE_URL = 'https://stash-config.edgeone.dev/update.json'
TIMEOUT = 15

# gitproxy 代理前缀，用于把 getstash.py 内部直连 github 的地址换成代理
PROXY = 'https://api.gitproxy.dev/'

# StaSh 运行所需依赖（pyte 及其依赖 wcwidth）
DEPS = (
    ('selectel/pyte', 'pyte', '0.8.2'),
    ('jquast/wcwidth', 'wcwidth', '0.2.13'),
)

# 新版 pip.py（修复 python3 下 "unknown command: None" 问题）
PIP_URL = 'https://raw.githubusercontent.com/ywangd/stash/master/bin/pip.py'

# 补丁状态 JSON（stash 目录下，字段如 {"pip": true, "launch_stash": true}）
STASH_DIR = os.path.expanduser('~/Documents/site-packages/stash')
PATCH_STATE_PATH = os.path.join(STASH_DIR, 'patch_state.json')

# 下载源选项：1 = 加速，2 = 不加速
SOURCES = {"1": "fast", "2": "normal"}
BRANCHES = {"1": "master", "2": "dev"}


# ───────────────────────── 界面美化 ─────────────────────────

# 仅在终端（如 StaSh）启用颜色，Pythonista 交互控制台自动降级为纯文本
USE_COLOR = hasattr(sys.stdout, 'isatty') and sys.stdout.isatty()

C = {
    'r': '\033[91m',  # 红
    'g': '\033[92m',  # 绿
    'y': '\033[93m',  # 黄
    'b': '\033[94m',  # 蓝
    'm': '\033[95m',  # 品红
    'c': '\033[96m',  # 青
    'd': '\033[2m',   # 暗
    '0': '\033[0m',   # 复位
}


def c(code, text):
    """颜色包装（无颜色环境原样返回）"""
    return ('%s%s%s' % (C.get(code, ''), text, C['0'])) if USE_COLOR else text


def line(ch='─', n=44):
    return ch * n


def title():
    print()
    print(c('c', '┌' + line('─', 42) + '┐'))
    print(c('c', '│') + c('b', '  🐚  StaSh 安装程序').ljust(46) + c('c', '│'))
    print(c('c', '│') + c('d', '  Pythonista 终端环境一键部署').ljust(46) + c('c', '│'))
    print(c('c', '└' + line('─', 42) + '┘'))
    print()


def ok(msg):
    print(c('g', '  [✓] ') + msg)


def fail(msg):
    print(c('r', '  [✗] ') + msg)


def info(msg):
    print(c('b', '  [→] ') + msg)


def ask(msg):
    return input(c('y', '  [?] ') + msg).strip()


def banner(text):
    print(c('m', '  ── ' + text + ' ' + line('─', 40 - len(text))))


def menu(prompt, options, default=None):
    """循环菜单，直到输入合法选项"""
    print()
    for opt in options:
        print('     ' + c('c', '[' + opt[0] + ']') + '  ' + opt[1])
    while True:
        raw = input(c('y', prompt)).strip().lower()
        if not raw and default is not None:
            raw = default
        for opt in options:
            if raw == opt[0].lower():
                return opt[0]
        print(c('r', '     ⚠ 无效输入，请重新选择。'))


# ───────────────────────── 配置加载 ─────────────────────────

def load_config():
    """读取配置：云端优先（版本更高则更新本地），本地次之。

    两者都不可用时直接报错退出，不使用任何内置 JSON。
    """
    info('正在读取配置...')

    local = None
    if os.path.exists(CONFIG_PATH):
        try:
            local = json.load(open(CONFIG_PATH, 'r', encoding='utf-8'))
        except (IOError, ValueError):
            pass

    remote = None
    if REMOTE_CONFIG_URL:
        try:
            remote = requests.get(REMOTE_CONFIG_URL, timeout=TIMEOUT).json()
        except Exception:
            remote = None

    # 云端版本更高则更新本地
    if remote and remote.get('version', 0) > (local or {}).get('version', 0):
        try:
            json.dump(remote, open(CONFIG_PATH, 'w', encoding='utf-8'),
                      ensure_ascii=False, indent=2)
        except (IOError, OSError):
            pass
        ok('检测到云端更新，已同步本地配置 (v%s)' % remote.get('version', 0))
        return remote

    if local:
        ok('本地配置就绪 (v%s)' % local.get('version', 0))
        return local

    if remote:
        ok('已下载云端配置 (v%s)' % remote.get('version', 0))
        return remote

    fail('无法获取配置：云端不可达，且本地无 download_config.json。')
    fail('请检查网络后重试，或将 download_config.json 与本脚本放在同目录。')
    raise SystemExit(1)


CONFIG = load_config()


# ───────────────────────── 更新检查 ─────────────────────────

def check_update():
    """读取 update.json，若云端脚本版本更高则提示。"""
    if not REMOTE_UPDATE_URL:
        return
    try:
        upd = requests.get(REMOTE_UPDATE_URL, timeout=TIMEOUT).json()
    except Exception:
        return

    ver = upd.get('version', 0)
    if ver > SCRIPT_VERSION:
        print()
        info('发现新版脚本 (v%s)，当前 (v%s)。' % (ver, SCRIPT_VERSION))
        note = upd.get('note')
        if note:
            info('更新说明：%s' % note)
        for ch in upd.get('changes', []):
            info('  · %s' % ch)
        print()


# ───────────────────────── 补丁状态 ─────────────────────────

def load_patch_state():
    """读取补丁状态 JSON，文件不存在或损坏返回空 dict。"""
    if os.path.exists(PATCH_STATE_PATH):
        try:
            with io.open(PATCH_STATE_PATH, 'r', encoding='utf-8') as f:
                return json.load(f)
        except (IOError, ValueError):
            pass
    return {}


def save_patch_state(state):
    """写回补丁状态 JSON。"""
    try:
        with io.open(PATCH_STATE_PATH, 'w', encoding='utf-8') as f:
            json.dump(state, f, ensure_ascii=False, indent=2)
    except IOError as e:
        fail('写入补丁状态 JSON 失败（%s）' % e)


def is_patched(key):
    """判断某个补丁是否已打（以 JSON 中是否存在该键为准）。"""
    return key in load_patch_state()


def mark_patched(key):
    """标记某个补丁已打。"""
    state = load_patch_state()
    state[key] = True
    save_patch_state(state)


# ───────────────────────── 依赖安装 ─────────────────────────

def install_deps():
    banner('补装依赖')
    site = os.path.expanduser('~/Documents/site-packages')
    os.makedirs(site, exist_ok=True)

    for repo, pkg, tag in DEPS:
        try:
            __import__(pkg)
            ok('%s 已存在，跳过' % pkg)
            continue
        except ImportError:
            pass

        url = PROXY + 'https://github.com/%s/archive/refs/tags/%s.zip' % (repo, tag)
        info('安装 %s %s ...' % (pkg, tag))
        try:
            r = requests.get(url, timeout=120)
            r.raise_for_status()
        except Exception as e:
            fail('下载 %s 失败（%s）' % (pkg, e))
            continue

        z = zipfile.ZipFile(io.BytesIO(r.content))
        tmp = tempfile.mkdtemp()
        z.extractall(tmp)

        src = None
        for root, _dirs, _files in os.walk(tmp):
            if os.path.basename(root) == pkg:
                src = root
                break
        if not src:
            fail('未找到包目录 %s' % pkg)
            continue

        dst = os.path.join(site, pkg)
        if os.path.exists(dst):
            shutil.rmtree(dst)
        shutil.copytree(src, dst)
        ok('%s 安装完成' % pkg)


# ───────────────────────── 启动崩溃修复 ─────────────────────────

def patch_stash_launch():
    """修复 iOS 15+ / Pythonista 3.4 上 StaSh 启动即崩溃的问题。

    崩溃信息：
      Modifications to the layout engine must not be performed from a
      background thread after it has been accessed from the main thread.

    修法（社区 issue #496 验证有效）：用 @ui.in_background 包裹启动入口，
    让 Pythonista 正确调度线程，避免与内部渲染线程竞态。

    判断依据：patch_state.json 中是否存在 launch_stash 键。
    """
    banner('修复启动崩溃')

    if is_patched('launch_stash'):
        ok('启动闪退补丁已打过（JSON 记录），跳过')
        return

    targets = [
        os.path.expanduser('~/Documents/site-packages/stash/__main__.py'),
        os.path.expanduser('~/Documents/launch_stash.py'),
    ]

    hit = False
    for path in targets:
        if not os.path.exists(path):
            continue
        try:
            with io.open(path, 'r', encoding='utf-8') as f:
                code = f.read()
        except (IOError, OSError):
            continue

        # 清理历史补丁残留，保证从干净状态重打
        code = _clean_launch_patches(code)

        m = re.search(r'^([ \t]*)_stash\.launch\(ns\.command\)[ \t]*$',
                      code, re.MULTILINE)
        if not m:
            continue

        indent = m.group(1)
        block = (
            indent + 'import ui\n'
            '\n'
            + indent + '@ui.in_background\n'
            + indent + 'def _launch_in_background(*args):\n'
            + indent + '    _stash.launch(*args)\n'
            '\n'
            + indent + '_launch_in_background(ns.command)\n'
        )

        code = code[:m.start()] + block + code[m.end():]

        with io.open(path, 'w', encoding='utf-8') as f:
            f.write(code)
        ok('已修复：%s' % path)
        hit = True

    if hit:
        mark_patched('launch_stash')
    else:
        info('未发现需修复的入口（可能已是最新版，无需处理）')


def _clean_launch_patches(code):
    """清掉历史补丁残留，恢复到干净状态。"""
    patterns = (
        # on_main_thread 版本
        r'[ \t]*from objc_util import on_main_thread\s*\n'
        r'\s*\n'
        r'[ \t]*@on_main_thread\s*\n'
        r'[ \t]*def _launch_main_thread\(\*args\):\s*\n'
        r'[ \t]*_stash\.launch\(\*args\)\s*\n'
        r'\s*\n'
        r'[ \t]*_launch_main_thread\(ns\.command\)\s*\n',

        # 旧 _launch_main_thread + ui.in_background 版本
        r'[ \t]*@ui\.in_background\s*\n'
        r'[ \t]*def _launch_main_thread\(\*args\):\s*\n'
        r'[ \t]*_stash\.launch\(\*args\)\s*\n'
        r'\s*\n'
        r'[ \t]*_launch_main_thread\(ns\.command\)\s*\n',

        # _launch_in_background 版本
        r'[ \t]*@ui\.in_background\s*\n'
        r'[ \t]*def _launch_in_background\(\*args\):\s*\n'
        r'[ \t]*_stash\.launch\(\*args\)\s*\n'
        r'\s*\n'
        r'[ \t]*_launch_in_background\(ns\.command\)\s*\n',
    )
    for p in patterns:
        code = re.sub(p, '', code, flags=re.MULTILINE)
    # 移除残留的 import ui（仅顶格那行）
    code = re.sub(r'^import ui\s*\n', '', code, flags=re.MULTILINE)
    return code


# ───────────────────────── pip 修复 ─────────────────────────

def patch_pip():
    """修复 StaSh 自带 pip.py 在 python3 下报 "unknown command: None" 的问题。

    判断依据：patch_state.json 中是否存在 pip 键。
    """
    banner('修复 pip')

    if is_patched('pip'):
        ok('pip 修复补丁已打过（JSON 记录），跳过')
        return

    target = os.path.join(STASH_DIR, 'bin', 'pip.py')
    if not os.path.exists(target):
        info('未找到 %s，跳过' % target)
        return

    # 新版 pip.py 要求 python 3.10+
    if sys.version_info[:2] < (3, 10):
        info('当前 Python %d.%d 低于 3.10，跳过（新版 pip.py 不支持）'
             % sys.version_info[:2])
        return

    url = PROXY + PIP_URL
    info('下载新版 pip.py ...')
    try:
        r = requests.get(url, timeout=60)
        r.raise_for_status()
    except Exception as e:
        fail('下载新版 pip.py 失败（%s）' % e)
        return

    code = r.text
    if 'distutils' in code:
        fail('下载到的 pip.py 疑似旧版，跳过覆盖')
        return

    with io.open(target, 'w', encoding='utf-8') as f:
        f.write(code)

    mark_patched('pip')
    ok('pip 修复完成')


# ───────────────────────── 安装流程 ─────────────────────────

def install(source, branch):
    url = CONFIG.get(source, {}).get(branch)
    if not url:
        fail('该版本暂无下载地址，请换一个版本试试。')
        return

    banner('开始安装')
    info('下载安装脚本...')
    try:
        r = requests.get(url, timeout=TIMEOUT)
        r.raise_for_status()
    except Exception as e:
        fail('下载安装脚本失败（%s）\n     地址：%s' % (e, url))
        return
    ok('安装脚本下载完成')

    code = r.text
    if branch == 'dev':
        code = code.replace('master', 'dev')

    # 关键：把 getstash.py 内部直连 github 的下载地址换成 gitproxy 代理，避免被墙断连
    code = code.replace('https://github.com/', PROXY + 'https://github.com/')
    code = code.replace('https://raw.githubusercontent.com/',
                        PROXY + 'https://raw.githubusercontent.com/')

    info('正在执行安装（请稍候）...')
    try:
        exec(code, {'__name__': '__main__'})
    except Exception as e:
        fail('执行安装脚本出错（%s）' % e)
        return
    ok('StaSh 安装完成')

    # 修复启动崩溃 + 修复 pip
    patch_stash_launch()
    patch_pip()

    # 询问是否补装依赖
    print()
    opt = ask('是否补装运行依赖（pyte/wcwidth，推荐）？[Y/n] ')
    if opt.lower() in ('', 'y', 'yes'):
        install_deps()
    else:
        info('已跳过依赖补装')

    print()
    print(c('g', '  ════════════════════════════════════'))
    print(c('g', '  🎉 全部完成！'))
    print(c('d', '  请完全退出并重启 Pythonista，'))
    print(c('d', '  然后运行 launch_stash.py 启动终端。'))
    print(c('g', '  ════════════════════════════════════'))
    print()


def main():
    title()
    check_update()

    source = menu(
        ' 请选择下载源：',
        [('1', '加速 —— 走 gitproxy 代理，速度快，推荐'), ('2', '不加速 —— 直连 GitHub，可能较慢')],
        default='1',
    )
    branch = menu(
        ' 请选择版本 / Branch：',
        [('1', '稳定版 master —— 稳定'), ('2', '开发版 dev —— 兼容新版，推荐')],
        default='2',
    )
    install(SOURCES[source], BRANCHES[branch])


if __name__ == '__main__':
    main()