use axum::{ Json, Router, extract::{Path, State}, routing::{get, post, put}, }; use chrono::{Datelike, NaiveDateTime, Utc, Weekday}; use serde::{Deserialize, Serialize}; use sqlx::PgPool; use std::collections::HashMap; use tower_http::cors::CorsLayer; use utoipa::OpenApi; use utoipa_swagger_ui::SwaggerUi; use crate::{ app::{ApiResult, AppState}, database::{ entities::{Company, CompanyUser, Schedule, User, WeeklyPreference}, queries::{ company as company_queries, employee as employee_queries, preference as preference_queries, schedule as schedule_queries, user as user_queries, }, types::{DayOfWeek, WorkTime}, }, models::company::{requests::CreateCompanyRequest, service as company_service}, }; #[derive(Serialize, utoipa::ToSchema)] struct HealthResponse { status: &'static str, } #[derive(Deserialize, utoipa::ToSchema)] struct CreateEmployeeRequest { company_id: i32, name: String, last_name: String, preference: WorkTime, } #[derive(Deserialize, utoipa::ToSchema)] struct RegisterPreferenceRequest { preference: WorkTime, } #[derive(Deserialize, utoipa::ToSchema)] struct WeeklyPreferenceRequest { preferences: Vec, } #[derive(Deserialize, utoipa::ToSchema)] struct DayPreferenceRequest { day_of_week: DayOfWeek, first: WorkTime, second: WorkTime, third: WorkTime, } #[derive(Deserialize, utoipa::ToSchema)] struct CreateScheduleRequest { user_id: i32, enter_time: NaiveDateTime, day_of_week: DayOfWeek, out_time: NaiveDateTime, } #[derive(Deserialize, utoipa::ToSchema)] struct ClockOutRequest { day_of_week: DayOfWeek, out_time: NaiveDateTime, } #[derive(Deserialize, utoipa::ToSchema)] struct GenerateWeeklyScheduleRequest { company_id: i32, } #[derive(Serialize, utoipa::ToSchema)] struct EmployeeResponse { id: i32, company_id: i32, name: String, last_name: String, preference: WorkTime, } #[derive(Serialize, utoipa::ToSchema)] struct ScheduleResponse { id: i32, user_id: i32, enter_time: NaiveDateTime, day_of_week: DayOfWeek, out_time: NaiveDateTime, } #[derive(Serialize, utoipa::ToSchema)] struct WeeklyScheduleResponse { schedules: Vec, } #[derive(Serialize, utoipa::ToSchema)] struct EmployeeNamesResponse { employees: Vec, } #[derive(OpenApi)] #[openapi( paths( health_check, create_company, list_companies, get_company, delete_company, create_employee, list_employees, assign_employee_to_company, delete_employee, register_employee_preference, register_weekly_employee_preferences, get_employee_names_by_company, create_schedule, get_weekly_schedule, generate_weekly_schedule, clock_in, clock_out ), components( schemas( Company, CompanyUser, CreateCompanyRequest, HealthResponse, CreateEmployeeRequest, RegisterPreferenceRequest, WeeklyPreferenceRequest, DayPreferenceRequest, CreateScheduleRequest, ClockOutRequest, GenerateWeeklyScheduleRequest, EmployeeResponse, EmployeeNamesResponse, ScheduleResponse, WeeklyScheduleResponse, WorkTime, DayOfWeek ) ), tags( (name = "health", description = "Application health endpoints"), (name = "companies", description = "Company management endpoints"), (name = "employees", description = "Employee endpoints to implement"), (name = "schedules", description = "Schedule endpoints to implement") ), servers( (url = "{{rust_url}}", description = "Rust API base URL") ) )] struct ApiDoc; pub fn build_router(pool: PgPool) -> Router { Router::new() .route("/health", get(health_check)) .route("/companies", post(create_company).get(list_companies)) .route("/companies/{id}", get(get_company).delete(delete_company)) .route("/employees", post(create_employee).get(list_employees)) .route( "/companies/{company_id}/employees/{employee_id}", post(assign_employee_to_company), ) .route("/employees/{id}", axum::routing::delete(delete_employee)) .route( "/employees/{id}/preference", put(register_employee_preference), ) .route( "/employees/{id}/weekly-preferences", put(register_weekly_employee_preferences), ) .route( "/companies/{company_name}/employees/names", get(get_employee_names_by_company), ) .route("/schedules", post(create_schedule)) .route("/schedules/weekly", get(get_weekly_schedule)) .route("/schedules/weekly/generate", post(generate_weekly_schedule)) .route("/employees/{id}/clock-in", post(clock_in)) .route("/employees/{id}/clock-out", put(clock_out)) .merge(SwaggerUi::new("/swagger-ui").url("/api-docs/openapi.json", ApiDoc::openapi())) .layer(CorsLayer::permissive()) .with_state(AppState { pool }) } fn day_of_week_from_weekday(weekday: Weekday) -> DayOfWeek { match weekday { Weekday::Mon => DayOfWeek::Monday, Weekday::Tue => DayOfWeek::Tuesday, Weekday::Wed => DayOfWeek::Wednesday, Weekday::Thu => DayOfWeek::Thursday, Weekday::Fri => DayOfWeek::Friday, Weekday::Sat => DayOfWeek::Saturday, Weekday::Sun => DayOfWeek::Sunday, } } fn schedule_response(schedule: Schedule) -> ScheduleResponse { ScheduleResponse { id: schedule.id, user_id: schedule.user_id, enter_time: schedule.enter_time, day_of_week: schedule.day_of_week, out_time: schedule.out_time, } } fn employee_response(employee: employee_queries::EmployeeRow) -> EmployeeResponse { EmployeeResponse { id: employee.id, company_id: employee.company_id, name: employee.name, last_name: employee.last_name, preference: employee.preference, } } fn all_days() -> [DayOfWeek; 7] { [ DayOfWeek::Monday, DayOfWeek::Tuesday, DayOfWeek::Wednesday, DayOfWeek::Thursday, DayOfWeek::Friday, DayOfWeek::Saturday, DayOfWeek::Sunday, ] } fn all_shifts() -> [WorkTime; 3] { [WorkTime::Morning, WorkTime::Afternoon, WorkTime::Evening] } fn preference_rank( preference: Option<&WeeklyPreference>, shift: WorkTime, fallback: WorkTime, ) -> i32 { match preference { Some(preference) if preference.first_preference == shift => 0, Some(preference) if preference.second_preference == shift => 1, Some(preference) if preference.third_preference == shift => 2, _ if fallback == shift => 3, _ => 4, } } #[utoipa::path( get, path = "/health", tag = "health", responses( (status = 200, description = "Application is healthy", body = HealthResponse) ) )] async fn health_check() -> Json { Json(HealthResponse { status: "ok" }) } #[utoipa::path( post, path = "/companies", tag = "companies", request_body = CreateCompanyRequest, responses( (status = 200, description = "Company created", body = Company), (status = 500, description = "Database error") ) )] async fn create_company( State(state): State, Json(payload): Json, ) -> ApiResult> { let company = company_service::create_company(&state.pool, payload).await?; Ok(Json(company)) } #[utoipa::path( get, path = "/companies", tag = "companies", responses( (status = 200, description = "Companies found", body = [Company]), (status = 500, description = "Database error") ) )] async fn list_companies(State(state): State) -> ApiResult>> { let companies = company_service::list_companies(&state.pool).await?; Ok(Json(companies)) } #[utoipa::path( get, path = "/companies/{id}", tag = "companies", params( ("id" = i32, Path, description = "Company id") ), responses( (status = 200, description = "Company found", body = Company), (status = 404, description = "Company not found"), (status = 500, description = "Database error") ) )] async fn get_company( State(state): State, Path(id): Path, ) -> ApiResult> { let company = company_service::get_company(&state.pool, id).await?; Ok(Json(company)) } #[utoipa::path( delete, path = "/companies/{id}", tag = "companies", params( ("id" = i32, Path, description = "Company id") ), responses( (status = 200, description = "Company deleted", body = Company), (status = 404, description = "Company not found"), (status = 500, description = "Database error") ) )] async fn delete_company( State(state): State, Path(id): Path, ) -> ApiResult> { let company = company_service::delete_company(&state.pool, id).await?; Ok(Json(company)) } #[utoipa::path( post, path = "/employees", tag = "employees", request_body = CreateEmployeeRequest, responses( (status = 200, description = "Employee created", body = EmployeeResponse), (status = 500, description = "Database error") ) )] async fn create_employee( State(state): State, Json(payload): Json, ) -> ApiResult> { let user = user_queries::add_user( &state.pool, payload.name.trim(), payload.last_name.trim(), payload.preference, ) .await?; company_queries::add_company_user(&state.pool, payload.company_id, user.id).await?; let employee = employee_queries::get_employee(&state.pool, user.id).await?; Ok(Json(employee_response(employee))) } #[utoipa::path( get, path = "/employees", tag = "employees", responses( (status = 200, description = "Employees found", body = [EmployeeResponse]), (status = 500, description = "Database error") ) )] async fn list_employees(State(state): State) -> ApiResult>> { let employees = employee_queries::list_employees(&state.pool).await?; Ok(Json(employees.into_iter().map(employee_response).collect())) } #[utoipa::path( post, path = "/companies/{company_id}/employees/{employee_id}", tag = "employees", params( ("company_id" = i32, Path, description = "Company id"), ("employee_id" = i32, Path, description = "Employee id") ), responses( (status = 200, description = "Employee assigned to company", body = CompanyUser), (status = 404, description = "Company or employee not found"), (status = 500, description = "Database error") ) )] async fn assign_employee_to_company( State(state): State, Path((company_id, employee_id)): Path<(i32, i32)>, ) -> ApiResult> { company_queries::get_company_by_id(&state.pool, company_id).await?; employee_queries::get_employee_user(&state.pool, employee_id).await?; let company_user = company_queries::add_company_user(&state.pool, company_id, employee_id).await?; Ok(Json(company_user)) } #[utoipa::path( delete, path = "/employees/{id}", tag = "employees", params( ("id" = i32, Path, description = "Employee id") ), responses( (status = 200, description = "Employee deleted", body = EmployeeResponse), (status = 404, description = "Employee not found"), (status = 500, description = "Database error") ) )] async fn delete_employee( State(state): State, Path(id): Path, ) -> ApiResult> { let employee = employee_queries::get_employee(&state.pool, id).await?; user_queries::remove_user(&state.pool, id).await?; Ok(Json(employee_response(employee))) } #[utoipa::path( put, path = "/employees/{id}/preference", tag = "employees", params( ("id" = i32, Path, description = "Employee id") ), request_body = RegisterPreferenceRequest, responses( (status = 200, description = "Employee preference saved", body = EmployeeResponse), (status = 404, description = "Employee not found"), (status = 500, description = "Database error") ) )] async fn register_employee_preference( State(state): State, Path(id): Path, Json(payload): Json, ) -> ApiResult> { user_queries::update_preference(&state.pool, id, payload.preference).await?; let employee = employee_queries::get_employee(&state.pool, id).await?; Ok(Json(employee_response(employee))) } #[utoipa::path( put, path = "/employees/{id}/weekly-preferences", tag = "employees", params( ("id" = i32, Path, description = "Employee id") ), request_body = WeeklyPreferenceRequest, responses( (status = 200, description = "Weekly preferences saved", body = EmployeeResponse), (status = 404, description = "Employee not found"), (status = 500, description = "Database error") ) )] async fn register_weekly_employee_preferences( State(state): State, Path(id): Path, Json(payload): Json, ) -> ApiResult> { employee_queries::get_employee_user(&state.pool, id).await?; for preference in payload.preferences { preference_queries::upsert_weekly_preference( &state.pool, id, preference.day_of_week, preference.first, preference.second, preference.third, ) .await?; } let employee = employee_queries::get_employee(&state.pool, id).await?; Ok(Json(employee_response(employee))) } #[utoipa::path( get, path = "/companies/{company_name}/employees/names", tag = "employees", params( ("company_name" = String, Path, description = "Company name") ), responses( (status = 200, description = "Employee names for company", body = EmployeeNamesResponse), (status = 500, description = "Database error") ) )] async fn get_employee_names_by_company( State(state): State, Path(company_name): Path, ) -> ApiResult> { let employees = employee_queries::get_employee_names_by_company_name(&state.pool, &company_name).await?; Ok(Json(EmployeeNamesResponse { employees })) } #[utoipa::path( post, path = "/schedules", tag = "schedules", request_body = CreateScheduleRequest, responses( (status = 200, description = "Schedule created", body = ScheduleResponse), (status = 500, description = "Database error") ) )] async fn create_schedule( State(state): State, Json(payload): Json, ) -> ApiResult> { let schedule = schedule_queries::add_schedule( &state.pool, payload.user_id, payload.enter_time, payload.day_of_week, payload.out_time, ) .await?; Ok(Json(schedule_response(schedule))) } #[utoipa::path( get, path = "/schedules/weekly", tag = "schedules", responses( (status = 200, description = "Weekly schedule found", body = WeeklyScheduleResponse), (status = 500, description = "Database error") ) )] async fn get_weekly_schedule( State(state): State, ) -> ApiResult> { let schedules = schedule_queries::list_schedules(&state.pool).await?; Ok(Json(WeeklyScheduleResponse { schedules: schedules.into_iter().map(schedule_response).collect(), })) } #[utoipa::path( post, path = "/schedules/weekly/generate", tag = "schedules", request_body = GenerateWeeklyScheduleRequest, responses( (status = 200, description = "Weekly schedule generated", body = WeeklyScheduleResponse), (status = 500, description = "Database error") ) )] async fn generate_weekly_schedule( State(state): State, Json(payload): Json, ) -> ApiResult> { let employees = company_queries::get_users_by_company(&state.pool, payload.company_id).await?; let preferences = preference_queries::get_preferences_for_company(&state.pool, payload.company_id).await?; let mut preferences_by_user_day = HashMap::new(); for preference in &preferences { preferences_by_user_day.insert((preference.user_id, preference.day_of_week), preference); } schedule_queries::delete_schedules_for_company(&state.pool, payload.company_id).await?; let mut assigned_days: HashMap<(i32, DayOfWeek), bool> = HashMap::new(); let mut day_count_by_user: HashMap = HashMap::new(); let mut generated = Vec::new(); let now = Utc::now().naive_utc(); for day in all_days() { for shift in all_shifts() { let mut candidates: Vec<&User> = employees .iter() .filter(|employee| !assigned_days.contains_key(&(employee.id, day))) .filter(|employee| day_count_by_user.get(&employee.id).copied().unwrap_or(0) < 5) .collect(); candidates.sort_by_key(|employee| { preference_rank( preferences_by_user_day.get(&(employee.id, day)).copied(), shift, employee.preference, ) }); for employee in candidates.into_iter().take(2) { let shift_end = now + chrono::Duration::hours(8); let schedule = schedule_queries::add_schedule(&state.pool, employee.id, now, day, shift_end) .await?; assigned_days.insert((employee.id, day), true); *day_count_by_user.entry(employee.id).or_insert(0) += 1; generated.push(schedule_response(schedule)); } } } Ok(Json(WeeklyScheduleResponse { schedules: generated, })) } #[utoipa::path( post, path = "/employees/{id}/clock-in", tag = "schedules", params( ("id" = i32, Path, description = "Employee id") ), responses( (status = 200, description = "Clock-in registered", body = ScheduleResponse), (status = 404, description = "Employee not found"), (status = 500, description = "Database error") ) )] async fn clock_in( State(state): State, Path(id): Path, ) -> ApiResult> { let now_utc = Utc::now(); let day_of_week = day_of_week_from_weekday(now_utc.weekday()); let schedule = schedule_queries::clock_in(&state.pool, id, now_utc.naive_utc(), day_of_week).await?; Ok(Json(schedule_response(schedule))) } #[utoipa::path( put, path = "/employees/{id}/clock-out", tag = "schedules", params( ("id" = i32, Path, description = "Employee id") ), request_body = ClockOutRequest, responses( (status = 200, description = "Clock-out registered", body = ScheduleResponse), (status = 404, description = "Open clock-in record not found"), (status = 500, description = "Database error") ) )] async fn clock_out( State(state): State, Path(id): Path, Json(payload): Json, ) -> ApiResult> { let schedule = schedule_queries::clock_out(&state.pool, id, payload.day_of_week, payload.out_time).await?; Ok(Json(schedule_response(schedule))) }