manager_employees_rust/src/router.rs

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<DayPreferenceRequest>,
}

#[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<ScheduleResponse>,
}

#[derive(Serialize, utoipa::ToSchema)]
struct EmployeeNamesResponse {
    employees: Vec<String>,
}

#[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<HealthResponse> {
    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<AppState>,
    Json(payload): Json<CreateCompanyRequest>,
) -> ApiResult<Json<Company>> {
    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<AppState>) -> ApiResult<Json<Vec<Company>>> {
    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<AppState>,
    Path(id): Path<i32>,
) -> ApiResult<Json<Company>> {
    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<AppState>,
    Path(id): Path<i32>,
) -> ApiResult<Json<Company>> {
    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<AppState>,
    Json(payload): Json<CreateEmployeeRequest>,
) -> ApiResult<Json<EmployeeResponse>> {
    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<AppState>) -> ApiResult<Json<Vec<EmployeeResponse>>> {
    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<AppState>,
    Path((company_id, employee_id)): Path<(i32, i32)>,
) -> ApiResult<Json<CompanyUser>> {
    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<AppState>,
    Path(id): Path<i32>,
) -> ApiResult<Json<EmployeeResponse>> {
    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<AppState>,
    Path(id): Path<i32>,
    Json(payload): Json<RegisterPreferenceRequest>,
) -> ApiResult<Json<EmployeeResponse>> {
    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<AppState>,
    Path(id): Path<i32>,
    Json(payload): Json<WeeklyPreferenceRequest>,
) -> ApiResult<Json<EmployeeResponse>> {
    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<AppState>,
    Path(company_name): Path<String>,
) -> ApiResult<Json<EmployeeNamesResponse>> {
    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<AppState>,
    Json(payload): Json<CreateScheduleRequest>,
) -> ApiResult<Json<ScheduleResponse>> {
    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<AppState>,
) -> ApiResult<Json<WeeklyScheduleResponse>> {
    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<AppState>,
    Json(payload): Json<GenerateWeeklyScheduleRequest>,
) -> ApiResult<Json<WeeklyScheduleResponse>> {
    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<i32, i32> = 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<AppState>,
    Path(id): Path<i32>,
) -> ApiResult<Json<ScheduleResponse>> {
    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<AppState>,
    Path(id): Path<i32>,
    Json(payload): Json<ClockOutRequest>,
) -> ApiResult<Json<ScheduleResponse>> {
    let schedule =
        schedule_queries::clock_out(&state.pool, id, payload.day_of_week, payload.out_time).await?;
    Ok(Json(schedule_response(schedule)))
}