1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
use std::collections::HashMap;

use pyo3::{prelude::*, types::PyDict, PyObject, ToPyObject};

use crate::{
	callbacks::{interface::EventHandler, PyEventCallbackInterface},
	comm::{ChannelHandler, PyMessageInterface},
};

/// Configuration struct representing the overall setup for a system.
#[derive(Debug, Clone)]
#[pyclass]
pub struct Config {
	/// Version of the configuration format.
	pub version: f32,
	/// Unique identifier for the querent (user/client).
	pub querent_id: String,
	/// Name of the querent.
	pub querent_name: String,
	/// Configuration for the workflow.
	pub workflow: WorkflowConfig,
	/// List of collector configurations.
	pub collectors: Vec<CollectorConfig>,
	/// List of engine configurations.
	pub engines: Vec<EngineConfig>,
	/// Optional resource configuration.
	pub resource: Option<ResourceConfig>,
}

impl ToPyObject for Config {
	/// Converts a Config to a Python object.
	fn to_object(&self, py: Python) -> PyObject {
		let config_dict = PyDict::new(py);
		config_dict.set_item("version", self.version).unwrap();
		config_dict.set_item("querent_id", &self.querent_id).unwrap();
		config_dict.set_item("querent_name", &self.querent_name).unwrap();
		config_dict.set_item("workflow", &self.workflow).unwrap();
		config_dict.set_item("collectors", &self.collectors).unwrap();
		config_dict.set_item("engines", &self.engines).unwrap();
		config_dict.set_item("resource", &self.resource).unwrap();

		config_dict.to_object(py)
	}
}

impl Default for Config {
	/// Creates a default configuration.
	fn default() -> Self {
		Config {
			version: 0.1,
			querent_id: "querent".to_string(),
			querent_name: "Querent".to_string(),
			workflow: WorkflowConfig {
				name: "workflow".to_string(),
				id: "workflow".to_string(),
				config: HashMap::new(),
				inner_channel: None,
				channel: None,
				inner_event_handler: Some(EventHandler::new(None)),
				event_handler: None,
				inner_tokens_feader: None,
				tokens_feader: None,
			},
			collectors: vec![],
			engines: vec![],
			resource: None,
		}
	}
}

/// Configuration for a workflow.
#[derive(Debug, Clone)]
#[pyclass]
pub struct WorkflowConfig {
	/// Name of the workflow.
	pub name: String,
	/// Unique identifier for the workflow.
	pub id: String,
	/// Additional configuration options for the workflow.
	pub config: HashMap<String, String>,
	/// Internal channel handler in rust will be wrapped and marshalled into python.
	pub inner_channel: Option<ChannelHandler>,
	/// PyObject for the channel handler.
	/// This is a workaround for the fact that PyMessageInterface is not a PyObject.
	#[pyo3(get, set)]
	pub channel: Option<PyObject>,
	/// Inner EventHandler for workflow to get events from python
	pub inner_event_handler: Option<EventHandler>,
	/// PyObject for the event handler.
	#[pyo3(get, set)]
	pub event_handler: Option<PyObject>,
	/// Token feader for the engine for live tokens
	pub inner_tokens_feader: Option<ChannelHandler>,
	/// Token feeder for the engine for live tokens
	#[pyo3(get, set)]
	pub tokens_feader: Option<PyObject>,
}

impl ToPyObject for WorkflowConfig {
	/// Converts a WorkflowConfig to a Python object.
	fn to_object(&self, py: Python) -> PyObject {
		let workflow_dict = PyDict::new(py);
		workflow_dict.set_item("name", &self.name).unwrap();
		workflow_dict.set_item("id", &self.id).unwrap();
		workflow_dict.set_item("config", &self.config).unwrap();
		// convert channel handler to python object
		if let Some(inner_channel) = &self.inner_channel {
			let channel_interface = PyMessageInterface::new(inner_channel.clone());
			let channel: PyObject =
				Py::new(py, channel_interface).expect("Unable to create class").into_py(py);
			workflow_dict.set_item("channel", channel).unwrap();
		}
		// convert event handler to python object
		if let Some(inner_event_handler) = &self.inner_event_handler {
			let event_interface = PyEventCallbackInterface::new(inner_event_handler.clone());
			let event_handler: PyObject =
				Py::new(py, event_interface).expect("Unable to create class").into_py(py);
			workflow_dict.set_item("event_handler", event_handler).unwrap();
		}
		// convert token feeder to python object
		if let Some(inner_tokens_feader) = &self.inner_tokens_feader {
			let channel_interface = PyMessageInterface::new(inner_tokens_feader.clone());
			let tokens_feader: PyObject =
				Py::new(py, channel_interface).expect("Unable to create class").into_py(py);
			workflow_dict.set_item("tokens_feader", tokens_feader).unwrap();
		}
		workflow_dict.to_object(py)
	}
}

