use axum::{Json, extract::State, http::StatusCode, response::IntoResponse};
use common_contracts::registry::{Heartbeat, ServiceRegistration};
use std::time::Instant;
use tracing::{info, warn};

use crate::{
    error::Result,
    state::{AppState, RegistryEntry},
};

/// [POST /v1/registry/register]
pub async fn register_service(
    State(state): State<AppState>,
    Json(payload): Json<ServiceRegistration>,
) -> Result<impl IntoResponse> {
    info!(
        "Registering service: {} ({}) at {}",
        payload.service_id, payload.service_name, payload.base_url
    );

    let entry = RegistryEntry {
        registration: payload.clone(),
        last_heartbeat: Instant::now(),
    };

    let mut registry = state.registry.write().unwrap();
    registry.insert(payload.service_id.clone(), entry);

    Ok(StatusCode::OK)
}

/// [POST /v1/registry/heartbeat]
pub async fn heartbeat(
    State(state): State<AppState>,
    Json(payload): Json<Heartbeat>,
) -> Result<impl IntoResponse> {
    let mut registry = state.registry.write().unwrap();

    if let Some(entry) = registry.get_mut(&payload.service_id) {
        entry.last_heartbeat = Instant::now();
        Ok(StatusCode::OK)
    } else {
        // This is the key part for self-healing: tell the provider we don't know them
        warn!(
            "Received heartbeat from unknown service: {}",
            payload.service_id
        );
        Ok(StatusCode::NOT_FOUND)
    }
}

/// [POST /v1/registry/deregister]
pub async fn deregister_service(
    State(state): State<AppState>,
    Json(payload): Json<serde_json::Value>,
) -> Result<impl IntoResponse> {
    let service_id = payload
        .get("service_id")
        .and_then(|v| v.as_str())
        .ok_or_else(|| crate::error::AppError::BadRequest("Missing service_id".into()))?;

    info!("Deregistering service: {}", service_id);

    let mut registry = state.registry.write().unwrap();
    registry.remove(service_id);

    Ok(StatusCode::OK)
}
