{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Speckle Automate notebook example\n",
    "\n",
    "This is an example jupyter notebook, that will be executed by speckle automate\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Parameters\n",
    "\n",
    "These values are supplied to the notebook, when its being run by automate.\n",
    "\n",
    "**WARNING**: You really shouldn't modify this block, unless its following upstream changes from automate.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {
    "tags": [
     "parameters"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "from speckle_automate import AutomationContext, AutomationRunData\n",
    "\n",
    "function_inputs = '{\"forbiddenSpeckleType\": \"Objects.Geometry.Brep\"}'\n",
    "token_env_var = \"SPECKLE_TOKEN\"\n",
    "automation_run_data = ''\n",
    "\n",
    "execute_function = False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Local testing\n",
    "\n",
    "The cells below are here to make it easier to locally test the notebook.\n",
    "\n",
    "The cell has an option to register a new automation for the project on every run.\n",
    "\n",
    "With the local testing flag enabled, `local_testing = True`, the subsequent cell will override parameters from the cell above.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "local test registered\n"
     ]
    }
   ],
   "source": [
    "local_testing = False\n",
    "register_automation = True\n",
    "\n",
    "if local_testing:\n",
    "    execute_function = True\n",
    "    speckle_token = \"\"\n",
    "    from speckle_automate.helpers import register_new_automation, crypto_random_string\n",
    "\n",
    "    if not register_automation:\n",
    "        automation_id = \"\"\n",
    "        automation_revision_id = \"\"\n",
    "        automation_run_id = \"\"\n",
    "        function_id = \"\"\n",
    "        function_release = \"\"\n",
    "    else:\n",
    "        automation_id = crypto_random_string(10)\n",
    "        automation_revision_id = crypto_random_string(10)\n",
    "        automation_run_id = crypto_random_string(11)\n",
    "        function_id = crypto_random_string(10)\n",
    "        function_release = crypto_random_string(10)\n",
    "\n",
    "    automation_context = AutomationContext.initialize(\n",
    "        AutomationRunData(\n",
    "            project_id=\"3354d7be31\",\n",
    "            model_id=\"724617c963\",\n",
    "            branch_name=\"example.ifc\",\n",
    "            version_id=\"2a514e194f\",\n",
    "            speckle_server_url=\"https://latest.speckle.systems\",\n",
    "            automation_id=automation_id,\n",
    "            automation_revision_id=automation_revision_id,\n",
    "            automation_run_id=automation_run_id,\n",
    "            function_id=function_id,\n",
    "            function_release=function_release,\n",
    "        ),\n",
    "        speckle_token,\n",
    "    )\n",
    "\n",
    "    register_new_automation(\n",
    "        automation_context.speckle_client,\n",
    "        automation_context.automation_run_data.project_id,\n",
    "        automation_context.automation_run_data.model_id,\n",
    "        automation_context.automation_run_data.automation_id,\n",
    "        \"Jupyter local test\",\n",
    "        automation_context.automation_run_data.automation_revision_id,\n",
    "    )\n",
    "    print(\"local test registered\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Function inputs\n",
    "\n",
    "The `FunctionInputs` class defines the schema for the values, this function requires to run.\n",
    "These values will be provided by the users of the function.\n",
    "\n",
    "Automate uses the Json Schema, generated from the class, to validate user provided values.\n",
    "\n",
    "The schema block is tagged with the `function_input` tag, do not remove that!\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {
    "tags": [
     "function_inputs"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\"description\": \"These are function author defined values.\\n\\nAutomate will make sure to supply them matching the types specified here.\\nPlease use the pydantic model schema to define your inputs:\\nhttps://docs.pydantic.dev/latest/usage/models/\", \"properties\": {\"forbiddenSpeckleType\": {\"description\": \"If a object has the following speckle_type, it will be marked with an error.\", \"title\": \"Forbidden speckle type\", \"type\": \"string\"}}, \"required\": [\"forbiddenSpeckleType\"], \"title\": \"FunctionInputs\", \"type\": \"object\"}\n"
     ]
    }
   ],
   "source": [
    "from speckle_automate import AutomateBase\n",
    "from pydantic import Field\n",
    "import json\n",
    "\n",
    "\n",
    "class FunctionInputs(AutomateBase):\n",
    "    \"\"\"These are function author defined values.\n",
    "\n",
    "    Automate will make sure to supply them matching the types specified here.\n",
    "    Please use the pydantic model schema to define your inputs:\n",
    "    https://docs.pydantic.dev/latest/usage/models/\n",
    "    \"\"\"\n",
    "\n",
    "    forbidden_speckle_type: str = Field(\n",
    "        title=\"Forbidden speckle type\",\n",
    "        description=(\n",
    "            \"If a object has the following speckle_type,\"\n",
    "            \" it will be marked with an error.\"\n",
    "        ),\n",
    "    )\n",
    "\n",
    "\n",
    "print(json.dumps(FunctionInputs.model_json_schema()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Function logic\n",
    "\n",
    "in this block we're defining the actual business logic of the function.\n",
    "\n",
    "By all means modify this block but do keep the function signature compatible with the executing main function.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [],
   "source": [
    "from specklepy.objects import Base\n",
    "from typing import Iterable\n",
    "\n",
    "from speckle_automate import AutomationContext\n",
    "\n",
    "\n",
    "def flatten_base(base: Base) -> Iterable[Base]:\n",
    "    \"\"\"Take a base and flatten it to an iterable of bases.\"\"\"\n",
    "    if hasattr(base, \"elements\"):\n",
    "        for element in base[\"elements\"]:\n",
    "            yield from flatten_base(element)\n",
    "    yield base\n",
    "\n",
    "\n",
    "def automate_function(\n",
    "    automate_context: AutomationContext,\n",
    "    function_inputs: FunctionInputs,\n",
    "):\n",
    "    version_root_object = automate_context.receive_version()\n",
    "\n",
    "    count = 0\n",
    "    for b in flatten_base(version_root_object):\n",
    "        if b.speckle_type == function_inputs.forbidden_speckle_type:\n",
    "            if not b.id:\n",
    "                raise ValueError(\"Cannot operate on objects without their id's.\")\n",
    "            automate_context.add_object_error(\n",
    "                b.id,\n",
    "                \"This project should not contain the type: \" f\"{b.speckle_type}\",\n",
    "            )\n",
    "            count += 1\n",
    "\n",
    "    if count > 0:\n",
    "        # this is how a run is marked with a failure cause\n",
    "        automate_context.mark_run_failed(\n",
    "            \"Automation failed: \"\n",
    "            f\"Found {count} object that have one of the forbidden speckle types: \"\n",
    "            f\"{function_inputs.forbidden_speckle_type}\"\n",
    "        )\n",
    "\n",
    "    else:\n",
    "        automate_context.mark_run_success(\"No forbidden types found.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Function execution\n",
    "\n",
    "This last step executes the actual function using the automation context. The calculation cell must be tagged with `function_execution` tag, this allows the output to be discovered and reported to Speckle Automate.\n",
    "Do not change the below unless you know what you are doing!\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {
    "tags": [
     "function_execution"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "It took 2.040990 seconds to receive  the speckle version 2a514e194f\n",
      "Automation run SUCCEEDED after 2.044367 seconds.\n",
      "No forbidden types found.\n",
      "Run completed with status: AutomationStatus.SUCCEEDED, message: No forbidden types found.\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "from speckle_automate.runner import run_function\n",
    "\n",
    "if execute_function:\n",
    "    automation_context = run_function(\n",
    "        AutomationContext.initialize(automation_run_data, os.environ[token_env_var]),\n",
    "        automate_function,\n",
    "        FunctionInputs.model_validate_json(function_inputs),\n",
    "    )\n",
    "\n",
    "    # WARNING: Make sure the last print statement is this!\n",
    "    # automate relies on this convention, to process the automation results\n",
    "    print(automation_context._automation_result.model_dump_json())"
   ]
  }
 ],
 "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",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.4"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}