"""Exercise GPU replacement preflight/cancellation with local images, no real GPU jobs."""
import base64, json, os, pathlib, re, sqlite3, subprocess, threading, time, urllib.error, urllib.request
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

root = pathlib.Path(__file__).resolve().parent.parent
data = root / 'integration-work' / ('replacement-test-' + time.strftime('%Y%m%d-%H%M%S'))
data.mkdir(); (data / 'static/images').mkdir(parents=True)
png = base64.b64decode('iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+aN1kAAAAASUVORK5CYII=')
(data / 'static/images/reference.png').write_bytes(png)
stage = root / 'integration-work/stage-4.0.11/app'
source = (root / 'huobao-source/desktop/build/backend.mjs').read_text('utf-8')
source = source.replace('const require = __hbCreateRequire(import.meta.url)', 'const require = __hbCreateRequire(' + json.dumps(str(stage / 'resources/app.asar/build/backend.mjs')) + ')', 1)
source = re.sub(r'^import (\w+) from "(better-sqlite3|sharp)";', r'const \1 = require("\2");', source, flags=re.M)
needle = 'target.hostname === "192.168.110.111" && ["18792", "18794"].includes(target.port)'
assert source.count(needle) == 1
source = source.replace(needle, 'target.hostname === "127.0.0.1" && target.port === "18935"')
source, count = re.subn(r'(url\d*)\.hostname === "192\.168\.110\.111" && \["18792", "18794"\]\.includes\(\1.port\)',
    r'\1.hostname === "127.0.0.1" && \1.port === "18935"', source)
assert count == 1
# Fake IDs are private to the loopback fixture.
source, count = re.subn(r'&& !!taskRouting\(\{ provider: config\d*\.provider, taskId \}\)', '&& !!taskId', source)
assert count == 1
(data / 'backend.mjs').write_text(source, 'utf-8')
events = []; reject = False; process = None; cancel_modes = {}; job_states = {}

class Mock(BaseHTTPRequestHandler):
    def log_message(self, *args): pass
    def reply(self, status, body):
        raw = json.dumps(body).encode(); self.send_response(status)
        self.send_header('Content-Type', 'application/json'); self.end_headers(); self.wfile.write(raw)
    def do_POST(self):
        body = json.loads(self.rfile.read(int(self.headers['Content-Length'])))
        events.append((self.path, body))
        if self.path.endswith('/validate'):
            images = [x['image_url']['url'] for x in body['content'] if x['type'] == 'image_url']
            if reject or not images or any(not x.startswith('data:image/') for x in images):
                self.reply(400, {'error': {'message': 'Local video requires inline image data'}})
            else: self.reply(200, {'valid': True})
        elif self.path.endswith('/cancel'):
            jid=self.path.split('/')[-2];mode=cancel_modes.get(jid,'cancelled')
            if mode=='wrong-id':self.reply(200,{'task':{'id':'different-job','status':'cancelled'}})
            elif mode=='transport':self.reply(503,{'error':'transient'})
            else:
                job_states[jid]=mode
                self.reply(200, {'task': {'id':jid, 'status':mode}})
        else:
            self.reply(200, {'task_id': 'new-' + str(len(events)), 'routing': {'gpu': body['gpu'], 'port': 8191, 'providerId': 'gpu3', 'legacy': False}})
    def do_GET(self):
        if self.path.startswith('/v1/images/task/'):
            jid=self.path.split('/')[-1];state=job_states.get(jid,'running')
            self.reply(200,{'task_id':jid,'status':'failed' if state=='cancelled' else 'processing',
                'cancelled':state=='cancelled','worker_state':'running','error':{'message':'cancelled'}});return
        self.reply(200, {'task': {'id': self.path.split('/')[-1], 'status': 'running', 'worker_state': 'running',
            'progress': {'phase':'sampling','step':4,'totalSteps':20,'observedAt':round(time.time()*1000),'samplingMinSeconds':120,'samplingMaxSeconds':180}}})

server = ThreadingHTTPServer(('127.0.0.1', 18935), Mock)
threading.Thread(target=server.serve_forever, daemon=True).start()
env = {**os.environ, 'ELECTRON_RUN_AS_NODE': '1', 'PORT': '18934', 'HUOBAO_DATA_DIR': str(data),
       'SQLITE_PATH': str(data / 'huobao.sqlite3'), 'FRONTEND_DIST': str(root / 'huobao-source/frontend/.output/public'),
       'WORKSPACE_PATH': str(stage / 'resources/workspace-template')}
