From 64b85676ae9cdf428c337e1465a9e130c52dee64 Mon Sep 17 00:00:00 2001
From: Uri Shaked <uri@urishaked.com>
Date: Wed, 12 Feb 2025 11:37:36 +0200
Subject: [PATCH] feat(tools): add idf.py wokwi tool

Provides a quick way to launch Wokwi simulation for any esp-idf project.

This also provides an `idf.py wokwi-install` tool to install/upgrade Wokwi CLI automatically.
---
 tools/idf_py_actions/wokwi_ext.py | 86 +++++++++++++++++++++++++++++++
 1 file changed, 86 insertions(+)
 create mode 100644 tools/idf_py_actions/wokwi_ext.py

diff --git a/tools/idf_py_actions/wokwi_ext.py b/tools/idf_py_actions/wokwi_ext.py
new file mode 100644
index 0000000000..eb12377ecd
--- /dev/null
+++ b/tools/idf_py_actions/wokwi_ext.py
@@ -0,0 +1,86 @@
+#!/usr/bin/env python
+# SPDX-FileCopyrightText: 2025 Espressif Systems (Shanghai) CO LTD
+# SPDX-License-Identifier: Apache-2.0
+import os
+import re
+import subprocess
+import sys
+from typing import Dict
+from typing import List
+
+from click.core import Context
+from idf_py_actions.errors import FatalError
+from idf_py_actions.tools import PropertyDict
+from idf_py_actions.tools import red_print
+from packaging.version import Version
+
+WOKWI_CLI_PATH = os.getenv('WOKWI_CLI_PATH', 'wokwi-cli')
+WOKWI_CLI_MINIMUM_VERSION = '0.15.0'
+
+
+def check_wokwi_cli() -> bool:
+    try:
+        stdout = subprocess.check_output([WOKWI_CLI_PATH, '--version']).decode()
+        match = re.match(r'Wokwi CLI v(\d+\.\d+\.\d+)', stdout)
+        if not match:
+            raise FileNotFoundError('wokwi-cli not found.')
+        version = match.group(1)
+        if Version(version) < Version(WOKWI_CLI_MINIMUM_VERSION):
+            red_print(
+                f'wokwi-cli version {version} is less than minimum required version {WOKWI_CLI_MINIMUM_VERSION}'
+            )
+            red_print('Please update wokwi-cli by running: idf.py wokwi-install')
+            return False
+    except (subprocess.CalledProcessError, FileNotFoundError):
+        red_print('wokwi-cli not found.')
+        red_print('Please install it using: idf.py wokwi-install')
+        return False
+    return True
+
+
+def action_extensions(base_actions: Dict, project_path: List) -> Dict:
+    def install_wokwi_cli(action: str, ctx: Context, args: PropertyDict) -> None:
+        if sys.platform == 'win32':
+            cli_setup_cmd = 'irm https://wokwi.com/ci/install.ps1 | iex'
+        else:
+            cli_setup_cmd = 'curl -L https://wokwi.com/ci/install.sh | sh'
+        print('Running command: ', cli_setup_cmd)
+        if sys.platform == 'win32':
+            subprocess.run(['powershell', '-Command', cli_setup_cmd], check=True)
+        else:
+            subprocess.run(cli_setup_cmd, check=True, shell=True)
+
+    def wokwi(action: str, ctx: Context, args: PropertyDict, diagram_file: str) -> None:
+        if not check_wokwi_cli():
+            raise FatalError('Please run: idf.py wokwi-install')
+
+        wokwi_cmd = [WOKWI_CLI_PATH, args.project_dir]
+        if diagram_file:
+            wokwi_cmd.append('--diagram-file')
+            wokwi_cmd.append(diagram_file)
+
+        try:
+            subprocess.run(wokwi_cmd, check=True)
+        except subprocess.CalledProcessError as e:
+            raise FatalError(f'Error running wokwi-cli: {e}')
+
+    return {
+        'actions': {
+            'wokwi': {
+                'callback': wokwi,
+                'help': 'Launch current project in Wokwi simulator',
+                'options': [
+                    {
+                        'names': ['--diagram-file'],
+                        'help': 'Allows to specify a custom diagram file to use (e.g. --diagram-file diagram.esp32c3.json)',
+                    }
+                ],
+                'dependencies': ['all'],
+            },
+            'wokwi-install': {
+                'callback': install_wokwi_cli,
+                'help': 'Install wokwi-cli',
+                'options': [],
+            },
+        },
+    }