"""Install and collect a selected Android decision probe with diagnostic evidence."""
from __future__ import annotations
import argparse, json, re, shutil, subprocess, time
from pathlib import Path
from common import ROOT, CPUS, ProbeError, safe_relative, sha256, stamp, write_json

DECISION_PROFILES={'oq_host_view','oq_host_callback'}

class Adb:
    def __init__(self,serial:str|None): self.base=['adb']+(['-s',serial] if serial else [])
    def run(self,*args:str,check:bool=True,timeout:int=120)->subprocess.CompletedProcess:
        p=subprocess.run(self.base+list(args),capture_output=True,timeout=timeout)
        if check and p.returncode: raise ProbeError('ADB command failed: '+' '.join(args)+'\n'+p.stderr.decode(errors='replace'))
        return p
    def text(self,*args:str,check:bool=True)->str: return self.run(*args,check=check).stdout.decode(errors='replace').strip()

def own_report_paths(text:str)->list[str]:
    paths=[]
    for line in text.splitlines():
        name=line.strip()[2:] if line.strip().startswith('./') else line.strip()
        try: safe_relative(name)
        except ProbeError: continue
        if name.endswith('kam-architecture-probe/runtime.json'): paths.append(name)
    return paths

def validate_runtime(data:dict,package:dict,require_frames:bool=True)->list[str]:
    errors=[]
    for key,value in (('schema','kam-android-runtime/1'),('build_id',package['build_id']),('profile',package['profile']),('cpu_expected',package['cpu'])):
        if data.get(key)!=value: errors.append('Wrong/stale runtime '+key)
    if data.get('automatic_smoke_status')!='PASS': errors.append('Automatic smoke did not pass')
    tests=data.get('tests')
    if not isinstance(tests,list) or not tests: return errors+['No runtime tests']
    by_id={x.get('id'):x for x in tests if isinstance(x,dict)}; profile=package['profile']; required=['android.startup']; observed=[]
    if profile not in DECISION_PROFILES:
        required += ['abi.pointer','abi.scalar_layout','abi.endianness','files.native_utf8','files.seek','files.case_sensitive','fpc.zstream','fpc.thread']
        observed += ['decision.persistence.rename_replace','decision.persistence.append_reopen']
    if profile=='oq_host_view': required += ['decision.host.view']
    elif profile=='oq_host_callback': required += ['decision.host.callback']
    if profile=='platform': required += ['decision.text.cwstring_unit']; observed += ['decision.text.cp1251_decode','decision.text.cp1251_encode']
    if require_frames: required += ['render.frame_callbacks','render.context']
    for name in required:
        if by_id.get(name,{}).get('status')!='PASS': errors.append('Required runtime check did not pass: '+name)
    for name in observed:
        if name not in by_id: errors.append('Decision observation missing: '+name)
    if any(t.get('status')=='FAIL' for t in tests if isinstance(t,dict)): errors.append('One or more runtime checks failed')
    if require_frames and data.get('session',{}).get('frames',0)<30: errors.append('Fewer than 30 frames')
    return errors

def read_runtime(adb:Adb,pkg:str)->tuple[dict|None,str|None]:
    paths=own_report_paths(adb.text('shell','run-as',pkg,'find','.','-name','runtime.json',check=False))
    if len(paths)!=1: return None,None
    raw=adb.run('exec-out','run-as',pkg,'cat',paths[0],check=False).stdout
    try: return json.loads(raw),paths[0]
    except (ValueError,UnicodeDecodeError): return None,None

def wait_runtime(adb:Adb,pkg:str,predicate,timeout:int)->tuple[dict|None,str|None]:
    end=time.monotonic()+timeout
    while time.monotonic()<end:
        data,path=read_runtime(adb,pkg)
        if data is not None and predicate(data): return data,path
        time.sleep(1)
    return read_runtime(adb,pkg)

def save_command(out:Path,name:str,p:subprocess.CompletedProcess)->None:
    out.joinpath(name).write_text('exit='+str(p.returncode)+'\nstdout:\n'+p.stdout.decode(errors='replace')+'\nstderr:\n'+p.stderr.decode(errors='replace'),encoding='utf-8')

