"""Connect an employee's Codex CLI to OneOS without changing existing configuration."""
import argparse
import getpass
import json
import os
from pathlib import Path
import re
import shutil
import subprocess
import sys
import urllib.error
import urllib.request

BASE_URL = 'https://aigateway.lnoneos.com/v1'
ENV_KEY = 'ONEOS_GATEWAY_API_KEY'
MODEL = re.compile(r'[A-Za-z0-9][A-Za-z0-9._:/-]{0,127}\Z')
KEY = re.compile(r'sk-gateway-[A-Za-z0-9_-]{40,100}\Z')


class ConnectionProblem(Exception):
    pass


class NoRedirect(urllib.request.HTTPRedirectHandler):
    def redirect_request(self, request, response, code, message, headers, new_url):
        return None  # Never forward an employee credential to another endpoint.


def available_models(key, base_url=BASE_URL):
    if not KEY.fullmatch(key):
        raise ConnectionProblem('员工密钥格式不正确，请核对完整密钥。')
    # The deployment URL is fixed by the administrator. Never put the key in a URL.
    request = urllib.request.Request(base_url + '/models', headers={'Authorization':'Bearer ' + key})
    try:
        with urllib.request.build_opener(NoRedirect()).open(request, timeout=15) as response:
            content = response.read(1024 * 1024 + 1)
            if len(content) > 1024 * 1024:
                raise ConnectionProblem('模型列表过大，请联系管理员检查网关。')
            payload = json.loads(content)
    except urllib.error.HTTPError as error:
        error.close()
        messages = {401:'员工密钥无效或已撤销。', 403:'员工没有访问权限。',
                    429:'请求已达限制，请稍后重试。', 502:'上游暂时不可达。',
                    503:'上游模型尚未就绪。', 504:'上游响应超时。'}
        raise ConnectionProblem(messages.get(error.code, '网关请求失败。')) from None
    except (OSError, ValueError):
        raise ConnectionProblem('无法连接网关或响应格式异常，请检查网络。') from None
    if not isinstance(payload, dict) or not isinstance(payload.get('data'), list):
        raise ConnectionProblem('模型列表格式异常，请联系管理员。')
    models = sorted({item['id'] for item in payload['data'] if isinstance(item, dict)
                     and isinstance(item.get('id'), str) and MODEL.fullmatch(item['id'])})
    if not models:
        raise ConnectionProblem('员工密钥有效，但尚无可用模型。请管理员完成上游账号和网络接入。')
    return models


def configuration(model):
    if not MODEL.fullmatch(model):
        raise ConnectionProblem('模型名称不符合要求。')
    return (f'model = {json.dumps(model)}\nmodel_provider = "oneos_gateway"\nweb_search = "disabled"\n\n'
            '[model_providers.oneos_gateway]\nname = "OneOS Gateway"\n'
            f'base_url = "{BASE_URL}"\nwire_api = "responses"\nenv_key = "{ENV_KEY}"\n'
            'requires_openai_auth = false\n')


def codex_arguments(model):
    configuration(model)  # Validate before creating command-line arguments.
    settings = {'model':model, 'model_provider':'oneos_gateway', 'web_search':'disabled',
        'model_providers.oneos_gateway.name':'OneOS Gateway',
        'model_providers.oneos_gateway.base_url':BASE_URL,
        'model_providers.oneos_gateway.wire_api':'responses',
        'model_providers.oneos_gateway.env_key':ENV_KEY,
        'model_providers.oneos_gateway.requires_openai_auth':False}
    arguments = []
    for name, value in settings.items():
        arguments.extend(['-c', name + '=' + json.dumps(value)])
    return arguments


def launch(model, key, executable=None):
    binary = executable or shutil.which('codex')
    if not binary:
        raise ConnectionProblem('未找到 Codex CLI，请先安装公司批准的客户端。')
    environment = dict(os.environ)
    environment[ENV_KEY] = key
    # Only the child process receives the key; no config, shell history or global env change.
    return subprocess.run([binary,*codex_arguments(model)], env=environment, check=False).returncode


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--check', action='store_true', help='Verify access and list models without starting Codex')
    parser.add_argument('--model', help='Choose an ID returned by the gateway')
    parser.add_argument('--credential-file', type=Path, help='Private employee JSON exported by gatewayctl.py')
    parser.add_argument('--write-config', type=Path, help='Write a key-free fragment to a NEW file only')
    args = parser.parse_args()
    try:
        if args.credential_file:
            if args.credential_file.is_symlink():
                raise ConnectionProblem('员工凭证文件不能是符号链接。')
            if os.name != 'nt' and args.credential_file.stat().st_mode & 0o077:
                raise ConnectionProblem('员工凭证文件权限必须限制为本人可读。')
            credential = json.loads(args.credential_file.read_text(encoding='utf-8-sig'))
            if credential.get('base_url') != BASE_URL:
                raise ConnectionProblem('凭证文件的网关地址与公司配置不一致。')
            key = credential['api_key']
        else:
            if not sys.stdin.isatty():
                raise ConnectionProblem('请在交互终端运行以隐藏输入密钥，或使用受保护的员工凭证文件。')
            key = getpass.getpass('员工密钥（输入不显示）：').strip()
        models = available_models(key)
        if args.check:
            print(json.dumps({'connected':True,'base_url':BASE_URL,'models':models},ensure_ascii=False))
            return 0
        if args.model:
            model = args.model
        elif len(models) == 1:
            model = models[0]
        else:
            for index, name in enumerate(models,1):
                print(f'{index}. {name}')
            choice = input('选择模型编号：').strip()
            if not choice.isdecimal() or not 1 <= int(choice) <= len(models):
                raise ConnectionProblem('模型编号无效。')
            model = models[int(choice)-1]
        if model not in models:
            raise ConnectionProblem('所选模型不在网关当前返回的列表中。')
        if args.write_config:
            descriptor = os.open(args.write_config, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
            with os.fdopen(descriptor,'w',encoding='utf-8',newline='\n') as stream:
                stream.write(configuration(model))
            print('已写入不含密钥的配置片段；请合并到已有配置，不要覆盖现有内容。')
            return 0
        print(f'启动 OneOS 网关客户端，模型：{model}。退出后本次密钥不保留。')
        return launch(model,key)
    except ConnectionProblem as error:
        print(str(error),file=sys.stderr)
        return 2
    except (OSError, ValueError, KeyError, TypeError):
        print('无法读取凭证、写入新文件或启动客户端；未覆盖现有配置。',file=sys.stderr)
        return 2


if __name__ == '__main__':
    raise SystemExit(main())