import os
import secrets
from concurrent.futures import ThreadPoolExecutor
import asyncio
import pytest
import concurrent
import functools
from hailtop.utils import secret_alnum_string, bounded_gather2
from hailtop.aiotools import LocalAsyncFS, RouterAsyncFS
from hailtop.aiocloud.aiogoogle import GoogleStorageClient, GoogleStorageAsyncFS


@pytest.fixture(params=['gs', 'router/gs'])
async def gs_filesystem(request):
    token = secret_alnum_string()

    with ThreadPoolExecutor() as thread_pool:
        if request.param.startswith('router/'):
            fs = RouterAsyncFS(
                'file', [LocalAsyncFS(thread_pool),
                         GoogleStorageAsyncFS()])
        else:
            assert request.param.endswith('gs')
            fs = GoogleStorageAsyncFS()
        async with fs:
            bucket = os.environ['HAIL_TEST_GCS_BUCKET']
            base = f'gs://{bucket}/tmp/{token}/'

            await fs.mkdir(base)
            sema = asyncio.Semaphore(50)
            async with sema:
                yield (sema, fs, base)
                await fs.rmtree(sema, base)
            assert not await fs.isdir(base)


@pytest.mark.asyncio
async def test_get_object_metadata():
    bucket = os.environ['HAIL_TEST_GCS_BUCKET']
    file = secrets.token_hex(16)

    async with GoogleStorageClient() as client:
        async with await client.insert_object(bucket, file) as f:
            await f.write(b'foo')
        metadata = await client.get_object_metadata(bucket, file)
        assert 'etag' in metadata
        assert metadata['md5Hash'] == 'rL0Y20zC+Fzt72VPzMSk2A=='
        assert metadata['crc32c'] == 'z8SuHQ=='
        assert int(metadata['size']) == 3


@pytest.mark.asyncio
async def test_get_object_headers():
    bucket = os.environ['HAIL_TEST_GCS_BUCKET']
    file = secrets.token_hex(16)

    async with GoogleStorageClient() as client:
        async with await client.insert_object(bucket, file) as f:
            await f.write(b'foo')
        async with await client.get_object(bucket, file) as f:
            headers = f.headers()
            assert 'ETag' in headers
            assert headers['X-Goog-Hash'] == 'crc32c=z8SuHQ==,md5=rL0Y20zC+Fzt72VPzMSk2A=='
            assert await f.read() == b'foo'


@pytest.mark.asyncio
async def test_compose():
    bucket = os.environ['HAIL_TEST_GCS_BUCKET']
    token = secret_alnum_string()

    part_data = [b'a', b'bb', b'ccc']

    async with GoogleStorageClient() as client:
        for i, b in enumerate(part_data):
            async with await client.insert_object(bucket, f'{token}/{i}') as f:
                await f.write(b)
        await client.compose(bucket, [f'{token}/{i}' for i in range(len(part_data))], f'{token}/combined')

        expected = b''.join(part_data)
        async with await client.get_object(bucket, f'{token}/combined') as f:
            actual = await f.read()
        assert actual == expected


@pytest.mark.asyncio
async def test_multi_part_create_many_two_level_merge(gs_filesystem):
    # This is a white-box test.  compose has a maximum of 32 inputs,
    # so if we're composing more than 32 parts, the
    # GoogleStorageAsyncFS does a multi-level hierarhical merge.
    try:
        sema, fs, base = gs_filesystem

        # > 32 so we perform at least 2 levels of merging
        part_data_size = [100 for _ in range(40)]
        part_data = [secrets.token_bytes(s) for s in part_data_size]

        s = 0
        part_start = []
        for b in part_data:
            part_start.append(s)
            s += len(b)

        path = f'{base}a'
        async with await fs.multi_part_create(sema, path, len(part_data)) as c:
            async def create_part(i):
                async with await c.create_part(i, part_start[i]) as f:
                    await f.write(part_data[i])

            # do in parallel
            await bounded_gather2(sema, *[
                functools.partial(create_part, i) for i in range(len(part_data))])

        expected = b''.join(part_data)
        actual = await fs.read(path)
        assert expected == actual
    except (concurrent.futures._base.CancelledError, asyncio.CancelledError) as err:
        raise AssertionError('uncaught cancelled error') from err

@@ -194,17 +194,41 @@
 _gather2
+, retry_transient_errors
 %0A
-
 from hai
@@ -1344,32 +1344,64 @@
 nt() as client:%0A
+        async def upload():%0A    
         async wi
@@ -1454,32 +1454,36 @@
  f:%0A            
+    
 await f.write(b'
@@ -1480,32 +1480,77 @@
 f.write(b'foo')%0A
+        await retry_transient_errors(upload)%0A
         metadata
@@ -1970,32 +1970,64 @@
 nt() as client:%0A
+        async def upload():%0A    
         async wi
@@ -2080,32 +2080,36 @@
  f:%0A            
+    
 await f.write(b'
@@ -2106,32 +2106,77 @@
 f.write(b'foo')%0A
+        await retry_transient_errors(upload)%0A
         async wi
@@ -2660,39 +2660,29 @@
     
-for i, b in enumerate(part_data
+async def upload(i, b
 ):%0A 
@@ -2752,32 +2752,32 @@
 en%7D/%7Bi%7D') as f:%0A
-
                 
@@ -2793,16 +2793,113 @@
 rite(b)%0A
+        for i, b in enumerate(part_data):%0A            await retry_transient_errors(upload, i, b)%0A
         

make aiogoogle tests resilient (#10947)