/// Configuration for a collector.
#[derive(Debug, Clone)]
#[pyclass]
pub struct CollectorConfig {
	/// Unique identifier for the collector.
	pub id: String,
	/// Name of the collector.
	pub name: String,
	/// Backend used by the collector.
	pub backend: String,
	/// Additional configuration options for the collector.
	pub config: HashMap<String, String>,
	/// Internal channel handler in rust will be wrapped and marshalled into python.
	pub inner_channel: Option<ChannelHandler>,
	/// PyObject for the channel handler.
	#[pyo3(get, set)]
	pub channel: Option<PyObject>,
}

impl ToPyObject for CollectorConfig {
	/// Converts a CollectorConfig to a Python object.
	fn to_object(&self, py: Python) -> PyObject {
		let collector_dict = PyDict::new(py);
		collector_dict.set_item("id", &self.id).unwrap();
		collector_dict.set_item("name", &self.name).unwrap();
		collector_dict.set_item("backend", &self.backend).unwrap();
		collector_dict.set_item("config", &self.config).unwrap();
		// convert channel handler to python object
		if let Some(inner_channel) = &self.inner_channel {
			let channel_interface = PyMessageInterface::new(inner_channel.clone());
			let channel: PyObject =
				Py::new(py, channel_interface).expect("Unable to create class").into_py(py);
			collector_dict.set_item("channel", channel).unwrap();
		}
		collector_dict.to_object(py)
	}
}

/// Configuration for an engine.
#[derive(Debug, Clone)]
#[pyclass]
pub struct EngineConfig {
	/// Unique identifier for the engine.
	pub id: String,
	/// Name of the engine.
	pub name: String,
	/// Config for the engine.
	pub config: HashMap<String, String>,
	/// Internal channel handler in rust will be wrapped and marshalled into python.
	pub inner_channel: Option<ChannelHandler>,
	/// PyObject for the channel handler.
	#[pyo3(get, set)]
	pub channel: Option<PyObject>,
}

impl ToPyObject for EngineConfig {
	/// Converts an EngineConfig to a Python object.
	fn to_object(&self, py: Python) -> PyObject {
		let engine_dict = PyDict::new(py);
		engine_dict.set_item("id", &self.id).unwrap();
		engine_dict.set_item("name", &self.name).unwrap();
		engine_dict.set_item("config", &self.config).unwrap();
		// convert channel handler to python object
		if let Some(inner_channel) = &self.inner_channel {
			let channel_interface = PyMessageInterface::new(inner_channel.clone());
			let channel: PyObject =
				Py::new(py, channel_interface).expect("Unable to create class").into_py(py);
			engine_dict.set_item("channel", channel).unwrap();
		}
		engine_dict.to_object(py)
	}
}

/// Configuration for resource constraints.
#[derive(Debug, Clone)]
pub struct ResourceConfig {
	/// Unique identifier for the resource.
	pub id: String,
	/// Maximum number of workers allowed (optional).
	pub max_workers_allowed: Option<u32>,
	/// Maximum number of workers per collector (optional).
	pub max_workers_per_collector: Option<u32>,
	/// Maximum number of workers per engine (optional).
	pub max_workers_per_engine: Option<u32>,
	/// Maximum number of workers per querent (optional).
	pub max_workers_per_querent: Option<u32>,
}

// Implementation of conversion traits for ResourceConfig.
impl<'a> FromPyObject<'a> for ResourceConfig {
	/// Extracts a ResourceConfig from a Python object.
	fn extract(ob: &'a PyAny) -> PyResult<Self> {
		let id = ob.getattr("id")?.extract()?;
		let max_workers_allowed = ob.getattr("max_workers_allowed")?.extract()?;
		let max_workers_per_collector = ob.getattr("max_workers_per_collector")?.extract()?;
		let max_workers_per_engine = ob.getattr("max_workers_per_engine")?.extract()?;
		let max_workers_per_querent = ob.getattr("max_workers_per_querent")?.extract()?;
		Ok(ResourceConfig {
			id,
			max_workers_allowed,
			max_workers_per_collector,
			max_workers_per_engine,
			max_workers_per_querent,
		})
	}
}

impl ToPyObject for ResourceConfig {
	/// Converts a ResourceConfig to a Python object.
	fn to_object(&self, py: Python) -> PyObject {
		let resource_dict = PyDict::new(py);
		resource_dict.set_item("id", &self.id).unwrap();
		resource_dict
			.set_item("max_workers_allowed", &self.max_workers_allowed)
			.unwrap();
		resource_dict
			.set_item("max_workers_per_collector", &self.max_workers_per_collector)
			.unwrap();
		resource_dict
			.set_item("max_workers_per_engine", &self.max_workers_per_engine)
			.unwrap();
		resource_dict
			.set_item("max_workers_per_querent", &self.max_workers_per_querent)
			.unwrap();

		resource_dict.to_object(py)
	}
}

#[pymethods]
impl Config {
	/// Constructor for creating a new Config instance.
	#[new]
	fn new(
		version: f32,
		querent_id: String,
		querent_name: String,
		workflow: WorkflowConfig,
		collectors: Vec<CollectorConfig>,
		engines: Vec<EngineConfig>,
		resource: Option<ResourceConfig>,
	) -> Self {
		Config { version, querent_id, querent_name, workflow, collectors, engines, resource }
	}
}