#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Codex 语音副驾驶 —— 本机 Codex 桥接客户端
==========================================

在**你自己的电脑**上运行这个脚本，它会主动连到服务器的 WebSocket 桥接端点，
把"干活"的任务交给你本机的 Codex CLI 执行，执行过程与结果实时回传服务器，
于是语音通话里那个副驾驶指挥的就是你电脑上的 Codex。

特点：
* 只用标准库（Windows / macOS / Linux 都能跑，Python 3.6+）
* 心跳保活：定时发心跳，收到服务器 ping 立刻回 pong；超时自动重连（指数退避）
* 支持打断：服务器说停，本机立刻杀掉正在跑的 codex 进程
* 断线不丢状态：重连后自动重新登记设备信息

用法（在你自己的电脑上）：
    python codex-voice-bridge.py --server wss://ai.jicreate.cn --token <你的登录token> --workspace D:\\codex-work

参数：
    --server     服务器地址，例如 wss://ai.jicreate.cn 或 ws://127.0.0.1:8188
    --token      登录 token（网页里 localStorage.token，或语音页"本机 Codex"帮助里复制）
    --workspace  本机 Codex 的工作目录（默认 ./codex-work）
    --codex      codex 可执行文件路径（默认自动在 PATH 里找 codex / codex.cmd）
    --name       设备名（默认取主机名）
    --heartbeat  心跳间隔秒数（默认 10）
    --env        给本机 codex 追加环境变量，可重复，例如 --env DEEPSEEK_API_KEY=sk-xxx
    --insecure   跳过 TLS 证书校验（自签证书时才用）

注意：本机需要已经装好并配置好 Codex CLI（直接敲 codex 能跑起来）。
     如果本机 codex 依赖某个环境变量，用 --env 传给它，例如 --env DEEPSEEK_API_KEY=sk-xxx
