#!/usr/bin/env python3
"""IQ Lab local companion. Inspect this file before running. Python 3.9+.

Uses a fresh temporary Codex home for official device-code login. Credentials
are never uploaded. Only this test's final answer, elapsed time, and pairing
code are submitted. Local submissions are explicitly labelled unverified.
"""
import argparse
import json
import os
import pathlib
import shutil
import signal
import subprocess
import sys
import tempfile
import time
import urllib.parse
import urllib.request


def child_run(argv, *, cwd, env, timeout, capture=False, input=None):
    """Stop the full CLI process group on timeout, including launcher children."""
    proc=subprocess.Popen(argv,cwd=cwd,env=env,text=True,
                          stdin=subprocess.PIPE if input is not None else None,
                          stdout=subprocess.PIPE if capture else None,
                          stderr=subprocess.PIPE if capture else None,
                          start_new_session=os.name=='posix')
    try:
        stdout,stderr=proc.communicate(input=input,timeout=timeout)
        return subprocess.CompletedProcess(argv,proc.returncode,stdout,stderr)
    finally:
        if proc.poll() is None:
            if os.name=='posix':
                try:os.killpg(proc.pid,signal.SIGKILL)
                except ProcessLookupError:pass
            else:proc.kill()
            proc.wait()


def main():
    p=argparse.ArgumentParser(description='IQ Lab · 本地 Codex 脑力体检')
    p.add_argument('--site',default='https://iq.sjtuagent.cn/')
    p.add_argument('--session',required=True,help='网页生成的10分钟配对码')
    args=p.parse_args()
    site=args.site.rstrip('/')+'/'
    u=urllib.parse.urlsplit(site)
    if u.scheme!='https' or u.hostname not in ('iq.sjtuagent.cn','sjtuagent.cn') or u.username or u.password:
        sys.exit('仅允许连接 https://iq.sjtuagent.cn/ 或 https://sjtuagent.cn/iq/。')
    codex=shutil.which('codex')
    if not codex:
        sys.exit('未找到 Codex CLI。请先从官方文档安装：https://learn.chatgpt.com/docs/codex-cli')
    def request(path,body=None):
        req=urllib.request.Request(site+path,data=json.dumps(body).encode() if body is not None else None,headers={'Content-Type':'application/json','User-Agent':'IQ-Lab-Local/1.0'})
        # The allowlisted origin must not redirect a credential-bearing pairing request.
        class NoRedirect(urllib.request.HTTPRedirectHandler):
            def redirect_request(self,*a,**kw):return None
        opener=urllib.request.build_opener(NoRedirect())
        with opener.open(req,timeout=20) as r:return json.load(r)
    try:
        challenge=request('api/local/challenge?code='+urllib.parse.quote(args.session))
    except Exception:
        sys.exit('配对码无效或网络不可用。请在网页重新创建配对。')
    print('\nIQ Lab · 本地脑力体检\n')
    print('接下来会在本机启动官方设备码登录。请只在官方授权页面输入终端显示的设备码。')
    print('本次上传：8道题的最终答案、耗时与配对码。不会上传登录凭据。\n')
    with tempfile.TemporaryDirectory(prefix='iq-lab-') as temp:
        home=pathlib.Path(temp)/'codex';home.mkdir(mode=0o700)
        output=pathlib.Path(temp)/'answer.txt'
        env=os.environ.copy()
        env['CODEX_HOME']=str(home)
        # Do not inherit an API key or injected account token as an auth override.
        for key in ('OPENAI_API_KEY','CODEX_API_KEY','CODEX_ACCESS_TOKEN','OPENAI_BASE_URL'):
            env.pop(key,None)
        env['NO_COLOR']='1'
        config=home/'config.toml'
        config.write_text('cli_auth_credentials_store = "file"\nweb_search = "disabled"\n[features]\nshell_tool = false\n',encoding='utf-8')
        try:
            login=child_run([codex,'login','--device-auth'],cwd=temp,env=env,timeout=480)
            if login.returncode:sys.exit('登录未完成。请确认账号允许设备码登录，或改用 API 模式。')
            print('\n登录成功，开始8道小题（本次不使用工具）。\n')
            started=time.monotonic()
            result=child_run([codex,'exec','--ephemeral','--skip-git-repo-check','--sandbox','read-only','--color','never','--model',challenge['model'],'--output-last-message',str(output),'-'],input=challenge['prompt'],cwd=temp,env=env,capture=True,timeout=65)
            if result.returncode or not output.exists():
                sys.exit('模型未完成作答。请检查模型名称、账号额度和 CLI 版本。本次未上传结果。')
            answer=output.read_text(encoding='utf-8')
            if len(answer)>16000:sys.exit('模型返回内容过长，本次未上传。')
            request('api/local/submit',dict(code=args.session,answer=answer,elapsed=round(time.monotonic()-started,1)))
            print('体检完成！回到网页查看你的脑力报告。临时登录凭据随程序退出删除。')
        except subprocess.TimeoutExpired:
            sys.exit('授权或测试超时，本次未上传。临时登录目录已清理；调用可能仍产生账号用量。')
        except KeyboardInterrupt:
            sys.exit('\n已退出，临时登录目录已清理。')
        except Exception:
            sys.exit('未能上传结果，请重新创建配对。临时登录目录已清理。')


if __name__=='__main__':main()