import os
from tempfile import NamedTemporaryFile

from orionclient.session import APISession
from orionclient.types import Shard, ShardCollection

# APISession uses the default profile, or the profile name defined in
# the ORION_PROFILE environment variable.
session = APISession

# create a new collection
collection = ShardCollection.create(
    session, "Descriptive name of new shard collection", metadata={"size_bytes": 0}
)
print(
    f"Collection '{collection.name}' can be referenced later using id={collection.id}"
)

for i in range(10):
    with NamedTemporaryFile("w+t") as temp:
        temp.write(f"data{i}")
        temp.flush()

        # create a new shard
        shard = Shard.create(
            collection,
            f"Descriptive name of shard {i}",
            metadata={"size_bytes": os.stat(temp.name).st_size},
        )
        print(
            f"Shard {i} in collection {collection.id} can be referenced later using id={shard.id}"
        )

        # upload a file to the newly created shard
        # file size of individual shard is currently limited to 5 GB
        shard.upload_file(temp.name)

        # by default, a shard is in an open state
        # a shard changes to the temporary state when data has been uploaded
        # shards can be marked ready immediately after upload in a serial cube,
        # however shards should be marked ready downstream of upload in a parallel cube
        shard.close()

        metadata = {
            "size_bytes": collection.metadata["size_bytes"]
            + shard.metadata["size_bytes"]
        }
        # Updates collection in place
        collection.update(metadata=metadata)

# prevent any additional shards from being added to the collection
collection.close()