// Copyright(c) 2022 Intel Corporation. All rights reserved.
 //
 //
+#include <linux/bitfield.h>
 #include <uapi/sound/sof/tokens.h>
 #include <sound/pcm_params.h>
 #include <sound/sof/ext_manifest4.h>
 static int sof_ipc4_widget_setup_comp_process(struct snd_sof_widget *swidget)
 {
        struct snd_soc_component *scomp = swidget->scomp;
+       struct sof_ipc4_fw_module *fw_module;
        struct sof_ipc4_process *process;
-       int cfg_size;
        void *cfg;
        int ret;
 
        if (ret)
                goto err;
 
-       cfg_size = sizeof(struct sof_ipc4_base_module_cfg);
+       ret = sof_ipc4_widget_setup_msg(swidget, &process->msg);
+       if (ret)
+               goto err;
 
-       cfg = kzalloc(cfg_size, GFP_KERNEL);
+       /* parse process init module payload config type from module info */
+       fw_module = swidget->module_info;
+       process->init_config = FIELD_GET(SOF_IPC4_MODULE_INIT_CONFIG_MASK,
+                                        fw_module->man4_module_entry.type);
+
+       process->ipc_config_size = sizeof(struct sof_ipc4_base_module_cfg);
+
+       /* allocate memory for base config extension if needed */
+       if (process->init_config == SOF_IPC4_MODULE_INIT_CONFIG_TYPE_BASE_CFG_WITH_EXT) {
+               struct sof_ipc4_base_module_cfg_ext *base_cfg_ext;
+               u32 ext_size = struct_size(base_cfg_ext, pin_formats,
+                                               swidget->num_input_pins + swidget->num_output_pins);
+
+               base_cfg_ext = kzalloc(ext_size, GFP_KERNEL);
+               if (!base_cfg_ext) {
+                       ret = -ENOMEM;
+                       goto free_available_fmt;
+               }
+
+               base_cfg_ext->num_input_pin_fmts = swidget->num_input_pins;
+               base_cfg_ext->num_output_pin_fmts = swidget->num_output_pins;
+               process->base_config_ext = base_cfg_ext;
+               process->base_config_ext_size = ext_size;
+               process->ipc_config_size += ext_size;
+       }
+
+       cfg = kzalloc(process->ipc_config_size, GFP_KERNEL);
        if (!cfg) {
                ret = -ENOMEM;
-               goto free_available_fmt;
+               goto free_base_cfg_ext;
        }
 
        process->ipc_config_data = cfg;
-       process->ipc_config_size = cfg_size;
-       ret = sof_ipc4_widget_setup_msg(swidget, &process->msg);
-       if (ret)
-               goto free_cfg_data;
 
        sof_ipc4_widget_update_kcontrol_module_id(swidget);
 
        return 0;
-free_cfg_data:
-       kfree(process->ipc_config_data);
-       process->ipc_config_data = NULL;
+free_base_cfg_ext:
+       kfree(process->base_config_ext);
+       process->base_config_ext = NULL;
 free_available_fmt:
        sof_ipc4_free_audio_fmt(&process->available_fmt);
 err:
                return;
 
        kfree(process->ipc_config_data);
+       kfree(process->base_config_ext);
        sof_ipc4_free_audio_fmt(&process->available_fmt);
        kfree(swidget->private);
        swidget->private = NULL;
        return 0;
 }
 
