use axum::{
extract::{Path, State},
response::Json,
};
use serde::{Deserialize, Serialize};
use super::state::MultiUserMemoryManager;
use crate::ab_testing;
use crate::errors::{AppError, ValidationErrorExt};
use crate::relevance;
use crate::validation;
use std::sync::Arc;
type AppState = Arc<MultiUserMemoryManager>;
fn default_traffic_split() -> f32 {
0.5
}
fn default_min_impressions() -> u64 {
100
}
#[derive(Debug, Deserialize)]
pub struct CreateABTestRequest {
pub name: String,
#[serde(default)]
pub description: Option<String>,
#[serde(default)]
pub control_weights: Option<relevance::LearnedWeights>,
#[serde(default)]
pub treatment_weights: Option<relevance::LearnedWeights>,
#[serde(default = "default_traffic_split")]
pub traffic_split: f32,
#[serde(default = "default_min_impressions")]
pub min_impressions: u64,
#[serde(default)]
pub max_duration_hours: Option<u64>,
#[serde(default)]
pub tags: Vec<String>,
}
#[derive(Debug, Serialize)]
pub struct ABTestResponse {
pub success: bool,
pub test_id: Option<String>,
pub message: String,
}
#[derive(Debug, Deserialize)]
pub struct RecordImpressionRequest {
pub user_id: String,
#[serde(default)]
pub relevance_score: Option<f64>,
#[serde(default)]
pub latency_us: Option<u64>,
}
#[derive(Debug, Deserialize)]
pub struct RecordClickRequest {
pub user_id: String,
pub memory_id: uuid::Uuid,
}
#[derive(Debug, Deserialize)]
pub struct RecordFeedbackRequest {
pub user_id: String,
pub positive: bool,
}
pub async fn list_ab_tests(
State(state): State<AppState>,
) -> Result<Json<serde_json::Value>, AppError> {
let tests = state.ab_test_manager.list_tests();
let summary = state.ab_test_manager.summary();
Ok(Json(serde_json::json!({
"success": true,
"tests": tests.iter().map(|t| serde_json::json!({
"id": t.id,
"name": t.config.name,
"description": t.config.description,
"status": format!("{:?}", t.status),
"traffic_split": t.config.traffic_split,
"control_impressions": t.control_metrics.impressions,
"treatment_impressions": t.treatment_metrics.impressions,
"created_at": t.created_at.to_rfc3339(),
})).collect::<Vec<_>>(),
"summary": {
"total_active": summary.total_active,
"draft": summary.draft,
"running": summary.running,
"paused": summary.paused,
"completed": summary.completed,
"archived": summary.archived,
}
})))
}
pub async fn create_ab_test(
State(state): State<AppState>,
Json(req): Json<CreateABTestRequest>,
) -> Result<Json<ABTestResponse>, AppError> {
let mut builder = ab_testing::ABTest::builder(&req.name)
.with_traffic_split(req.traffic_split)
.with_min_impressions(req.min_impressions);
if let Some(desc) = req.description {
builder = builder.with_description(&desc);
}
if let Some(control) = req.control_weights {
builder = builder.with_control(control);
}
if let Some(treatment) = req.treatment_weights {
builder = builder.with_treatment(treatment);
}
if let Some(hours) = req.max_duration_hours {
builder = builder.with_max_duration_hours(hours);
}
if !req.tags.is_empty() {
builder = builder.with_tags(req.tags);
}
let test = builder.build();
match state.ab_test_manager.create_test(test) {
Ok(test_id) => Ok(Json(ABTestResponse {
success: true,
test_id: Some(test_id),
message: "A/B test created successfully".to_string(),
})),
Err(e) => Ok(Json(ABTestResponse {
success: false,
test_id: None,
message: format!("Failed to create test: {}", e),
})),
}
}
pub async fn get_ab_test(
State(state): State<AppState>,
Path(test_id): Path<String>,
) -> Result<Json<serde_json::Value>, AppError> {
match state.ab_test_manager.get_test(&test_id) {
Some(test) => Ok(Json(serde_json::json!({
"success": true,
"test": {
"id": test.id,
"name": test.config.name,
"description": test.config.description,
"status": format!("{:?}", test.status),
"traffic_split": test.config.traffic_split,
"min_impressions": test.config.min_impressions,
"max_duration_hours": test.config.max_duration_hours,
"control_weights": test.config.control_weights,
"treatment_weights": test.config.treatment_weights,
"control_metrics": {
"impressions": test.control_metrics.impressions,
"clicks": test.control_metrics.clicks,
"ctr": if test.control_metrics.impressions > 0 {
test.control_metrics.clicks as f64 / test.control_metrics.impressions as f64
} else { 0.0 },
"positive_feedback": test.control_metrics.positive_feedback,
"negative_feedback": test.control_metrics.negative_feedback,
},
"treatment_metrics": {
"impressions": test.treatment_metrics.impressions,
"clicks": test.treatment_metrics.clicks,
"ctr": if test.treatment_metrics.impressions > 0 {
test.treatment_metrics.clicks as f64 / test.treatment_metrics.impressions as f64
} else { 0.0 },
"positive_feedback": test.treatment_metrics.positive_feedback,
"negative_feedback": test.treatment_metrics.negative_feedback,
},
"created_at": test.created_at.to_rfc3339(),
"started_at": test.started_at.map(|t| t.to_rfc3339()),
"completed_at": test.completed_at.map(|t| t.to_rfc3339()),
}
}))),
None => Ok(Json(serde_json::json!({
"success": false,
"message": format!("Test not found: {}", test_id)
}))),
}
}
pub async fn delete_ab_test(
State(state): State<AppState>,
Path(test_id): Path<String>,
) -> Result<Json<ABTestResponse>, AppError> {
match state.ab_test_manager.delete_test(&test_id) {
Ok(()) => Ok(Json(ABTestResponse {
success: true,
test_id: Some(test_id),
message: "Test deleted successfully".to_string(),
})),
Err(e) => Ok(Json(ABTestResponse {
success: false,
test_id: None,
message: format!("Failed to delete test: {}", e),
})),
}
}
pub async fn start_ab_test(
State(state): State<AppState>,
Path(test_id): Path<String>,
) -> Result<Json<ABTestResponse>, AppError> {
match state.ab_test_manager.start_test(&test_id) {
Ok(()) => Ok(Json(ABTestResponse {
success: true,
test_id: Some(test_id),
message: "Test started successfully".to_string(),
})),
Err(e) => Ok(Json(ABTestResponse {
success: false,
test_id: None,
message: format!("Failed to start test: {}", e),
})),
}
}
pub async fn pause_ab_test(
State(state): State<AppState>,
Path(test_id): Path<String>,
) -> Result<Json<ABTestResponse>, AppError> {
match state.ab_test_manager.pause_test(&test_id) {
Ok(()) => Ok(Json(ABTestResponse {
success: true,
test_id: Some(test_id),
message: "Test paused successfully".to_string(),
})),
Err(e) => Ok(Json(ABTestResponse {
success: false,
test_id: None,
message: format!("Failed to pause test: {}", e),
})),
}
}
pub async fn resume_ab_test(
State(state): State<AppState>,
Path(test_id): Path<String>,
) -> Result<Json<ABTestResponse>, AppError> {
match state.ab_test_manager.resume_test(&test_id) {
Ok(()) => Ok(Json(ABTestResponse {
success: true,
test_id: Some(test_id),
message: "Test resumed successfully".to_string(),
})),
Err(e) => Ok(Json(ABTestResponse {
success: false,
test_id: None,
message: format!("Failed to resume test: {}", e),
})),
}
}
pub async fn complete_ab_test(
State(state): State<AppState>,
Path(test_id): Path<String>,
) -> Result<Json<serde_json::Value>, AppError> {
match state.ab_test_manager.complete_test(&test_id) {
Ok(results) => Ok(Json(serde_json::json!({
"success": true,
"test_id": test_id,
"results": {
"is_significant": results.is_significant,
"confidence_level": results.confidence_level,
"chi_squared": results.chi_squared,
"p_value": results.p_value,
"winner": results.winner.map(|w| format!("{:?}", w)),
"relative_improvement": results.relative_improvement,
"control_ctr": results.control_ctr,
"treatment_ctr": results.treatment_ctr,
"confidence_interval": results.confidence_interval,
"recommendations": results.recommendations,
}
}))),
Err(e) => Ok(Json(serde_json::json!({
"success": false,
"message": format!("Failed to complete test: {}", e)
}))),
}
}
pub async fn analyze_ab_test(
State(state): State<AppState>,
Path(test_id): Path<String>,
) -> Result<Json<serde_json::Value>, AppError> {
match state.ab_test_manager.analyze_test(&test_id) {
Ok(results) => Ok(Json(serde_json::json!({
"success": true,
"test_id": test_id,
"analysis": {
"is_significant": results.is_significant,
"confidence_level": results.confidence_level,
"chi_squared": results.chi_squared,
"p_value": results.p_value,
"winner": results.winner.map(|w| format!("{:?}", w)),
"relative_improvement": results.relative_improvement,
"control_ctr": results.control_ctr,
"treatment_ctr": results.treatment_ctr,
"confidence_interval": results.confidence_interval,
"recommendations": results.recommendations,
}
}))),
Err(e) => Ok(Json(serde_json::json!({
"success": false,
"message": format!("Failed to analyze test: {}", e)
}))),
}
}
pub async fn record_ab_impression(
State(state): State<AppState>,
Path(test_id): Path<String>,
Json(req): Json<RecordImpressionRequest>,
) -> Result<Json<serde_json::Value>, AppError> {
validation::validate_user_id(&req.user_id).map_validation_err("user_id")?;
let relevance_score = req.relevance_score.unwrap_or(0.0);
let latency_us = req.latency_us.unwrap_or(0);
match state.ab_test_manager.record_impression(
&test_id,
&req.user_id,
relevance_score,
latency_us,
) {
Ok(()) => {
let variant = state
.ab_test_manager
.get_variant(&test_id, &req.user_id)
.ok();
Ok(Json(serde_json::json!({
"success": true,
"variant": variant.map(|v| format!("{:?}", v)),
})))
}
Err(e) => Ok(Json(serde_json::json!({
"success": false,
"message": format!("Failed to record impression: {}", e)
}))),
}
}
pub async fn record_ab_click(
State(state): State<AppState>,
Path(test_id): Path<String>,
Json(req): Json<RecordClickRequest>,
) -> Result<Json<serde_json::Value>, AppError> {
validation::validate_user_id(&req.user_id).map_validation_err("user_id")?;
match state
.ab_test_manager
.record_click(&test_id, &req.user_id, req.memory_id)
{
Ok(()) => Ok(Json(serde_json::json!({
"success": true,
}))),
Err(e) => Ok(Json(serde_json::json!({
"success": false,
"message": format!("Failed to record click: {}", e)
}))),
}
}
pub async fn record_ab_feedback(
State(state): State<AppState>,
Path(test_id): Path<String>,
Json(req): Json<RecordFeedbackRequest>,
) -> Result<Json<serde_json::Value>, AppError> {
validation::validate_user_id(&req.user_id).map_validation_err("user_id")?;
match state
.ab_test_manager
.record_feedback(&test_id, &req.user_id, req.positive)
{
Ok(()) => Ok(Json(serde_json::json!({
"success": true,
}))),
Err(e) => Ok(Json(serde_json::json!({
"success": false,
"message": format!("Failed to record feedback: {}", e)
}))),
}
}
pub async fn get_ab_summary(
State(state): State<AppState>,
) -> Result<Json<serde_json::Value>, AppError> {
let summary = state.ab_test_manager.summary();
let expired = state.ab_test_manager.check_expired_tests();
Ok(Json(serde_json::json!({
"success": true,
"summary": {
"total_active": summary.total_active,
"draft": summary.draft,
"running": summary.running,
"paused": summary.paused,
"completed": summary.completed,
"archived": summary.archived,
},
"expired_tests": expired,
})))
}