import base64,json,sqlite3,struct,sys,urllib.request
from pathlib import Path
root=Path(__file__).resolve().parent.parent
cfg=json.loads((root/'model-bridge/client.json').read_text())
base=cfg['url'].rstrip('/')
def request(path,body=None):
    req=urllib.request.Request(base+path,data=json.dumps(body).encode() if body else None,headers={'Content-Type':'application/json','Authorization':'Bearer '+cfg['token']})
    with urllib.request.urlopen(req,timeout=90) as response:return json.load(response)
report=root/'integration-work/image-pool-smoke.json'
if '--submit' in sys.argv:
    if report.exists():raise RuntimeError('Smoke test already submitted; poll the existing job')
    pool=request('/v1/images/pool')
    assert pool['auto_balance'] and len(pool['configured_workers'])>=2,pool
    result=request('/v1/images/generations',{'model':'qwen_image_2_1','prompt':'A photorealistic ceramic teacup on a simple wooden table, soft window light, no text, no watermark.','size':'512x512','n':1,'response_format':'url'})
    data={'pool':pool,'task_id':result['task_id'],'assigned_worker':result.get('provider_id')}
    report.write_text(json.dumps(data,indent=2))
else:
    data=json.loads(report.read_text())
    result=request('/v1/images/task/'+data['task_id'])
    data.update(status=result['status'],worker=result.get('provider_id'))
    if result['status']=='completed':
        with urllib.request.urlopen(result['image_url'],timeout=60) as response:content=response.read()
        target=root/'integration-work/image-pool-smoke.png';target.write_bytes(content)
        assert content.startswith(b'\x89PNG\r\n\x1a\n')
        data['dimensions']=list(struct.unpack('>II',content[16:24])); data['file']=str(target)
    if result.get('error'):data['error']=result['error']
    report.write_text(json.dumps(data,indent=2))
print(json.dumps(data))
with sqlite3.connect((root/'huobao-data/huobao.sqlite3').as_uri()+'?mode=ro',uri=True) as db:
    rows=db.execute("SELECT id,task_id,status FROM sys_task WHERE type='image' AND model='qwen_image_2_1' AND status='completed'").fetchall()
providers={}
for tid,jid,status in rows:
    try:
        worker=base64.urlsafe_b64decode(jid+'='*(-len(jid)%4)).decode().split(':',1)[0]
        providers[worker]=providers.get(worker,0)+1
    except (ValueError,TypeError):pass
print('Previously completed Qwen images by worker:',providers)