1use axum::{
6 extract::{Path, Query, State},
7 http::StatusCode,
8 response::Json,
9};
10use mockforge_core::consistency::ConsistencyEngine;
11use mockforge_core::snapshots::{ProtocolStateExporter, SnapshotComponents, SnapshotManager};
12use mockforge_core::workspace_persistence::WorkspacePersistence;
13use serde::Deserialize;
14use serde_json::Value;
15use std::sync::Arc;
16use tracing::{debug, error, info, warn};
17
18#[derive(Clone)]
20pub struct SnapshotState {
21 pub manager: Arc<SnapshotManager>,
23 pub consistency_engine: Option<Arc<ConsistencyEngine>>,
25 pub workspace_persistence: Option<Arc<WorkspacePersistence>>,
27 pub vbr_engine: Option<Arc<dyn ProtocolStateExporter>>,
30 pub recorder: Option<Arc<dyn ProtocolStateExporter>>,
33}
34
35#[derive(Debug, Deserialize)]
37pub struct SaveSnapshotRequest {
38 pub name: String,
40 pub description: Option<String>,
42 pub components: Option<SnapshotComponents>,
44}
45
46#[derive(Debug, Deserialize)]
48pub struct LoadSnapshotRequest {
49 pub components: Option<SnapshotComponents>,
51}
52
53#[derive(Debug, Deserialize)]
55pub struct SnapshotQuery {
56 #[serde(default = "default_workspace")]
58 pub workspace: String,
59}
60
61fn default_workspace() -> String {
62 "default".to_string()
63}
64
65async fn extract_vbr_state(vbr_engine: &Option<Arc<dyn ProtocolStateExporter>>) -> Option<Value> {
67 if let Some(engine) = vbr_engine {
68 match engine.export_state().await {
69 Ok(state) => {
70 let summary = engine.state_summary().await;
71 info!("Extracted VBR state from {} engine: {}", engine.protocol_name(), summary);
72 Some(state)
73 }
74 Err(e) => {
75 warn!("Failed to extract VBR state: {}", e);
76 None
77 }
78 }
79 } else {
80 debug!("No VBR engine available for state extraction");
81 None
82 }
83}
84
85async fn extract_recorder_state(
87 recorder: &Option<Arc<dyn ProtocolStateExporter>>,
88) -> Option<Value> {
89 if let Some(rec) = recorder {
90 match rec.export_state().await {
91 Ok(state) => {
92 let summary = rec.state_summary().await;
93 info!("Extracted Recorder state from {} engine: {}", rec.protocol_name(), summary);
94 Some(state)
95 }
96 Err(e) => {
97 warn!("Failed to extract Recorder state: {}", e);
98 None
99 }
100 }
101 } else {
102 debug!("No Recorder available for state extraction");
103 None
104 }
105}
106
107pub async fn save_snapshot(
111 State(state): State<SnapshotState>,
112 Query(params): Query<SnapshotQuery>,
113 Json(request): Json<SaveSnapshotRequest>,
114) -> Result<Json<Value>, StatusCode> {
115 let components = request.components.unwrap_or_else(SnapshotComponents::all);
116
117 let consistency_engine = state.consistency_engine.as_deref();
118 let workspace_persistence = state.workspace_persistence.as_deref();
119
120 let vbr_state = if components.vbr_state {
122 extract_vbr_state(&state.vbr_engine).await
123 } else {
124 None
125 };
126
127 let recorder_state = if components.recorder_state {
129 extract_recorder_state(&state.recorder).await
130 } else {
131 None
132 };
133 let manifest = state
134 .manager
135 .save_snapshot(
136 request.name.clone(),
137 request.description,
138 params.workspace.clone(),
139 components,
140 consistency_engine,
141 workspace_persistence,
142 vbr_state,
143 recorder_state,
144 )
145 .await
146 .map_err(|e| {
147 error!("Failed to save snapshot: {}", e);
148 StatusCode::INTERNAL_SERVER_ERROR
149 })?;
150
151 info!("Saved snapshot '{}' for workspace '{}'", request.name, params.workspace);
152 Ok(Json(serde_json::json!({
153 "success": true,
154 "manifest": manifest,
155 })))
156}
157
158pub async fn load_snapshot(
162 State(state): State<SnapshotState>,
163 Path(name): Path<String>,
164 Query(params): Query<SnapshotQuery>,
165 Json(request): Json<LoadSnapshotRequest>,
166) -> Result<Json<Value>, StatusCode> {
167 let consistency_engine = state.consistency_engine.as_deref();
168 let workspace_persistence = state.workspace_persistence.as_deref();
169 let (manifest, vbr_state, recorder_state) = state
170 .manager
171 .load_snapshot(
172 name.clone(),
173 params.workspace.clone(),
174 request.components,
175 consistency_engine,
176 workspace_persistence,
177 )
178 .await
179 .map_err(|e| {
180 error!("Failed to load snapshot: {}", e);
181 StatusCode::INTERNAL_SERVER_ERROR
182 })?;
183
184 info!("Loaded snapshot '{}' for workspace '{}'", name, params.workspace);
185 Ok(Json(serde_json::json!({
186 "success": true,
187 "manifest": manifest,
188 "vbr_state": vbr_state,
189 "recorder_state": recorder_state,
190 })))
191}
192
193pub async fn list_snapshots(
197 State(state): State<SnapshotState>,
198 Query(params): Query<SnapshotQuery>,
199) -> Result<Json<Value>, StatusCode> {
200 let snapshots = state.manager.list_snapshots(¶ms.workspace).await.map_err(|e| {
201 error!("Failed to list snapshots: {}", e);
202 StatusCode::INTERNAL_SERVER_ERROR
203 })?;
204
205 Ok(Json(serde_json::json!({
206 "workspace": params.workspace,
207 "snapshots": snapshots,
208 "count": snapshots.len(),
209 })))
210}
211
212pub async fn get_snapshot_info(
216 State(state): State<SnapshotState>,
217 Path(name): Path<String>,
218 Query(params): Query<SnapshotQuery>,
219) -> Result<Json<Value>, StatusCode> {
220 let manifest = state
221 .manager
222 .get_snapshot_info(name.clone(), params.workspace.clone())
223 .await
224 .map_err(|e| {
225 error!("Failed to get snapshot info: {}", e);
226 StatusCode::NOT_FOUND
227 })?;
228
229 Ok(Json(serde_json::json!({
230 "success": true,
231 "manifest": manifest,
232 })))
233}
234
235pub async fn delete_snapshot(
239 State(state): State<SnapshotState>,
240 Path(name): Path<String>,
241 Query(params): Query<SnapshotQuery>,
242) -> Result<Json<Value>, StatusCode> {
243 state
244 .manager
245 .delete_snapshot(name.clone(), params.workspace.clone())
246 .await
247 .map_err(|e| {
248 error!("Failed to delete snapshot: {}", e);
249 StatusCode::INTERNAL_SERVER_ERROR
250 })?;
251
252 info!("Deleted snapshot '{}' for workspace '{}'", name, params.workspace);
253 Ok(Json(serde_json::json!({
254 "success": true,
255 "message": format!("Snapshot '{}' deleted successfully", name),
256 })))
257}
258
259pub async fn validate_snapshot(
263 State(state): State<SnapshotState>,
264 Path(name): Path<String>,
265 Query(params): Query<SnapshotQuery>,
266) -> Result<Json<Value>, StatusCode> {
267 let is_valid = state
268 .manager
269 .validate_snapshot(name.clone(), params.workspace.clone())
270 .await
271 .map_err(|e| {
272 error!("Failed to validate snapshot: {}", e);
273 StatusCode::INTERNAL_SERVER_ERROR
274 })?;
275
276 Ok(Json(serde_json::json!({
277 "success": true,
278 "valid": is_valid,
279 "snapshot": name,
280 "workspace": params.workspace,
281 })))
282}
283
284pub fn snapshot_router(state: SnapshotState) -> axum::Router {
286 use axum::routing::{get, post};
287
288 axum::Router::new()
289 .route("/api/v1/snapshots", get(list_snapshots).post(save_snapshot))
290 .route("/api/v1/snapshots/{name}", get(get_snapshot_info).delete(delete_snapshot))
291 .route("/api/v1/snapshots/{name}/load", post(load_snapshot))
292 .route("/api/v1/snapshots/{name}/validate", get(validate_snapshot))
293 .with_state(state)
294}