"""Single-session SSH+SFTP batch fetcher for AD2 -> AD1 Engineering share."""
import os, sys, time, base64, paramiko

import subprocess, yaml as _yaml

HOST = '192.168.0.6'
USER = 'sysadmin'

def _pwd():
    r = subprocess.run(['sops','-d','D:/vault/clients/dataforth/ad2.sops.yaml'],
                       capture_output=True, text=True, timeout=30, check=True)
    return _yaml.safe_load(r.stdout)['credentials']['password'].replace('\\','')

PWD = _pwd()

LOCAL_ROOT = r'D:\claudetools\projects\dataforth-dos\datasheet-pipeline\scmvas-hvas-research'
REMOTE_BASE = r'\\AD1\Engineering\ENGR\ATE\High Voltage Input Module Test'
AD2_STAGE = r'C:\Users\sysadmin\Documents\scmvas_stage'

def connect():
    c = paramiko.SSHClient()
    c.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    c.connect(HOST, username=USER, password=PWD, timeout=30, banner_timeout=30, look_for_keys=False, allow_agent=False)
    return c

def run(c, cmd, timeout=120):
    stdin, stdout, stderr = c.exec_command(cmd, timeout=timeout)
    out = stdout.read().decode('utf-8', errors='replace')
    err = stderr.read().decode('utf-8', errors='replace')
    rc = stdout.channel.recv_exit_status()
    return rc, out, err

def ps(c, command):
    enc = base64.b64encode(command.encode('utf-16-le')).decode()
    return run(c, f'powershell -NoProfile -EncodedCommand {enc}', timeout=300)

def copy_one(c, remote_src, stage_dir):
    script = f'Copy-Item -LiteralPath "{remote_src}" -Destination "{stage_dir}\\" -Force -ErrorAction Stop; Write-Host "OK"'
    rc, out, err = ps(c, script)
    status = 'OK' if 'OK' in out and rc == 0 else f'FAIL: {err.strip() or out.strip()}'
    print(f'[{status}] {os.path.basename(remote_src)}')

def main():
    c = connect()
    try:
        ps(c, f'New-Item -ItemType Directory -Force -Path "{AD2_STAGE}" | Out-Null')

        files = [
            f'{REMOTE_BASE}\\TESTHV3.BAS',
            f'{REMOTE_BASE}\\LIBATE3.BAS',
            f'{REMOTE_BASE}\\DBHV.BAS',
            f'{REMOTE_BASE}\\Readme.txt',
            f'{REMOTE_BASE}\\HVDATA\\hvin.dat',
            f'{REMOTE_BASE}\\HVDATA\\hvsort.dat',
            f'{REMOTE_BASE}\\Released\\TESTHV3.BAS',
            f'{REMOTE_BASE}\\Released\\NLIBATE3.BAS',
            f'{REMOTE_BASE}\\Released\\TESTHV4.BAS',
            f'{REMOTE_BASE}\\Released\\TESTHV3.MAK',
        ]
        for f in files:
            copy_one(c, f, AD2_STAGE)

        # Rename to disambiguate source locations
        rc, out, err = ps(c, f'Get-ChildItem -LiteralPath "{AD2_STAGE}" | Select-Object Name,Length | Format-Table -AutoSize | Out-String')
        print('\n=== AD2 stage ===')
        print(out)

        # SFTP pull everything
        out_dir = os.path.join(LOCAL_ROOT, 'source')
        os.makedirs(out_dir, exist_ok=True)
        sftp = c.open_sftp()
        stage_posix = AD2_STAGE.replace('\\', '/')
        for e in sftp.listdir(stage_posix):
            sftp.get(f'{stage_posix}/{e}', os.path.join(out_dir, e))
            print(f'[pulled] {e}')
        sftp.close()
    finally:
        c.close()

if __name__ == '__main__':
    main()