"""Pull the entire DFWDS tree from AD1 to local."""
import base64, os, posixpath, paramiko, subprocess, yaml

LOCAL = r'D:\claudetools\projects\dataforth-dos\dfwds-research\source'
REMOTE_BASE = r'\\AD1\Engineering\ENGR\ATE\Test Datasheets\DFWDS'
AD2_STAGE = r'C:\Users\sysadmin\Documents\dfwds_stage'

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

c = paramiko.SSHClient(); c.set_missing_host_key_policy(paramiko.AutoAddPolicy())
c.connect('192.168.0.6', username='sysadmin', password=PWD,
          timeout=30, banner_timeout=45, look_for_keys=False, allow_agent=False)

def ps(cmd, to=300):
    enc = base64.b64encode(cmd.encode('utf-16-le')).decode()
    _, o, e = c.exec_command(f'powershell -NoProfile -EncodedCommand {enc}', timeout=to)
    return o.read().decode('utf-8','replace'), e.read().decode('utf-8','replace')

# 1. Stage on AD2 (recursive copy)
print('[1] copying tree to AD2 stage', flush=True)
out, err = ps(f'''
if (Test-Path "{AD2_STAGE}") {{ Remove-Item -Recurse -Force "{AD2_STAGE}" }}
Copy-Item -LiteralPath "{REMOTE_BASE}" -Destination "{AD2_STAGE}" -Recurse -Force -ErrorAction Stop
$f = Get-ChildItem -LiteralPath "{AD2_STAGE}" -Recurse -File
"copied: $($f.Count) files, $(($f | Measure-Object Length -Sum).Sum) bytes"
''')
print(out.rstrip())
if err.strip() and 'CLIXML' not in err: print('[stderr]', err[:400])

# 2. SFTP recursive download
print('\n[2] SFTP-pulling to local', flush=True)
sftp = c.open_sftp()

def walk_remote(base):
    """Yield (relpath, is_dir, size) for every entry under base."""
    stack = ['']
    while stack:
        rel = stack.pop()
        for entry in sftp.listdir_attr(f'{base}/{rel}' if rel else base):
            entry_rel = f'{rel}/{entry.filename}' if rel else entry.filename
            from stat import S_ISDIR
            is_dir = S_ISDIR(entry.st_mode)
            yield entry_rel, is_dir, entry.st_size
            if is_dir:
                stack.append(entry_rel)

stage_posix = AD2_STAGE.replace('\\','/')
n_files = 0
n_bytes = 0
for rel, is_dir, sz in walk_remote(stage_posix):
    local_path = os.path.join(LOCAL, rel.replace('/', os.sep))
    if is_dir:
        os.makedirs(local_path, exist_ok=True)
    else:
        os.makedirs(os.path.dirname(local_path), exist_ok=True)
        sftp.get(f'{stage_posix}/{rel}', local_path)
        n_files += 1
        n_bytes += sz

print(f'  pulled {n_files} files, {n_bytes:,} bytes', flush=True)
sftp.close()

# 3. Cleanup AD2 stage
print('\n[3] cleanup AD2 stage', flush=True)
ps(f'Remove-Item -Recurse -Force "{AD2_STAGE}"')
print('[OK]', flush=True)

c.close()