log = (data / 'backend.log').open('w', encoding='utf-8')
def api(method, path, body=None, expected=200):
    request = urllib.request.Request('http://127.0.0.1:18934/api/v1' + path, method=method,
        data=None if body is None else json.dumps(body).encode(), headers={'Content-Type': 'application/json'})
    try:
        with urllib.request.urlopen(request, timeout=40) as response: status = response.status; result = json.load(response)
    except urllib.error.HTTPError as error: status = error.code; result = json.load(error)
    assert status == expected, (path, status, result)
    return result.get('data', result)
def sql(statement, args=()):
    with sqlite3.connect(data / 'huobao.sqlite3') as db:
        db.row_factory = sqlite3.Row
        return [dict(row) for row in db.execute(statement, args)]
def replace(shot, old, **extra):
    return api('POST', '/tasks', {'type': 'video', 'storyboard_id': shot, 'config_id': 1,
        'model': 'MiniMax-H3', 'gpu': 'gpu3', 'replace_task_id': old, 'prompt': 'reference fixture',
        'reference_image_urls': ['static/images/reference.png'], **extra.pop('body', {})}, **extra)
try:
    process = subprocess.Popen([str(stage / 'HuobaoDrama.exe'), str(data / 'backend.mjs')], cwd=data,
        env=env, stdout=log, stderr=log, creationflags=subprocess.CREATE_NO_WINDOW)
    for _ in range(150):
        assert process.poll() is None, 'Backend exited'
        try: api('GET', '/health'); break
        except OSError: time.sleep(.1)
    ts = '2026-11-12T00:00:00Z'
    sql('INSERT INTO dramas(id,title,created_at,updated_at) VALUES(1,?,?,?)', ('fixture',ts,ts))
    sql('INSERT INTO episodes(id,drama_id,episode_number,title,created_at,updated_at) VALUES(1,1,1,?,?,?)', ('fixture',ts,ts))
    sql('INSERT INTO ai_service_configs(id,service_type,provider,name,base_url,api_key,model,is_active,created_at,updated_at) VALUES(1,?,?,?,?,?,?,1,?,?)',
        ('video','minimax','fixture','http://127.0.0.1:18935','test-key',json.dumps(['MiniMax-H3']),ts,ts))
    for i in range(1,5):
        sql('INSERT INTO storyboards(id,episode_id,storyboard_number,created_at,updated_at) VALUES(?,1,?,?,?)',(i,i,ts,ts))
        sql('INSERT INTO sys_task(id,type,storyboard_id,provider,model,status,stage,task_id,config_id,config_base_url,created_at,updated_at) VALUES(?,?,?,?,?,?,?,?,?,?,?,?)',
            (i,'video',i,'minimax','MiniMax-H3','processing','running','old-'+str(i),1,'http://127.0.0.1:18935',ts,ts))
    # Real local image bytes must reach preflight before cancelling exactly one old job.
    result = replace(1,1,expected=201)
    for _ in range(100):
        if len(events) >= 3: break
        time.sleep(.1)
    assert [p for p,b in events[:3]] == ['/v2/video_generation/validate','/v2/video_generation/old-1/cancel','/v2/video_generation']
    assert events[0][1]['content'] == events[2][1]['content']
    assert base64.b64decode(events[0][1]['content'][1]['image_url']['url'].split(',')[1]) == png
    assert events[0][1]['gpu'] == events[2][1]['gpu'] == 'gpu3'
    assert sql('SELECT error_code FROM sys_task WHERE id=1')[0]['error_code'] == 'user_cancelled'
    assert sql('SELECT status FROM sys_task WHERE id=2')[0]['status'] == 'processing'
    # Missing reference fails before any upstream call, leaving the old job intact.
    before = len(events)
    response = replace(2,2,expected=400,body={'reference_image_urls':['static/images/missing.png']})
    assert '参考图读取失败' in response['message'] and len(events) == before
    assert sql('SELECT status FROM sys_task WHERE id=2')[0]['status'] == 'processing'
    # Unready destination also leaves old task intact; no silent duplicate submission.
    reject = True; before = len(events)
    replace(3,3,expected=400)
    assert len(events) == before+1 and events[-1][0].endswith('/validate')
    assert sql('SELECT status FROM sys_task WHERE id=3')[0]['status'] == 'processing'
    # First-frame replacements use the same conversion and role.
    reject = False; before = len(events)
    replace(4,4,expected=201,body={'reference_image_urls':[], 'first_frame_url':'/static/images/reference.png'})
    for _ in range(100):
        if len(events) >= before+3: break
        time.sleep(.1)
    image = events[before][1]['content'][1]
    assert image['role'] == 'first_frame' and image['image_url']['url'].startswith('data:image/')
    assert events[before][1]['content'] == events[before+2][1]['content']
    for _ in range(140):
        row=sql('SELECT params FROM sys_task WHERE id=?',(result['id'],))[0]
        progress=json.loads(row['params']).get('progress')
        if progress: break
        time.sleep(.1)
    assert progress['step']==4 and progress['totalSteps']==20
    assert progress['samplingMaxSeconds']==180
    # Standalone cancellation never creates another generation or touches neighbors.
    before=len(events);assert api('POST','/tasks/2/cancel',{})['outcome']=='cancelled'
    assert events[-1][0]=='/v2/video_generation/old-2/cancel' and len(events)==before+1
    assert sql('SELECT status FROM sys_task WHERE id=3')[0]['status']=='processing'
    before=len(events);assert api('POST','/tasks/2/cancel',{})['outcome']=='cancelled';assert len(events)==before
    sql('INSERT INTO ai_service_configs(id,service_type,provider,name,base_url,api_key,model,is_active,created_at,updated_at) VALUES(2,?,?,?,?,?,?,1,?,?)',
        ('image','openai','fixture-image','http://127.0.0.1:18935/v1','test-key',json.dumps(['qwen_image_2_1']),ts,ts))
    for i in [20,21,22,23]:
        sql('INSERT INTO sys_task(id,type,provider,model,status,stage,task_id,config_id,config_base_url,created_at,updated_at) VALUES(?,?,?,?,?,?,?,?,?,?,?)',
            (i,'image','openai','qwen_image_2_1','processing','submitting' if i==23 else 'queued',None if i==23 else 'image-'+str(i),2,'http://127.0.0.1:18935/v1',ts,ts))
    sql('INSERT INTO sys_task(id,type,status,stage,local_path,created_at,updated_at) VALUES(19,?,?,?,?,?,?)',('image','completed','completed','static/images/reference.png',ts,ts))
    assert api('GET','/tasks/20')['canCancel'] is True
    before=len(events);assert api('POST','/tasks/20/cancel',{})['outcome']=='cancelled'
    assert events[-1][0]=='/v1/images/task/image-20/cancel' and len(events)==before+1
    assert sql('SELECT local_path FROM sys_task WHERE id=19')[0]['local_path']=='static/images/reference.png'
    assert (data/'static/images/reference.png').read_bytes()==png
    # Wrong receipt / transport failures are unconfirmed, never falsely cancelled.
    cancel_modes['image-21']='wrong-id'
    assert api('POST','/tasks/21/cancel',{})['outcome']=='pending'
    assert sql('SELECT status,stage FROM sys_task WHERE id=21')[0]=={'status':'processing','stage':'cancelling'}
    cancel_modes['image-21']='transport'
    assert api('POST','/tasks/21/cancel',{})['outcome']=='pending'
    # A completion racing with cancellation is preserved for normal download.
    cancel_modes['image-22']='succeeded'
    assert api('POST','/tasks/22/cancel',{})['outcome']=='completed'
    assert sql('SELECT error_code,status FROM sys_task WHERE id=22')[0]=={'error_code':None,'status':'processing'}
    before=len(events);api('POST','/tasks/23/cancel',{},expected=400);api('POST','/tasks/19/cancel',{},expected=400)
    assert len(events)==before and api('GET','/tasks/23')['canCancel'] is False
    sql('INSERT INTO sys_task(id,type,status,stage,created_at,updated_at) VALUES(24,?,?,?,?,?)',('video','processing','waiting_predecessor',ts,ts))
    assert api('POST','/tasks/24/cancel',{})['outcome']=='cancelled' and len(events)==before
    # The unconfirmed request survives restart; polling confirms the same ID,
    # without another cancel or generation request.
    job_states['image-21']='cancelled'
    process.terminate();process.wait(timeout=10)
    process=subprocess.Popen([str(stage/'HuobaoDrama.exe'),str(data/'backend.mjs')],cwd=data,env=env,stdout=log,stderr=log,creationflags=subprocess.CREATE_NO_WINDOW)
    for _ in range(180):
        row=sql('SELECT status,error_code FROM sys_task WHERE id=21')[0]
        if row['error_code']=='user_cancelled':break
        time.sleep(.1)
    assert row=={'status':'failed','error_code':'user_cancelled'} and len(events)==before
    print(json.dumps({'passed':True,'inline_reference_preflight':True,'first_frame':True,'cancel_only_selected_shot':True,
        'preserve_old_on_missing_image':True,'preserve_old_on_validation_failure':True,'image_cancel':True,'cancel_idempotent':True,
        'completion_race_preserved':True,'bad_receipt_not_cancelled':True,'cancel_survives_restart':True,'fixture':str(data)}))
finally:
    if process and process.poll() is None: process.terminate(); process.wait(timeout=10)
    server.shutdown(); log.close()