"""Small live S3 checks; deletes only the smoke object versions created here."""
import hashlib
import io
import os
import uuid

from botocore.exceptions import ClientError
from stream_model import client, stream_object


def main():
    s3 = client('https://s3.nl-ams.scw.cloud', 'nl-ams')
    bucket = os.environ.get('FI_S3_BUCKET', 'railiance-fi-open-weight-reserve')
    prefix = f'staging/stream-smoke/{uuid.uuid4()}'
    data = b'fi-stream-smoke\n' * 450000  # ~6.4 MiB, two parts
    spec = dict(path='payload', bytes=len(data), source_algorithm='sha256',
                source_digest=hashlib.sha256(data).hexdigest())
    for storage_class in ('STANDARD', 'ONEZONE_IA', 'GLACIER'):
        key = f'{prefix}/{storage_class}'
        receipt = stream_object(s3, bucket, key, io.BytesIO(data), spec,
                                storage_class, 5 * 1024 * 1024, {'fi-smoke': 'true'})
        try:
            head = s3.head_object(Bucket=bucket, Key=key)
            assert head.get('StorageClass', 'STANDARD') == storage_class
            if storage_class != 'GLACIER':
                response = s3.get_object(Bucket=bucket, Key=key)
                with response['Body'] as body:
                    assert hashlib.sha256(body.read()).hexdigest() == spec['source_digest']
            print(f'PASS {storage_class}: multipart + HEAD' +
                  (' + full readback SHA256' if storage_class != 'GLACIER' else ' (no restore)'))
        finally:
            s3.delete_object(Bucket=bucket, Key=key, VersionId=receipt['version_id'])
    key = f'{prefix}/bad-md5'
    upload = s3.create_multipart_upload(Bucket=bucket, Key=key)['UploadId']
    try:
        try:
            result = s3.upload_part(Bucket=bucket, Key=key, UploadId=upload, PartNumber=1,
                                    Body=b'payload', ContentMD5='AAAAAAAAAAAAAAAAAAAAAA==')
        except ClientError as exc:
            assert exc.response['Error']['Code'] == 'BadDigest', exc.response['Error']['Code']
            print('PASS server rejects corrupted multipart Content-MD5')
        else:
            assert result['ETag'].strip('"') == hashlib.md5(b'payload').hexdigest()
            print('NOTE server ignores Content-MD5; returned part ETag matches actual bytes. Collector checks ETags explicitly.')
    finally:
        s3.abort_multipart_upload(Bucket=bucket, Key=key, UploadId=upload)


if __name__ == '__main__':
    main()