import subprocess, json, os, io, time, re
from base64 import b64decode, b64encode
from ckcc.client import ColdcardDevice
from ckcc.protocol import CCProtocolPacker
from ckcc.cli import real_file_upload

V='27.0'; DATA=os.path.expanduser('~/cksim/regtest-data')
CLI=os.path.expanduser('~/bitcoin-%s/bin/bitcoin-cli'%V)
def bc(*a, wallet=None):
    base=[CLI,'-datadir=%s'%DATA,'-rpcport=18999','-rpcuser=ck','-rpcpassword=ckms']
    if wallet: base.append('-rpcwallet=%s'%wallet)
    r=subprocess.run(base+list(a),capture_output=True,text=True)
    if r.returncode: raise SystemExit('%s FAIL %s'%(a,r.stderr.strip()))
    return r.stdout.strip()
SK={1:'/tmp/cksim-1.sock',2:'/tmp/cksim-2.sock',3:'/tmp/cksim-3.sock'}
def dev(n): return ColdcardDevice(sn=SK[n])
def E(d,c):
    r=d.send_recv(b'EXEC'+c.encode(),encrypt=False); return r.decode() if isinstance(r,(bytes,bytearray)) else r
def K(d,k): d.send_recv(CCProtocolPacker.sim_keypress(k.encode('ascii')))
def hstat(d):
    try:
        r=d.send_recv(CCProtocolPacker.hsm_status()); r=r.decode() if isinstance(r,(bytes,bytearray)) else r
        return json.loads(r)
    except Exception as e: return {'err':str(e)[:80]}
def story(d): return E(d,"RV.write(repr(sim_display.story))")

policy={"must_log":False,"period":60,"rules":[{"max_amount":300000000,"per_period":300000000,"wallet":"ckms23"}]}
pf=os.path.expanduser('~/cksim/policy.json'); open(pf,'w').write(json.dumps(policy))

print('== hsm-start on sims 1,2 (2-step approval: OK then random digit) ==')
for n in (1,2):
    d=dev(n)
    flen,sha=real_file_upload(open(pf,'rb'),d)
    d.send_recv(CCProtocolPacker.hsm_start(flen,sha)); time.sleep(1.2)
    s1=story(d)
    K(d,'y'); time.sleep(1.3)                 # screen 1: Press OK
    s2=story(d)
    m=re.search(r'Press \((\d)\) to save policy', s2)
    cc=m.group(1) if m else None
    if cc: K(d,cc); time.sleep(1.5)           # screen 2: the random digit
    st=hstat(d)
    print(' sim%d  confirm_digit=%s  HSM_active=%s' % (n, cc, st.get('active')))

print('== fund + build fresh PSBT ==')
dest=bc('getnewaddress',wallet='miner')
funded=json.loads(bc('walletcreatefundedpsbt','[]',json.dumps([{dest:0.5}]),'0',json.dumps({'fee_rate':2,'change_type':'bech32'}),wallet='ckms23-watch'))
psbt=b64decode(funded['psbt']); sigs=[]
print('== sign with NO keypress (HSM auto-cosign) ==')
for n in (1,2):
    d=dev(n)
    flen,sha=real_file_upload(io.BytesIO(psbt),d)
    d.send_recv(CCProtocolPacker.sign_transaction(flen,sha,flags=0x0),timeout=None)
    done=None
    for i in range(40):
        time.sleep(0.5)
        r=d.send_recv(CCProtocolPacker.get_signed_txn(),timeout=None)
        if r is not None: done=r; break
    if not done: print(' sim%d AUTO-SIGN TIMEOUT (hsm=%s)'%(n,hstat(d).get('active'))); continue
    rl,rsha=done; res=d.download_file(rl,rsha,file_number=1)
    sigs.append(b64encode(res).decode())
    print(' sim%d AUTO-SIGNED, NO keypress, bytes=%d' % (n, rl))
print('== combine + broadcast ==')
comb=bc('combinepsbt',json.dumps(sigs))
fin=json.loads(bc('finalizepsbt',comb)); print(' PSBT_COMPLETE:',fin['complete'])
txid=bc('sendrawtransaction',fin['hex']); print(' TXID:',txid)
bc('generatetoaddress','1',bc('getnewaddress',wallet='miner'),wallet='miner')
print(' BALANCE_AFTER:',bc('getbalance',wallet='ckms23-watch'))