"""

import argparse
import base64
import json
import os
import platform
import queue
import shutil
import socket
import ssl
import struct
import subprocess
import sys
import threading
import time
import uuid

DEFAULT_PORT = {'ws': 80, 'wss': 443}
BRIDGE_PATH = '/ws/codex-bridge'


# ====================== WebSocket 客户端（标准库手写） ======================

class WsClient:
    """极简 WebSocket 客户端：文本帧 + 心跳 + ping/pong，够用且零依赖。"""

    def __init__(self, url, insecure=False, on_text=None, on_open=None, on_close=None):
        self.url = url
        self.insecure = insecure
        self.on_text = on_text
        self.on_open = on_open
        self.on_close = on_close
        self.sock = None
        self.connected = False
        self.last_inbound = 0.0
        self._send_lock = threading.Lock()
        self._recv_buffer = b''

    # ---------- 连接 ----------
    def connect(self, timeout=15):
        parsed = parse_url(self.url)
        scheme, host, port, path = parsed
        raw = socket.create_connection((host, port), timeout=timeout)
        if scheme == 'wss':
            context = ssl.create_default_context()
            if self.insecure:
                context.check_hostname = False
                context.verify_mode = ssl.CERT_NONE
            raw = context.wrap_socket(raw, server_hostname=host)
        raw.settimeout(1.0)
        key = base64.b64encode(os.urandom(16)).decode()
        request = (
            'GET %s HTTP/1.1\r\n'
            'Host: %s:%d\r\n'
            'Upgrade: websocket\r\n'
            'Connection: Upgrade\r\n'
            'Sec-WebSocket-Key: %s\r\n'
            'Sec-WebSocket-Version: 13\r\n\r\n'
        ) % (path, host, port, key)
        raw.sendall(request.encode())
        head = b''
        while b'\r\n\r\n' not in head:
            chunk = raw.recv(4096)
            if not chunk:
                raise IOError('服务器在握手阶段就断开了')
            head += chunk
        status_line = head.split(b'\r\n')[0].decode('utf-8', 'ignore')
        if '101' not in status_line:
            raise IOError('WebSocket 握手失败: ' + status_line)
        self.sock = raw
        self.connected = True
        self.last_inbound = time.time()
        self._recv_buffer = head.split(b'\r\n\r\n', 1)[1]
        if self.on_open:
            self.on_open()
        return True

    def close(self):
        self.connected = False
        sock = self.sock
        self.sock = None
        if sock:
            try:
                sock.close()
            except Exception:
                pass

    # ---------- 发送 ----------
    def send_json(self, payload):
        return self._send_frame(0x1, json.dumps(payload, ensure_ascii=False).encode('utf-8'))

    def _send_frame(self, opcode, data):
        if not self.sock:
            return False
        header = bytearray([0x80 | opcode])
        mask = os.urandom(4)
        length = len(data)
        if length < 126:
            header.append(0x80 | length)
        elif length < 65536:
            header.append(0x80 | 126)
            header += struct.pack('>H', length)
        else:
            header.append(0x80 | 127)
            header += struct.pack('>Q', length)
        header += mask
        masked = bytes(b ^ mask[i % 4] for i, b in enumerate(data))
        try:
            with self._send_lock:
                self.sock.sendall(bytes(header) + masked)
            return True
        except Exception:
            self.connected = False
            return False

    # ---------- 接收 ----------
    def pump(self, max_seconds=1.0):
        """收一会儿消息；返回收到的文本消息条数。"""
        if not self.sock:
            return 0
        count = 0
        end = time.time() + max_seconds
        while time.time() < end:
            try:
                chunk = self.sock.recv(65536)
            except socket.timeout:
                continue
            except Exception:
                self.connected = False
                return count
            if not chunk:
                self.connected = False
                return count
            self.last_inbound = time.time()
            self._recv_buffer += chunk
            while True:
                frame = self._pop_frame()
                if frame is None:
                    break
                opcode, payload = frame
                if opcode == 0x1:
                    count += 1
                    if self.on_text:
                        try:
                            self.on_text(payload.decode('utf-8', 'ignore'))
                        except Exception as error:
                            log('处理服务端消息出错: %s' % error)
                elif opcode == 0x9:
                    self._send_frame(0xA, payload)          # ping → pong
                elif opcode == 0xA:
                    pass                                     # 服务器回了 pong
                elif opcode == 0x8:
                    self.connected = False
                    if self.on_close:
                        self.on_close(payload)
                    return count
        return count

    def _pop_frame(self):
        buf = self._recv_buffer
        if len(buf) < 2:
            return None
        opcode = buf[0] & 0x0F
        length = buf[1] & 0x7F
        index = 2
        if length == 126:
            if len(buf) < 4:
                return None
            length = struct.unpack('>H', buf[2:4])[0]
            index = 4
        elif length == 127:
            if len(buf) < 10:
                return None
            length = struct.unpack('>Q', buf[2:10])[0]
            index = 10
        if len(buf) < index + length:
            return None
        payload = buf[index:index + length]
        self._recv_buffer = buf[index + length:]
        return opcode, payload


def parse_url(url):
    """解析 ws/wss 地址；注意 http://host:8188?token=x 这种"没有 path 只有 query"的写法。"""
    if '://' not in url:
        url = 'ws://' + url
    scheme, rest = url.split('://', 1)
    scheme = scheme.lower()
    query = ''
    if '?' in rest:
        rest, _, query = rest.partition('?')
    rest = rest.rstrip('/')
    if '/' in rest:
        hostport, path = rest.split('/', 1)
        path = '/' + path
    else:
        hostport, path = rest, '/'
    if ':' in hostport:
        host, port_text = hostport.rsplit(':', 1)
        port = int(port_text)
    else:
        host, port = hostport, DEFAULT_PORT.get(scheme, 80)
    if query:
        path = path + '?' + query
    # 允许只给主机名（例如 wss://ai.jicreate.cn），自动补上桥接端点路径
    if path == '/' or path.startswith('/?'):
        path = BRIDGE_PATH + path[1:]
    return scheme, host, port, path


