"""Authenticated client: durable local write, SHA-256 verification, then server cleanup."""
import argparse
import getpass
import hashlib
import http.cookiejar
import json
import os
from pathlib import Path
import time
import urllib.error
import urllib.request


def generate_image(prompt, output, username=None, password=None, server='http://127.0.0.1:8189', size='1024x1024', seed=None):
    if not server.startswith(('http://127.0.0.1:','http://localhost:','https://')):
        raise ValueError('Use HTTPS or an SSH tunnel on localhost; passwords must not cross plain HTTP.')
    opener=urllib.request.build_opener(urllib.request.HTTPCookieProcessor(http.cookiejar.CookieJar()))
    csrf=''
    def call(path, body=None, binary=False):
        headers={'Content-Type':'application/json','X-CSRF-Token':csrf}
        req=urllib.request.Request(server.rstrip('/')+path, headers=headers,
            data=json.dumps(body).encode() if body is not None else None)
        try:
            with opener.open(req, timeout=60) as response:
                return response.read() if binary else json.load(response)
        except urllib.error.HTTPError as error:
            detail=error.read().decode('utf-8','replace')
            raise RuntimeError(f'HTTP {error.code}: {detail}') from None
    login=call('/auth/login',{'password':password or getpass.getpass('Access password: ')})
    csrf=login['csrf']
    body={'prompt':prompt,'size':size}
    if seed is not None:body['seed']=seed
    task=call('/api/jobs',body)
    image_id=task['id']
    print(f'Task: {image_id}',flush=True)
    for _ in range(1200):
        task=call('/api/jobs/'+image_id)
        if task['state']=='failed':raise RuntimeError(task['error'])
        if task['state']=='ready':break
        time.sleep(1)
    else:raise TimeoutError('Timed out; check pending tasks in the studio.')
    data=call(f'/api/jobs/{image_id}/image',binary=True)
    target=Path(output).expanduser().resolve()
    target.parent.mkdir(parents=True,exist_ok=True)
    # Avoid overwriting a different existing user file.
    expected=hashlib.sha256(data).hexdigest()
    with target.open('xb') as stream:
        stream.write(data)
        stream.flush()
        os.fsync(stream.fileno())
    if hashlib.sha256(target.read_bytes()).hexdigest()!=expected:
        raise RuntimeError('Local write verification failed; server cleanup was not acknowledged.')
    call(f'/api/jobs/{image_id}/ack',{'sha256':expected})
    call('/auth/logout',{})
    return {'output':str(target),'request_id':image_id,'server_cleaned':True,'seed':task['seed']}


if __name__=='__main__':
    parser=argparse.ArgumentParser()
    parser.add_argument('--server',default='http://127.0.0.1:8189')
    parser.add_argument('--username',help=argparse.SUPPRESS)
    parser.add_argument('--prompt',required=True)
    parser.add_argument('--output',required=True)
    parser.add_argument('--size',default='1024x1024')
    parser.add_argument('--seed',type=int)
    args=parser.parse_args()
    print(json.dumps(generate_image(**vars(args)),ensure_ascii=False,indent=2))