{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Quick Start Guide"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Bailo python client enables intuitive interaction with the Bailo service, from within a python environment. This example notebook will run through the following concepts:\n",
    "\n",
    "* Searching for existing models on Bailo.\n",
    "* Creating a new model and populating its model card.\n",
    "* Managing collaborators for a model.\n",
    "* Creating releases and uploading files.\n",
    "* Downloading files from a release.\n",
    "\n",
    "Prerequisites:\n",
    "\n",
    "* Python 3.10 or higher (including a notebook environment for this demo).\n",
    "* A local or remote Bailo service (see https://github.com/gchq/Bailo)."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Bailo python client is split into two sub-packages: **core** and **helper**.\n",
    "\n",
    "* **Core:** For direct interactions with the service endpoints.\n",
    "* **Helper:** For more intuitive interactions with the service, using classes (e.g. Model) to handle operations.\n",
    "\n",
    "In order to create helper classes, you will first need to instantiate a `Client()` object from the core. By default, this object will not support any authentication. However, Bailo also supports PKI authentication, which you can use from Python by passing a `PkiAgent()` or `TokenAgent()` object into the `Client()` object when you instantiate it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Necessary import statements\n",
    "from bailo import Model, Client\n",
    "\n",
    "# Instantiating the PkiAgent(), if using.\n",
    "from bailo import PkiAgent\n",
    "\n",
    "agent = PkiAgent(cert=\"\", key=\"\", auth=\"\")\n",
    "# Instantiating the TokenAgent(), if using.\n",
    "from bailo import TokenAgent\n",
    "\n",
    "agent = TokenAgent(access_key=\"\", secret_key=\"\")\n",
    "\n",
    "# Instantiating the Bailo client\n",
    "client = Client(\"http://127.0.0.1:8080\", agent)  # <- INSERT BAILO URL (if not hosting locally)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Searching for models"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Marketplace is Bailo's home page for browsing models. Using the Python client, you can search for models programmatically with `Model.search()`. This returns a list of `Model()` objects matching your parameters.\n",
    "\n",
    "You can filter by task, libraries, search text, and more. In the example below, we first list all available models, then demonstrate filtering by a search term."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# List all available models\n",
    "models = Model.search(client=client)\n",
    "print(f\"Found {len(models)} model(s)\")\n",
    "\n",
    "# Search with a filter\n",
    "filtered_models = Model.search(client=client, search=\"ResNet\")\n",
    "print(f\"Found {len(filtered_models)} model(s) matching 'ResNet'\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating a new model"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this section, we will create a new model using the `Model.create()` classmethod. On the Bailo service, a model must consist of at least **name** and **description** parameters upon creation. Other attributes like model cards, files, or releases are added later on. Below, we use the `Client()` object created before when instantiating the new `Model()` object.\n",
    "\n",
    "NOTE: This creates the model on your Bailo service too! The `model_id` is assigned by the backend, and we will use this later to retrieve the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Model.create(client=client, name=\"Demo Model\", description=\"A quick-start demo model for Bailo.\")\n",
    "\n",
    "model_id = model.model_id\n",
    "print(f\"Created model {model_id}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You may make changes to these attributes and then call the `update()` method to relay the changes to the service, as below:\n",
    "\n",
    "```python\n",
    "model.name = \"New Name\"\n",
    "model.update()\n",
    "```"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Selecting a schema and filling in the model card"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When creating a model card, first we need to generate an empty one using the `card_from_schema()` method. In this instance, we will use **minimal-general-v10**. You can manage custom schemas using the `Schema()` helper class, but this is out of scope for this demo."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bailo.core.enums import MinimalSchema\n",
    "\n",
    "model.card_from_schema(schema_id=MinimalSchema.MODEL)\n",
    "\n",
    "print(f\"Model card version is {model.model_card_version}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If successful, the above will have created a new model card, and the `model_card_version` attribute should be set to 1.\n",
    "\n",
    "Next, we can populate the model card using the `update_model_card()` method. This can be used any time you want to make changes, and the backend will create a new model card version each time.\n",
    "\n",
    "NOTE: Your model card must match the schema, otherwise a `BailoException` will be raised."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_card = {\n",
    "    \"overview\": {\n",
    "        \"tags\": [\"quickstart\", \"demo\"],\n",
    "        \"modelSummary\": \"A quick-start demo model for Bailo.\",\n",
    "    }\n",
    "}\n",
    "\n",
    "model.update_model_card(model_card=new_card)\n",
    "\n",
    "print(f\"Model card version is {model.model_card_version}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Managing collaborators"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Collaborators control who can access and contribute to your model. Each collaborator is defined using a `CollaboratorEntry`, which links an entity (e.g. a user) to one or more roles. The available roles are:\n",
    "\n",
    "* **Owner**: Full control over the model.\n",
    "* **Contributor**: Can edit the model card and upload files.\n",
    "* **Consumer**: Can download files and images (after access approval).\n",
    "\n",
    "After setting collaborators on the model, call `model.update()` to push the changes to the service."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bailo.core.enums import CollaboratorEntry, Role\n",
    "\n",
    "model.collaborators = [\n",
    "    CollaboratorEntry(entity=\"user:user\", roles=[Role.OWNER]),\n",
    "]\n",
    "model.update()\n",
    "\n",
    "print(f\"Collaborators updated for model {model.model_id}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating a release"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "On the Bailo service, different versions of the same model are managed using **releases**. In this section we will create a **release** and then upload a file to it.\n",
    "\n",
    "`Release()` is a separate helper class in itself, but we can use our `Model()` object to create and retrieve releases. Running the below code will create a new release of the model, and return an instantiated `Release()` object which we will use to upload files with.\n",
    "\n",
    "NOTE: If the model does not have a populated model card schema then a `BailoException` will be thrown and the release will fail."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "release = model.create_release(version=\"1.0.0\", notes=\"Initial release of the demo model.\")\n",
    "\n",
    "print(f\"Created release {release.version} with notes {release.notes!r}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Uploading files to a release"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To upload files for a release, we can use the release `upload()` method which takes a file path. In this demo, we use the existing **demo_file.txt** in the notebooks directory.\n",
    "\n",
    "NOTE: The `upload()` method optionally takes the `data` parameter of `BytesIO` type to allow for other integrations, such as with S3. If the `path` parameter is a directory then `upload()` will automatically create a zip of the directory, allowing for multiple files to be uploaded as an archive.\n",
    "\n",
    "Alternatively, you can upload files outside of a release context using `client.simple_upload()`, which takes a `model_id`, file `name`, and a `BytesIO` buffer directly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "release.upload(path=\"demo_file.txt\")\n",
    "\n",
    "print(f\"Uploaded file to release {release.version}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Uploading large files with multipart upload\n",
    "\n",
    "For large files, the client supports **multipart upload** which splits the file into chunks and uploads them individually. This is more resilient for large transfers and is handled through three steps:\n",
    "\n",
    "1. **Start** the upload — the server returns a `fileId`, `uploadId`, and the byte ranges (`chunks`) for each part.\n",
    "2. **Upload each part** — send each chunk of data and collect the returned `ETag` for each part.\n",
    "3. **Finish** the upload — send all the part ETags to complete the file.\n",
    "\n",
    "These methods are available directly on the `Client` object. Below is a complete example using in-memory data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from io import BytesIO\n",
    "\n",
    "# Create some sample data (in practice, this would be a large file read in chunks)\n",
    "file_data = b\"x\" * 2048  # 2 KB of sample data\n",
    "file_size = len(file_data)\n",
    "\n",
    "# Step 1: Start the multipart upload\n",
    "start_response = client.start_multipart_upload(\n",
    "    model_id=model_id,\n",
    "    name=\"large_model_weights.bin\",\n",
    "    size=file_size,\n",
    ")\n",
    "\n",
    "file_id = start_response[\"fileId\"]\n",
    "upload_id = start_response[\"uploadId\"]\n",
    "chunks = start_response[\"chunks\"]\n",
    "\n",
    "print(f\"Started multipart upload: fileId={file_id}, {len(chunks)} chunk(s)\")\n",
    "\n",
    "# Step 2: Upload each part using the byte ranges from the start response\n",
    "parts = []\n",
    "for i, chunk in enumerate(chunks):\n",
    "    chunk_data = file_data[chunk[\"startByte\"] : chunk[\"endByte\"]]\n",
    "    part_response = client.upload_multipart_part(\n",
    "        model_id=model_id,\n",
    "        file_id=file_id,\n",
    "        upload_id=upload_id,\n",
    "        part_number=i + 1,\n",
    "        data=chunk_data,\n",
    "    )\n",
    "    parts.append({\"ETag\": part_response[\"ETag\"], \"PartNumber\": i + 1})\n",
    "    print(f\"  Uploaded part {i + 1}/{len(chunks)}\")\n",
    "\n",
    "# Step 3: Finish the upload\n",
    "finish_response = client.finish_multipart_upload(\n",
    "    model_id=model_id,\n",
    "    file_id=file_id,\n",
    "    upload_id=upload_id,\n",
    "    parts=parts,\n",
    ")\n",
    "\n",
    "print(f\"Multipart upload complete: {finish_response['file']['name']}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Downloading files from a release"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can download specific files from a release using the `download()` method, or download all files using the `download_all()` method.\n",
    "\n",
    "NOTE: For `download()`, the `filename` parameter refers to the filename on Bailo, and `path` is the local destination. For `download_all()`, you can provide `include` and `exclude` lists with glob patterns for filtering, e.g. `include=[\"*.txt\", \"*.json\"]`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "release.download_all(path=\"downloads\")\n",
    "\n",
    "print(\"Downloaded all release files to 'downloads' directory\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Retrieving an existing model"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this section, we will retrieve our previous model using the `Model.from_id()` classmethod. This will create your `Model()` object as before, but using existing information retrieved from the service."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Model.from_id(client=client, model_id=model_id)\n",
    "\n",
    "print(f\"Model description: {model.description}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What's next?\n",
    "\n",
    "For more detailed examples and advanced features, see the following notebooks:\n",
    "\n",
    "* **models_and_releases_demo_pytorch.ipynb** - In-depth model management with PyTorch integration.\n",
    "* **datacards_demo.ipynb** - Managing datacards.\n",
    "* **schemas_demo.ipynb** - Creating and managing custom schemas.\n",
    "* **access_requests_demo.ipynb** - Managing access requests.\n",
    "* **experiment_tracking_demo.ipynb** - Experiment tracking with MLFlow integration."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbformat_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}