+static int
+sof_ipc4_process_set_pin_formats(struct snd_sof_widget *swidget, int pin_type)
+{
+       struct sof_ipc4_process *process = swidget->private;
+       struct sof_ipc4_base_module_cfg_ext *base_cfg_ext = process->base_config_ext;
+       struct sof_ipc4_available_audio_format *available_fmt = &process->available_fmt;
+       struct sof_ipc4_pin_format *pin_format, *format_list_to_search;
+       struct snd_soc_component *scomp = swidget->scomp;
+       int num_pins, format_list_count;
+       int pin_format_offset = 0;
+       int i, j;
+
+       /* set number of pins, offset of pin format and format list to search based on pin type */
+       if (pin_type == SOF_PIN_TYPE_INPUT) {
+               num_pins = swidget->num_input_pins;
+               format_list_to_search = available_fmt->input_pin_fmts;
+               format_list_count = available_fmt->num_input_formats;
+       } else {
+               num_pins = swidget->num_output_pins;
+               pin_format_offset = swidget->num_input_pins;
+               format_list_to_search = available_fmt->output_pin_fmts;
+               format_list_count = available_fmt->num_output_formats;
+       }
+
+       for (i = pin_format_offset; i < num_pins + pin_format_offset; i++) {
+               pin_format = &base_cfg_ext->pin_formats[i];
+
+               /* Pin 0 audio formats are derived from the base config input/output format */
+               if (i == pin_format_offset) {
+                       if (pin_type == SOF_PIN_TYPE_INPUT) {
+                               pin_format->buffer_size = process->base_config.ibs;
+                               pin_format->audio_fmt = process->base_config.audio_fmt;
+                       } else {
+                               pin_format->buffer_size = process->base_config.obs;
+                               pin_format->audio_fmt = process->output_format;
+                       }
+                       continue;
+               }
+
+               /*
+                * For all other pins, find the pin formats from those set in topology. If there
+                * is more than one format specified for a pin, this will pick the first available
+                * one.
+                */
+               for (j = 0; j < format_list_count; j++) {
+                       struct sof_ipc4_pin_format *pin_format_item = &format_list_to_search[j];
+
+                       if (pin_format_item->pin_index == i - pin_format_offset) {
+                               *pin_format = *pin_format_item;
+                               break;
+                       }
+               }
+
+               if (j == format_list_count) {
+                       dev_err(scomp->dev, "%s pin %d format not found for %s\n",
+                               (pin_type == SOF_PIN_TYPE_INPUT) ? "input" : "output",
+                               i - pin_format_offset, swidget->widget->name);
+                       return -EINVAL;
+               }
+       }
+
+       return 0;
+}
+
+static int sof_ipc4_process_add_base_cfg_extn(struct snd_sof_widget *swidget)
+{
+       int ret, i;
+
+       /* copy input and output pin formats */
+       for (i = 0; i <= SOF_PIN_TYPE_OUTPUT; i++) {
+               ret = sof_ipc4_process_set_pin_formats(swidget, i);
+               if (ret < 0)
+                       return ret;
+       }
+
+       return 0;
+}
+
 static int sof_ipc4_prepare_process_module(struct snd_sof_widget *swidget,
                                           struct snd_pcm_hw_params *fe_params,
                                           struct snd_sof_platform_stream_params *platform_params,
        if (ret < 0)
                return ret;
 
+       /* copy Pin 0 output format */
+       if (available_fmt->num_output_formats && ret < available_fmt->num_output_formats &&
+           !available_fmt->output_pin_fmts[ret].pin_index)
+               memcpy(&process->output_format, &available_fmt->output_pin_fmts[ret].audio_fmt,
+                      sizeof(struct sof_ipc4_audio_format));
+
        /* update pipeline memory usage */
        sof_ipc4_update_pipeline_mem_usage(sdev, swidget, &process->base_config);
 
-       /*
-        * ipc_config_data is composed of the base_config, optional output formats followed
-        * by the data required for module init in that order.
-        */
+       /* ipc_config_data is composed of the base_config followed by an optional extension */
        memcpy(cfg, &process->base_config, sizeof(struct sof_ipc4_base_module_cfg));
        cfg += sizeof(struct sof_ipc4_base_module_cfg);
 
+       if (process->init_config == SOF_IPC4_MODULE_INIT_CONFIG_TYPE_BASE_CFG_WITH_EXT) {
+               struct sof_ipc4_base_module_cfg_ext *base_cfg_ext = process->base_config_ext;
+
+               ret = sof_ipc4_process_add_base_cfg_extn(swidget);
+               if (ret < 0)
+                       return ret;
+
+               memcpy(cfg, base_cfg_ext, process->base_config_ext_size);
+       }
+
        return 0;
 }
 
 
 #define SOF_IPC4_MODULE_LL             BIT(5)
 #define SOF_IPC4_MODULE_DP             BIT(6)
 #define SOF_IPC4_MODULE_LIB_CODE               BIT(7)
+#define SOF_IPC4_MODULE_INIT_CONFIG_MASK       GENMASK(11, 8)
+
+#define SOF_IPC4_MODULE_INIT_CONFIG_TYPE_BASE_CFG              0
+#define SOF_IPC4_MODULE_INIT_CONFIG_TYPE_BASE_CFG_WITH_EXT     1
 
 #define SOF_IPC4_MODULE_INSTANCE_LIST_ITEM_SIZE 12
 #define SOF_IPC4_PIPELINE_OBJECT_SIZE 448
 /**
  * struct sof_ipc4_process - process config data
  * @base_config: IPC base config data
+ * @base_config_ext: Base config extension data for module init
  * @output_format: Output audio format
  * @available_fmt: Available audio format
  * @ipc_config_data: Process module config data
  * @ipc_config_size: Size of process module config data
  * @msg: IPC4 message struct containing header and data info
+ * @base_config_ext_size: Size of the base config extension data in bytes
+ * @init_config: Module init config type (SOF_IPC4_MODULE_INIT_CONFIG_TYPE_*)
  */
 struct sof_ipc4_process {
        struct sof_ipc4_base_module_cfg base_config;
+       struct sof_ipc4_base_module_cfg_ext *base_config_ext;
        struct sof_ipc4_audio_format output_format;
        struct sof_ipc4_available_audio_format available_fmt;
        void *ipc_config_data;
        uint32_t ipc_config_size;
        struct sof_ipc4_msg msg;
+       u32 base_config_ext_size;
+       u32 init_config;
 };
 
 #endif