import concurrent.futures,io,json,sqlite3,threading,time,uuid,hashlib
from email import policy
from email.parser import BytesParser
from http.server import BaseHTTPRequestHandler,ThreadingHTTPServer
from pathlib import Path
import requests
from PIL import Image

data=Path('D:/web/dj/integration-work/hd-test-references');base='http://127.0.0.1:18886/api/v1'
calls=[]
class Mock(BaseHTTPRequestHandler):
 def log_message(self,*args):pass
 def do_POST(self):
  raw=self.rfile.read(int(self.headers['Content-Length']));content_type=self.headers.get('Content-Type','')
  fields={};images=[]
  if 'multipart' in content_type:
   msg=BytesParser(policy=policy.default).parsebytes(('Content-Type: '+content_type+'\r\n\r\n').encode()+raw)
   for part in msg.iter_parts():
    name=part.get_param('name',header='content-disposition');value=part.get_payload(decode=True)
    if name=='image[]':
     with Image.open(io.BytesIO(value)) as im:images.append(list(im.convert('RGB').getpixel((50,50))))
    else:fields[name]=value.decode()
  else:fields=json.loads(raw)
  calls.append({'path':self.path,'fields':fields,'images':images})
  response=json.dumps({'data':[{'url':'http://127.0.0.1:18887/result.png'}]}).encode()
  self.send_response(200);self.send_header('Content-Type','application/json');self.end_headers();self.wfile.write(response)
 def do_GET(self):
  raw=(data/'green.png').read_bytes();self.send_response(200);self.send_header('Content-Type','image/png');self.end_headers();self.wfile.write(raw)
server=ThreadingHTTPServer(('127.0.0.1',18887),Mock)
threading.Thread(target=server.serve_forever,daemon=True).start()
def sql(query,args=()):
 with sqlite3.connect(data/'huobao.sqlite3') as db:
  db.row_factory=sqlite3.Row;return [dict(x) for x in db.execute(query,args)]
def upload(kind,id,color='red'):
 with (data/(color+'.png')).open('rb') as f:r=requests.post(f'{base}/{kind}/{id}/reference-images',files={'file':(color+'.png',f,'image/png')},timeout=10)
 return r
def refs(kind,id):return json.loads(sql(f'SELECT reference_images FROM {kind} WHERE id=?',(id,))[0]['reference_images'] or '[]')
def remove(kind,id,path):
 r=requests.delete(f'{base}/{kind}/{id}/reference-images',json={'path':path});assert r.status_code==200,r.text
def generate(kind,id,cfg=14):
 r=requests.post(f'{base}/{kind}/{id}/generate-image',json={'episode_id':7,'model':'qwen_image_2_1','config_id':cfg},timeout=20)
 assert r.status_code==200,r.text
 tid=r.json()['data']['image_generation_id'];return wait(tid)
def wait(tid):
 for _ in range(100):
  row=sql('SELECT * FROM sys_task WHERE id=?',(tid,))[0]
  if row['status']!='processing':return row
  time.sleep(.1)
 raise AssertionError(('task timeout',tid))
try:
 before={table:sql(f'SELECT image_url,local_path FROM {table} WHERE id=?',(id,))[0] for table,id in [('characters',26),('scenes',22),('props',1)]}
 for table,id in [('characters',26),('scenes',22),('props',1)]:
  for color in ['red','blue']:
   r=upload(table,id,color);assert r.status_code==200,r.text
  assert len(refs(table,id))==2
  assert sql(f'SELECT image_url,local_path FROM {table} WHERE id=?',(id,))[0]==before[table]
  row=generate(table,id);assert row['status']=='completed',row['error_msg']
  assert len(json.loads(row['params'])['referenceImages'])==2
  call=calls[-1];assert call['path']=='/v1/images/edits' and call['fields']['gpu']=='gpu1',call
  assert len(call['images'])==2 and call['images'][0][0]>180 and call['images'][1][2]>180,call
 print('PASS: character/scene/prop persist references, preserve current output on upload, send ordered multipart images with GPU1')
 # Independent assets and per-character batch snapshots.
 assert upload('characters',27,'green').status_code==200
 r=requests.post(base+'/characters/batch-generate-images',json={'character_ids':[26,27],'episode_id':7,'model':'qwen_image_2_1','config_id':16},timeout=20)
 assert r.status_code==200,r.text
 for tid in r.json()['data']['ids']:assert wait(tid)['status']=='completed'
 batch=calls[-2:];assert sorted(len(c['images']) for c in batch)==[1,2] and all(c['fields']['gpu']=='gpu3' for c in batch)
 print('PASS: batch retains different per-character references and selected GPU3')
 # Removing references does not erase files already snapshotted into past tasks.
 saved=refs('characters',26);oldpaths=sql('SELECT image_url,local_path FROM characters WHERE id=26')[0]
 for path in saved:remove('characters',26,path);assert (data/path).is_file()
 assert refs('characters',26)==[] and sql('SELECT image_url,local_path FROM characters WHERE id=26')[0]==oldpaths
 assert generate('characters',26)['status']=='completed' and calls[-1]['path']=='/v1/images/generations'
 # Invalid uploads, unknown/deleted assets and concurrent cap enforcement.
 r=requests.post(base+'/characters/26/reference-images',files={'file':('fake.png',b'not an image','image/png')});assert r.status_code==400
 assert upload('characters',999999).status_code==404
 for _ in range(5):assert upload('characters',26).status_code==200
 with concurrent.futures.ThreadPoolExecutor() as pool:statuses=list(pool.map(lambda _:upload('characters',26).status_code,range(2)))
 assert sorted(statuses)==[200,400] and len(refs('characters',26))==6,statuses
 # One unavailable reference must fail visibly, never fall back to text-only.
 broken=refs('characters',26)[0];(data/broken).rename(data/'temporarily-missing.png');count=len(calls)
 row=generate('characters',26);assert row['status']=='failed' and '参考图读取失败' in row['error_msg']
 assert len(calls)==count
 (data/'temporarily-missing.png').rename(data/broken)
 for path in refs('characters',26)[2:]:remove('characters',26,path)
 print('PASS: removal preserves output/history files, empty refs use text generation, invalid files rejected, concurrent cap enforced, missing refs never silently ignored')
 (data/'integration-result.json').write_text(json.dumps({'passed':True,'calls':len(calls),'asset_types':3,'maximum':6},indent=2))
finally:server.shutdown()