def main()->int:
    ap=argparse.ArgumentParser(description=__doc__)
    ap.add_argument('--profile',default='platform');ap.add_argument('--serial');ap.add_argument('--index',type=Path,default=ROOT/'dist/LATEST.json')
    ap.add_argument('--seconds',type=int,default=30);ap.add_argument('--exercise-lifecycle',action='store_true');ap.add_argument('--stage-file',type=Path)
    a=ap.parse_args()
    if a.seconds<5: ap.error('--seconds must be at least 5')
    out=ROOT/'reports/device'/stamp(); n=1
    while out.exists(): out=ROOT/'reports/device'/(stamp()+'-'+str(n));n+=1
    out.mkdir(parents=True)
    record={'schema':'kam-android-device/2','status':'BLOCKED','errors':[],'preflight':{'status':'NOT_RUN'},'manual_visual_touch_audio':'NOT_VERIFIED','world_simulation':'NOT_RUN','command_environment':{'host_adb':shutil.which('adb'),'timeout_seconds':a.seconds}}
    adb=None
    try:
        index=json.loads(a.index.read_text()); matches=[p for p in index.get('packages',[]) if p.get('profile')==a.profile]
        if len(matches)!=1: raise ProbeError('No unique successfully-built APK for '+a.profile)
        package=matches[0]; apk=ROOT/package['file']; pkg=package['package'];record.update(build_id=package['build_id'],profile=a.profile,package=package)
        if not apk.is_file() or sha256(apk)!=package['sha256']: raise ProbeError('APK missing/changed')
        if not re.fullmatch(r'org\.kammobile\.archprobe\.[a-z0-9]+',pkg): raise ProbeError('Refusing unrelated Android package')
        if not shutil.which('adb'): record['preflight']={'status':'BLOCKED','reason':'adb is not installed on host'};raise ProbeError(record['preflight']['reason'])
        listed=Adb(None).text('devices').splitlines()[1:]; devices=[line.split()[:2] for line in listed if line.strip()]
        record['preflight']={'adb':'PASS','devices':devices}
        if len(devices)!=1 or len(devices[0])!=2 or devices[0][1]!='device': record['preflight'].update(status='BLOCKED',reason='Require exactly one authorized device; found '+repr(devices));raise ProbeError(record['preflight']['reason'])
        if a.serial and a.serial!=devices[0][0]: record['preflight'].update(status='BLOCKED',reason='Requested serial is not the sole authorized device');raise ProbeError(record['preflight']['reason'])
        adb=Adb(devices[0][0]); record['device']={'serial':devices[0][0],'abi':adb.text('shell','getprop','ro.product.cpu.abilist'),'android':adb.text('shell','getprop','ro.build.version.release'),'sdk':adb.text('shell','getprop','ro.build.version.sdk'),'model':adb.text('shell','getprop','ro.product.model')}
        if CPUS[package['cpu']][0] not in record['device']['abi'].split(','): record['preflight'].update(status='BLOCKED',reason='Device ABI does not support arm64-v8a');raise ProbeError(record['preflight']['reason'])
        record['preflight']['status']='PASS'
        install=adb.run('install','-r',str(apk),check=False);save_command(out,'install.log',install);record['install']={'package_id':pkg,'initial_exit':install.returncode,'result':install.stdout.decode(errors='replace').strip()}
        if install.returncode and b'INSTALL_FAILED_UPDATE_INCOMPATIBLE' in install.stdout+install.stderr:
            uninstall=adb.run('uninstall',pkg,check=False);save_command(out,'uninstall.log',uninstall);record['install']['uninstall_for_ephemeral_key']={'exit':uninstall.returncode,'result':uninstall.stdout.decode(errors='replace').strip()}
            install=adb.run('install',str(apk),check=False);save_command(out,'install-retry.log',install);record['install'].update(retry_exit=install.returncode,retry_result=install.stdout.decode(errors='replace').strip())
        if install.returncode: raise ProbeError('APK install failed; see install.log')
        run_as=adb.run('shell','run-as',pkg,'id',check=False);save_command(out,'run-as.log',run_as);record['run_as']={'status':'PASS' if run_as.returncode==0 else 'BLOCKED','output':run_as.stdout.decode(errors='replace').strip(),'error':run_as.stderr.decode(errors='replace').strip()}
        if run_as.returncode: record['infrastructure_blocker']={'stage':'run_as_preflight','reason':'run-as unavailable for '+pkg};raise ProbeError(record['infrastructure_blocker']['reason'])
        if a.stage_file:
            if not a.stage_file.is_file(): raise ProbeError('Stage file not found: '+str(a.stage_file))
            rel='files/kam-architecture-probe/driver-staged-fixture.bin';adb.run('shell','run-as',pkg,'sh','-c','mkdir -p files/kam-architecture-probe')
            p=subprocess.run(adb.base+['exec-in','run-as',pkg,'sh','-c','cat > '+rel],input=a.stage_file.read_bytes(),stdout=subprocess.PIPE,stderr=subprocess.PIPE);save_command(out,'stage.log',p)
            size=adb.text('shell','run-as',pkg,'stat','-c','%s',rel,check=False);record['fixture_staging_driver']={'status':'PASS' if p.returncode==0 and size==str(a.stage_file.stat().st_size) else 'FAIL','source':str(a.stage_file),'bytes':a.stage_file.stat().st_size,'device_bytes':size}
        adb.run('shell','am','force-stop',pkg); listing=adb.text('shell','run-as',pkg,'find','.','-name','runtime.json',check=False)
        for name in own_report_paths(listing): adb.run('shell','run-as',pkg,'rm','--',name)
        adb.run('logcat','-c');record['launch_marker_epoch_ms']=int(time.time()*1000)
        component=adb.text('shell','cmd','package','resolve-activity','--brief','-a','android.intent.action.MAIN','-c','android.intent.category.LAUNCHER',pkg); candidates=[x.strip() for x in component.splitlines() if x.strip().startswith(pkg+'/')]
        if len(candidates)!=1 or not re.fullmatch(r'[\w.$/]+',candidates[0]): raise ProbeError('Cannot resolve probe launch Activity')
        launch=adb.run('shell','am','start','-W','-n',candidates[0],check=False);save_command(out,'launch.log',launch);record['launch']={'activity':candidates[0],'exit':launch.returncode,'result':launch.stdout.decode(errors='replace').strip()}
        if launch.returncode: raise ProbeError('Launch failed; see launch.log')
        data,path=wait_runtime(adb,pkg,lambda d:any(x.get('id')=='android.startup' and x.get('status')=='PASS' for x in d.get('tests',[])),a.seconds)
        if data is None: raise ProbeError('android.startup was not observed before timeout; see logcat.log and process.log')
        if a.exercise_lifecycle:
            adb.run('shell','input','keyevent','KEYCODE_HOME');time.sleep(1);adb.run('shell','am','start','-W','-n',candidates[0])
        expected=[]
        if a.exercise_lifecycle and a.profile=='platform': expected=['decision.schedule.callback_progress','decision.schedule.worker_progress','decision.schedule.lifecycle']
        if a.profile=='platform': expected += ['render.frame_callbacks','decision.text.cwstring_unit','decision.text.cp1251_decode','decision.text.cp1251_encode']
        if expected: data,path=wait_runtime(adb,pkg,lambda d: all(any(x.get('id')==name for x in d.get('tests',[])) for name in expected),a.seconds)
        if data is None: raise ProbeError('Runtime report disappeared; see logcat.log and process.log')
        out.joinpath('runtime.json').write_text(json.dumps(data,ensure_ascii=False,indent=2)+'\n',encoding='utf-8');record['runtime_path']=path
        record['errors']+=validate_runtime(data,package,require_frames=package['profile'] not in DECISION_PROFILES)
        if expected and not all(any(x.get('id')==name for x in data.get('tests',[])) for name in expected): record['errors'].append('Timed out waiting for observations: '+', '.join(expected))
        if a.stage_file and record.get('fixture_staging_driver',{}).get('status')!='PASS': record['errors'].append('Existing server fixture staging size mismatch')
        record['status']='FAIL' if record['errors'] else 'AUTOMATIC_SMOKE_PASS'
    except (ProbeError,OSError,ValueError,KeyError,subprocess.TimeoutExpired) as e:
        record['errors'].append(str(e))
        if record['preflight'].get('status')=='BLOCKED': record['infrastructure_blocker']={'stage':'device_preflight','reason':record['preflight'].get('reason',str(e))}
    finally:
        if adb:
            p=adb.run('shell','pidof',record.get('package',{}).get('package',''),check=False);save_command(out,'process.log',p);record['process_after_launch']=p.stdout.decode(errors='replace').strip()
            p=adb.run('logcat','-d','-v','threadtime',check=False);save_command(out,'logcat.log',p)
    write_json(out/'device-result.json',record);print(record['status']+'\n'+str(out));[print(x) for x in record['errors']]
    return 0 if record['status']=='AUTOMATIC_SMOKE_PASS' else 1 if record['status']=='FAIL' else 2
if __name__=='__main__': raise SystemExit(main())