#!/usr/bin/env python3
"""Submit jobs to Michel's persistent queue over WinRM and poll result.
"""
import argparse, base64, json, os, sys, time, uuid
from pathlib import Path
try:
    import winrm
except ModuleNotFoundError:  # local Michel profile may not have pywinrm installed
    winrm = None

HOST = 'https://89.117.79.202:5986/wsman'
USER = 'Administrator'
PASS = 'pineapple'
REMOTE_IN = r'C:\MCSJAgent\MichelBridge\queue\in'
REMOTE_OUT = r'C:\MCSJAgent\MichelBridge\queue\out'
LOCAL_IN = Path(REMOTE_IN)
LOCAL_OUT = Path(REMOTE_OUT)


def sess():
    return winrm.Session(HOST, auth=(USER, PASS), transport='basic', server_cert_validation='ignore', operation_timeout_sec=30, read_timeout_sec=90)


def ps_quote(s):
    return "'" + s.replace("'", "''") + "'"


def submit_local(job, timeout=30):
    job.setdefault('id', 'job-' + uuid.uuid4().hex[:12])
    LOCAL_IN.mkdir(parents=True, exist_ok=True)
    LOCAL_OUT.mkdir(parents=True, exist_ok=True)
    name = job['id'] + '.json'
    in_path = LOCAL_IN / name
    out_path = LOCAL_OUT / name
    in_path.write_text(json.dumps(job, ensure_ascii=False, separators=(',', ':')), encoding='utf-8')
    deadline = time.time() + timeout
    while time.time() < deadline:
        if out_path.exists():
            return json.loads(out_path.read_text(encoding='utf-8'))
        time.sleep(0.25)
    return {'ok': False, 'job_id': job['id'], 'error': f'timeout waiting for {out_path}', 'submitted': job}


def submit(job, timeout=30):
    if winrm is None or os.environ.get('MICHELBRIDGE_LOCAL') == '1':
        return submit_local(job, timeout=timeout)
    job.setdefault('id', 'job-' + uuid.uuid4().hex[:12])
    name = job['id'] + '.json'
    data = json.dumps(job, ensure_ascii=False, separators=(',', ':')).encode('utf-8')
    b64 = base64.b64encode(data).decode('ascii')
    s = sess()
    remote_b64 = REMOTE_IN + '\\' + name + '.b64'
    remote_json = REMOTE_IN + '\\' + name
    # write b64 in chunks to avoid command-line length
    s.run_ps(f"New-Item -ItemType Directory -Force {ps_quote(REMOTE_IN)} | Out-Null; New-Item -ItemType Directory -Force {ps_quote(REMOTE_OUT)} | Out-Null; Set-Content -Path {ps_quote(remote_b64)} -Value '' -NoNewline")
    for i in range(0, len(b64), 900):
        r = s.run_ps(f"Add-Content -Path {ps_quote(remote_b64)} -Value {ps_quote(b64[i:i+900])} -NoNewline")
        if r.status_code:
            raise RuntimeError((r.std_err or b'').decode('utf-8','replace'))
    r = s.run_ps(f"$b=Get-Content -Raw {ps_quote(remote_b64)}; [IO.File]::WriteAllBytes({ps_quote(remote_json)}, [Convert]::FromBase64String($b)); Remove-Item {ps_quote(remote_b64)} -Force")
    if r.status_code:
        raise RuntimeError((r.std_err or b'').decode('utf-8','replace'))
    out_path = REMOTE_OUT + '\\' + name
    deadline = time.time() + timeout
    while time.time() < deadline:
        # Fetch as base64 so PowerShell/WinRM cannot mangle JSON escape sequences/newlines.
        r = s.run_ps(f"if(Test-Path {ps_quote(out_path)}) {{ [Convert]::ToBase64String([IO.File]::ReadAllBytes({ps_quote(out_path)})) }}")
        txt = (r.std_out or b'').decode('ascii','replace').strip()
        if txt:
            return json.loads(base64.b64decode(txt).decode('utf-8'))
        time.sleep(0.4)
    return {'ok': False, 'job_id': job['id'], 'error': f'timeout waiting for {out_path}', 'submitted': job}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('action')
    ap.add_argument('--job-json')
    ap.add_argument('--job-file')
    ap.add_argument('--timeout', type=int, default=30)
    args = ap.parse_args()
    if args.job_file:
        job = json.loads(Path(args.job_file).read_text())
    elif args.job_json:
        job = json.loads(args.job_json)
    else:
        job = {}
    job['action'] = args.action
    print(json.dumps(submit(job, timeout=args.timeout), ensure_ascii=False, indent=2))

if __name__ == '__main__':
    main()