"""Prepare one file translation and resume it safely with a saved job ID.

Preparation only: python3 minimal_api_example.py source.docx --to zh-CN
Start/resume: add --start after reviewing the estimate.
Set EQUALANG_API_KEY in the environment.
"""
import argparse
import hashlib
import json
import os
import time
import uuid
from pathlib import Path
from urllib import error, request
from urllib.parse import urlsplit


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('source', type=Path)
    parser.add_argument('--to', required=True)
    parser.add_argument('--output', type=Path, default=Path('translated'))
    parser.add_argument('--state', type=Path, default=Path('api-job-state.json'))
    parser.add_argument('--start', action='store_true')
    parser.add_argument('--max-credits', type=float, default=1)
    args = parser.parse_args()
    base = os.environ.get('EQUALANG_BASE_URL', 'https://equalang.com/v1').rstrip('/')
    endpoint = urlsplit(base)
    if endpoint.scheme != 'https' or not endpoint.netloc or endpoint.username or endpoint.password or endpoint.query or endpoint.fragment:
        raise ValueError('Use a trusted HTTPS API endpoint without credentials or query parameters')
    key = os.environ.get('EQUALANG_API_KEY')
    if not key:
        raise ValueError('Set EQUALANG_API_KEY in the environment')

    def api(method, path, body=None, content_type='application/json', request_id=None):
        headers = {'Authorization': f'Bearer {key}', 'Content-Type': content_type}
        if request_id:
            headers['Idempotency-Key'] = request_id
        req = request.Request(base + path, data=body, headers=headers, method=method)
        with request.urlopen(req, timeout=90) as response:
            return json.loads(response.read())['data'], response.headers

    digest = hashlib.sha256(args.source.read_bytes()).hexdigest()
    identity = {'base': base, 'source_sha256': digest, 'target_language': args.to}
    state = json.loads(args.state.read_text()) if args.state.exists() else dict(identity)
    if any(state.get(k) != v for k, v in identity.items()):
        raise ValueError('State belongs to another input, language or endpoint; use another --state')

    def save():
        args.state.parent.mkdir(parents=True, exist_ok=True)
        args.state.write_text(json.dumps(state, indent=2) + '\n')

    if not state.get('file_id'):
        boundary = 'equalang-' + uuid.uuid4().hex
        filename = args.source.name.replace('"', '').replace('\r', '').replace('\n', '')
        body = (f'--{boundary}\r\nContent-Disposition: form-data; name="file"; filename="{filename}"\r\n'
                'Content-Type: application/octet-stream\r\n\r\n').encode() + args.source.read_bytes()
        body += f'\r\n--{boundary}--\r\n'.encode()
        upload, _ = api('POST', '/files', body, f'multipart/form-data; boundary={boundary}')
        state.update(file_id=upload['file_id'], estimate=upload['quote']['translate'])
        save()
    print(json.dumps({'estimated_credits': state['estimate']}))
    if not args.start:
        return
    if not isinstance(state['estimate'], (int, float)) or not 0 <= state['estimate'] <= args.max_credits:
        raise ValueError('Estimate exceeds --max-credits')
    if not state.get('job_id'):
        state.setdefault('request_id', str(uuid.uuid4()))
        save()  # Preserve idempotency even if the create response is lost.
        payload = json.dumps({'file_id': state['file_id'], 'target_language': args.to}).encode()
        created, _ = api('POST', '/jobs/translate', payload, request_id=state['request_id'])
        state['job_id'] = created['job_id']
        save()
    deadline = time.monotonic() + 120
    while True:
        job, headers = api('GET', '/jobs/' + state['job_id'])
        print(json.dumps({'job_id': job['job_id'], 'status': job['status']}))
        if job['finished']:
            break
        if time.monotonic() >= deadline:
            print('Still running. Repeat the same command to resume this job.')
            return
        time.sleep(min(30, max(1, int(headers.get('Retry-After', '5')))))
    if job['status'] != 'SUCCEEDED':
        raise RuntimeError('Job did not succeed; inspect its status before retrying')
    args.output.mkdir(parents=True, exist_ok=True)
    for item in job['outputs']:
        # Do not forward the API Authorization header to the signed download URL.
        with request.urlopen(item['download_url'], timeout=90) as response:
            content = response.read()
        path = args.output / Path(item['filename']).name
        if path.exists() and path.read_bytes() != content:
            raise ValueError('Output already exists with different content; use another --output')
        if not path.exists():
            path.write_bytes(content)
        print(json.dumps({'saved': str(path), 'sha256': hashlib.sha256(content).hexdigest()}))
    print(json.dumps({'credits_charged': job['usage']['credits_charged']}))


if __name__ == '__main__':
    try:
        main()
    except error.HTTPError as exc:
        raise SystemExit(f'API or download HTTP {exc.code}; saved state can be reused')
    except error.URLError:
        raise SystemExit('Network request failed; inspect saved state before retrying')