def log(message):
    sys.stdout.write('[%s] %s\n' % (time.strftime('%H:%M:%S'), message))
    sys.stdout.flush()


def limit_text(text, max_bytes=3000):
    """服务端单条文本消息有上限，codex 的事件行可能超长，按 UTF-8 字节截断。"""
    if not text:
        return ''
    data = text.encode('utf-8')
    if len(data) <= max_bytes:
        return text
    return data[:max_bytes].decode('utf-8', 'ignore') + '…'


# ====================== 任务执行 ======================

class CodexRunner:
    """在本机跑 codex exec --json，把事件流交给回调。"""

    def __init__(self, codex_bin, workspace, extra_env=None, sandbox='danger-full-access'):
        self.codex_bin = codex_bin
        self.sandbox = sandbox
        self.workspace = workspace
        self.extra_env = dict(extra_env or {})
        self.process = None
        self.lock = threading.Lock()
        self.last_lines = []

    def interrupt(self):
        with self.lock:
            process = self.process
        if process is None or process.poll() is not None:
            return False
        log('收到打断指令，正在结束本机 codex 进程…')
        try:
            if os.name == 'nt':
                subprocess.call(['taskkill', '/F', '/T', '/PID', str(process.pid)],
                                stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
            else:
                try:
                    os.killpg(os.getpgid(process.pid), 15)
                except Exception:
                    process.terminate()
        except Exception as error:
            log('结束进程失败: %s' % error)
        try:
            process.wait(timeout=10)
        except Exception:
            try:
                process.kill()
            except Exception:
                pass
        return True

    def run(self, task_id, prompt, on_event):
        """同步执行一个任务；返回 (status, message, exit_code)。"""
        if not os.path.isdir(self.workspace):
            os.makedirs(self.workspace)
        last_message_file = os.path.join(self.workspace, 'last-message-%s.txt' % task_id)
        command = [self.codex_bin, 'exec', '--json', '--skip-git-repo-check',
                   '-s', self.sandbox, '-C', self.workspace,
                   '-o', last_message_file, prompt]
        log('开始执行任务 %s（工作区 %s）' % (task_id, self.workspace))
        creationflags = 0
        popen_kwargs = {}
        if os.name == 'nt':
            creationflags = getattr(subprocess, 'CREATE_NEW_PROCESS_GROUP', 0)
        else:
            popen_kwargs['start_new_session'] = True
        env = None
        if self.extra_env:
            env = dict(os.environ)
            env.update(self.extra_env)
        self.last_lines = []
        try:
            process = subprocess.Popen(
                command,
                stdin=subprocess.DEVNULL,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                cwd=self.workspace,
                bufsize=1,
                universal_newlines=True,
                env=env,
                creationflags=creationflags,
                **popen_kwargs
            )
        except Exception as error:
            return 'ERROR', '启动本机 codex 失败：%s' % error, -1
        with self.lock:
            self.process = process

        for line in iter(process.stdout.readline, ''):
            text = (line or '').rstrip('\n')
            if text:
                self.remember(text)
                on_event(text)
        process.stdout.close()
        exit_code = process.wait()
        with self.lock:
            self.process = None

        message = ''
        try:
            if os.path.isfile(last_message_file):
                with open(last_message_file, 'r', encoding='utf-8', errors='ignore') as handle:
                    message = handle.read().strip()
                os.remove(last_message_file)
        except Exception:
            pass
        if exit_code != 0 and not message:
            # 把 codex 真正的报错回传，否则用户只看到"失败"没法排查
            message = self.failure_reason(exit_code)
        status = 'DONE' if exit_code == 0 or message and '失败' not in message[:2] else 'ERROR'
        if exit_code != 0:
            status = 'ERROR'
        log('任务 %s 结束：状态=%s 退出码=%s' % (task_id, status, exit_code))
        return status, message, exit_code

    def remember(self, line):
        self.last_lines.append(line)
        if len(self.last_lines) > 8:
            self.last_lines.pop(0)

    def failure_reason(self, exit_code):
        """从 codex 的事件流里挑一条最有用的错误说明。"""
        candidates = []
        for line in self.last_lines:
            text = line.strip()
            if not text:
                continue
            try:
                node = json.loads(text)
            except Exception:
                candidates.append(text)
                continue
            message = node.get('message')
            if not message and isinstance(node.get('error'), dict):
                message = node['error'].get('message')
            if not message and isinstance(node.get('item'), dict):
                message = node['item'].get('message')
            if message:
                candidates.append(str(message))
        for text in reversed(candidates):
            lowered = text.lower()
            if 'error' in lowered or 'missing' in lowered or 'failed' in lowered or '失败' in text:
                return text.replace('\n', ' ')[:300]
        if candidates:
            return candidates[-1].replace('\n', ' ')[:300]
        return '本机 codex 退出码 %s（无输出）' % exit_code


# ====================== 桥接主循环 ======================

class Bridge:
    def __init__(self, args):
        self.args = args
        self.runner = CodexRunner(args.codex, args.workspace, parse_env(args.env), args.sandbox)
        self.tasks = queue.Queue()
        self.busy = threading.Event()
        self._running = set()
        self._uploads = {}
        self._running_lock = threading.Lock()
        self.ws = None
        self.should_run = True
        self.current_task = None

    # ---------- 网络 ----------
    def build_url(self):
        separator = '&' if '?' in self.args.server else '?'
        if self.args.key:
            auth = 'key=' + quote(self.args.key)
        else:
            auth = 'token=' + quote(self.args.token)
        return '%s%s%s&device=%s' % (
            self.args.server.rstrip('/'),
            separator,
            auth,
            quote(self.args.name)
        )

    def connect_forever(self):
        backoff = 1
        while self.should_run:
            url = self.build_url()
            self.ws = WsClient(url, insecure=self.args.insecure,
                               on_text=self.on_message, on_close=self.on_closed)
            try:
                log('连接服务器 %s …' % sanitize(url))
                self.ws.connect()
                log('已连接（设备名 %s）' % self.args.name)
                backoff = 1
                self.send_hello()
                self.loop()
            except Exception as error:
                log('连接失败/断开：%s' % error)
            finally:
                if self.ws:
                    self.ws.close()
            if not self.should_run:
                break
            log('%d 秒后重连…' % backoff)
            time.sleep(backoff)
            backoff = min(30, backoff * 2)

    def loop(self):
        last_heartbeat = 0.0
        while self.should_run and self.ws and self.ws.connected:
            self.ws.pump(1.0)
            now = time.time()
            if now - last_heartbeat >= self.args.heartbeat:
                last_heartbeat = now
                self.ws.send_json({'type': 'heartbeat', 'ts': int(now * 1000)})
            # 服务器静默太久：TCP 可能已经半开，主动重连
            if now - self.ws.last_inbound > max(30.0, self.args.heartbeat * 3):
                log('服务器静默超过 %d 秒，判定连接已死，主动重连' % max(30, self.args.heartbeat * 3))
                self.ws.close()
                return
            # 从队列取任务：并行跑，来一个起一个（最多 max_concurrent 个）
            if self.running_count() < self.args.max_concurrent:
                try:
                    task = self.tasks.get_nowait()
                except queue.Empty:
                    task = None
                if task:
                    # 任务放到独立线程里跑，主循环继续收发与发心跳，
                    # 否则长任务期间心跳会断，服务器会以为本机掉线了
                    worker = threading.Thread(target=self.run_task, args=(task,),
                                              name='codex-task-%s' % task.get('taskId'))
                    worker.setDaemon(True)
                    worker.start()

    def running_count(self):
        with self._running_lock:
            return len(self._running)

    def on_closed(self, payload):
        log('服务器关闭了连接' + ('（%s）' % payload.decode('utf-8', 'ignore') if payload else ''))

    # ---------- 协议 ----------
    def send_hello(self):
        self.ws.send_json({
            'type': 'hello',
            'device': self.args.name,
            'platform': '%s %s' % (platform.system(), platform.release()),
            'codexVersion': codex_version(self.args.codex),
            'workspace': os.path.abspath(self.args.workspace),
            'python': platform.python_version(),
            'client': 'codex-voice-bridge/1.0'
        })

    def on_message(self, raw):
        try:
            message = json.loads(raw)
        except Exception:
            return
        kind = message.get('type')
        if kind == 'ping':
            # 心跳是双向的：服务器 ping，我们立刻 pong
            self.ws.send_json({'type': 'pong', 'ts': message.get('ts')})
        elif kind == 'ready' or kind == 'welcome':
            log('服务器已就绪（心跳 %sms / 超时 %sms）'
                % (message.get('pingMs', self.args.heartbeat * 1000), message.get('timeoutMs', '-')))
        elif kind == 'task':
            task = {
                'taskId': message.get('taskId') or uuid.uuid4().hex[:16],
                'instruction': message.get('instruction') or '',
                'workspace': message.get('workspace') or self.args.workspace
            }
            if self.running_count() >= self.args.max_concurrent:
                # 并行已满，明确回绝，避免静默排队
                self.ws.send_json({'type': 'result', 'taskId': task['taskId'], 'status': 'ERROR',
                                   'message': '本机同时最多跑 %d 个任务，这条没接' % self.args.max_concurrent,
                                   'exitCode': -1})
                return
            log('开始执行任务 %s（并行 %d/%d）'
                % (task['taskId'], self.running_count() + 1, self.args.max_concurrent))
            self.tasks.put(task)
        elif kind == 'file_chunk':
            self.save_file_chunk(message)
        elif kind == 'interrupt':
            self.runner.interrupt()
        elif kind == 'error':
            log('服务器报错：%s' % message.get('message'))

    # ---------- 收文件 ----------
    def save_file_chunk(self, message):
        """服务器把用户上传的文件分片推过来，这里拼起来写进工作区。"""
        upload_id = str(message.get('uploadId') or '')
        if not upload_id:
            return
        raw_name = str(message.get('name') or '')
        name = raw_name.replace('\\', '/')
        if '/' in name:
            name = name.rsplit('/', 1)[-1]
        for ch in '\\/:*?"<>|':
            name = name.replace(ch, '_')
        name = name.strip().lstrip('.')
        if not name:
            name = 'upload-%s.bin' % upload_id
        target = os.path.join(self.args.workspace, name)
        entry = self._uploads.get(upload_id)
        if entry is None:
            handle = open(target, 'wb')
            entry = {'name': name, 'path': target, 'handle': handle, 'bytes': 0,
                     'total': int(message.get('total') or 1), 'index': 0}
            self._uploads[upload_id] = entry
        try:
            data = base64.b64decode(message.get('data') or '')
        except Exception:
            data = b''
        if data:
            entry['handle'].write(data)
            entry['bytes'] += len(data)
        entry['index'] += 1
        if entry['index'] >= entry['total']:
            try:
                entry['handle'].flush()
                entry['handle'].close()
            except Exception:
                pass
            log('已收到文件 %s（%d 字节）→ %s' % (entry['name'], entry['bytes'], entry['path']))
            self.ws.send_json({'type': 'file_result', 'uploadId': upload_id, 'ok': True,
                               'path': entry['path'],
                               'message': '已保存到本机：%s' % entry['path']})
            self._uploads.pop(upload_id, None)

    # ---------- 任务 ----------
    def run_task(self, task):
        with self._running_lock:
            self._running.add(task['taskId'])
        workspace = task.get('workspace') or self.args.workspace
        if os.path.abspath(workspace) != os.path.abspath(self.runner.workspace):
            self.runner.workspace = workspace
        try:
            def on_event(line):
                self.ws.send_json({'type': 'event', 'taskId': task['taskId'],
                                   'line': limit_text(line)})

            status, message, exit_code = self.runner.run(task['taskId'], task['instruction'], on_event)
            self.ws.send_json({'type': 'result', 'taskId': task['taskId'], 'status': status,
                               'message': limit_text(message, 6000), 'exitCode': exit_code})
        except Exception as error:
            self.ws.send_json({'type': 'result', 'taskId': task['taskId'], 'status': 'ERROR',
                               'message': '本机执行异常：%s' % error, 'exitCode': -1})
        finally:
            with self._running_lock:
                self._running.discard(task['taskId'])


def quote(value):
    try:
        from urllib.parse import quote as _quote
    except ImportError:
        from urllib import quote as _quote
    return _quote(value or '')


def sanitize(url):
    if 'token=' not in url:
        return url
    head, _, tail = url.partition('token=')
    token, _, rest = tail.partition('&')
    return '%stoken=%s…%s&%s' % (head, token[:4], token[-2:] if len(token) > 6 else '', rest)


def codex_version(codex_bin):
    try:
        output = subprocess.check_output([codex_bin, '--version'], stderr=subprocess.STDOUT, timeout=20)
        return output.decode('utf-8', 'ignore').strip().splitlines()[0]
    except Exception:
        return ''


def parse_env(items):
    """--env KEY=VALUE（可重复）→ dict；用于给本机 codex 传它需要的环境变量。"""
    result = {}
    for item in (items or []):
        if '=' not in item:
            continue
        key, _, value = item.partition('=')
        result[key.strip()] = value
    return result


def find_codex(explicit):
    if explicit:
        return explicit
    for name in (['codex.cmd', 'codex.exe', 'codex'] if os.name == 'nt' else ['codex']):
        found = shutil.which(name)
        if found:
            return found
    return 'codex'


def main():
    parser = argparse.ArgumentParser(description='Codex 语音副驾驶本机桥接客户端')
    parser.add_argument('--server', default='ws://127.0.0.1:8188', help='服务器地址，如 wss://ai.jicreate.cn')
    parser.add_argument('--token', default='', help='登录 token（会随重新登录失效，临时用）')
    parser.add_argument('--key', default='', help='长期桥接密钥 bk_…（推荐，预配到本机用这个）')
    parser.add_argument('--max-concurrent', type=int, default=4, help='本机同时跑几个 codex（默认 4）')
    parser.add_argument('--sandbox', default='danger-full-access',
                        help='传给 codex 的沙箱模式：danger-full-access / workspace-write / read-only')
    parser.add_argument('--workspace', default='./codex-work', help='本机工作目录')
    parser.add_argument('--codex', default='', help='codex 可执行文件路径（默认自动查找）')
    parser.add_argument('--name', default=platform.node() or '我的电脑', help='设备名')
    parser.add_argument('--heartbeat', type=float, default=10.0, help='心跳间隔秒数')
    parser.add_argument('--env', action='append', default=[],
                        help='给本机 codex 追加环境变量，可重复：--env KEY=VALUE')
    parser.add_argument('--insecure', action='store_true', help='跳过 TLS 证书校验')
    args = parser.parse_args()
    args.codex = find_codex(args.codex)
    if not args.token and not args.key:
        log('缺少身份：请用 --key（长期密钥）或 --token')
        return
    if not os.path.isdir(args.workspace):
        os.makedirs(args.workspace)
    log('本机 Codex 桥接启动')
    log('  codex    : %s' % args.codex)
    log('  版本     : %s' % (codex_version(args.codex) or '未知'))
    log('  工作区   : %s' % os.path.abspath(args.workspace))
    log('  服务器   : %s' % sanitize(args.server))
    log('  设备名   : %s' % args.name)
    log('  心跳     : %.0f 秒' % args.heartbeat)
    log('  沙箱     : %s' % args.sandbox)
    log('  并行上限 : %d 个任务' % args.max_concurrent)
    bridge = Bridge(args)
    try:
        bridge.connect_forever()
    except KeyboardInterrupt:
        log('已退出')


if __name__ == '__main__':
    main()