#[derive(Debug, Deserialize)]
struct WsQuery {
    token: Option<String>,
}
async fn websocket(
    State(state): State<AppState>,
    Query(query): Query<WsQuery>,
    ws: WebSocketUpgrade,
) -> Result<Response, AppError> {
    let expected = state.config.app_token.trim();
    if !expected.is_empty() && query.token.as_deref() != Some(expected) {
        return Err(AppError::Unauthorized);
    }
    Ok(ws.on_upgrade(move |socket| websocket_loop(state, socket)))
}

fn control_plan_ws_message(snapshot: &crate::state::ControlPlanSnapshot) -> Value {
    json!({
        "event": "control_plan.updated",
        "timestamp": Utc::now(),
        "data": {
            "revision": snapshot.revision,
            "plan": snapshot.plan.as_ref(),
        }
    })
}

async fn send_ws_bootstrap(state: &AppState, socket: &mut WebSocket) -> Result<Option<u64>, ()> {
    let (message, revision) = match build_bootstrap(state).await {
        Ok(data) => {
            let revision = Some(data.control_plan_revision);
            (
                json!({"event":"bootstrap","timestamp":Utc::now(),"data":data}),
                revision,
            )
        }
        Err(err) => (
            json!({"event":"error","timestamp":Utc::now(),"data":{"message":err.to_string()}}),
            None,
        ),
    };
    socket
        .send(Message::Text(message.to_string()))
        .await
        .map_err(|_| ())?;
    Ok(revision)
}

async fn websocket_loop(state: AppState, mut socket: WebSocket) {
    // Subscribe before building the bootstrap so state changes during bootstrap generation are
    // queued and can be applied immediately after the first frame.
    let mut receiver = state.events.subscribe();
    let mut control_plan = state.subscribe_control_plan();

    let bootstrap_revision = match send_ws_bootstrap(&state, &mut socket).await {
        Ok(revision) => revision,
        Err(()) => return,
    };
    let mut last_control_plan_revision = bootstrap_revision;

    // If the watch value is exactly the plan embedded in bootstrap, mark it seen to avoid a
    // duplicate control_plan.updated frame. A newer revision remains pending and is sent below.
    let current_revision = control_plan
        .borrow()
        .as_ref()
        .map(|snapshot| snapshot.revision);
    if current_revision == bootstrap_revision {
        control_plan.borrow_and_update();
    }

    loop {
        tokio::select! {
            event = receiver.recv() => {
                match event {
                    Ok(event) => {
                        if let Ok(text) = serde_json::to_string(&event) {
                            if socket.send(Message::Text(text)).await.is_err() { break; }
                        }
                    }
                    Err(tokio::sync::broadcast::error::RecvError::Lagged(skipped)) => {
                        tracing::warn!(skipped, "websocket client lagged; sending full bootstrap resync");
                        // Drop the retained stale backlog before taking the replacement snapshot.
                        // Events created while bootstrap is built are queued on this fresh receiver.
                        receiver = state.events.subscribe();
                        match send_ws_bootstrap(&state, &mut socket).await {
                            Ok(revision) => last_control_plan_revision = revision,
                            Err(()) => break,
                        }
                    }
                    Err(_) => break,
                }
            }
            changed = control_plan.changed() => {
                if changed.is_err() { break; }
                let snapshot = control_plan.borrow_and_update().clone();
                if let Some(snapshot) = snapshot {
                    if last_control_plan_revision == Some(snapshot.revision) {
                        continue;
                    }
                    let text = control_plan_ws_message(snapshot.as_ref()).to_string();
                    if socket.send(Message::Text(text)).await.is_err() { break; }
                    last_control_plan_revision = Some(snapshot.revision);
                }
            }
            message = socket.next() => {
                match message {
                    Some(Ok(Message::Ping(value))) => { if socket.send(Message::Pong(value)).await.is_err() { break; } }
                    Some(Ok(Message::Text(text))) if text == "ping" => { if socket.send(Message::Text("pong".into())).await.is_err() { break; } }
                    Some(Ok(Message::Close(_))) | None | Some(Err(_)) => break,
                    _ => {}
                }
            }
        }
    }
}