use axum::{
    extract::{
        ws::{Message, WebSocket, WebSocketUpgrade},
        Path, State,
    },
    response::IntoResponse,
};
use std::sync::Arc;

use crate::state::{AppState, RESET_EVENT};

const USER_TYPING_START_PREFIX: &str = "typing:user:start:";
const USER_TYPING_STOP_PREFIX: &str = "typing:user:stop:";
const SUPPORTER_TYPING_START_PREFIX: &str = "typing:supporter:start:";
const SUPPORTER_TYPING_STOP_PREFIX: &str = "typing:supporter:stop:";

/// WebSocket for user — pushes notifications when supporter replies
pub async fn user_ws(
    ws: WebSocketUpgrade,
    Path(session_id): Path<String>,
    State(state): State<Arc<AppState>>,
) -> impl IntoResponse {
    ws.on_upgrade(move |socket| handle_user_socket(socket, session_id, state))
}

fn is_typing_event_for_session(text: &str, session_id: &str) -> bool {
    text.strip_prefix(USER_TYPING_START_PREFIX) == Some(session_id)
        || text.strip_prefix(USER_TYPING_STOP_PREFIX) == Some(session_id)
        || text.strip_prefix(SUPPORTER_TYPING_START_PREFIX) == Some(session_id)
        || text.strip_prefix(SUPPORTER_TYPING_STOP_PREFIX) == Some(session_id)
}

async fn handle_user_socket(mut socket: WebSocket, session_id: String, state: Arc<AppState>) {
    let tx = state.get_or_create_notifier(&session_id);
    let mut rx = tx.subscribe();

    loop {
        tokio::select! {
            // Forward broadcast events to the user browser.
            Ok(event) = rx.recv() => {
                if event == session_id || event == RESET_EVENT || is_typing_event_for_session(&event, &session_id) {
                    if socket.send(Message::Text(event)).await.is_err() {
                        break;
                    }
                }
            }
            // Handle ping/close and typing events from client.
            msg = socket.recv() => {
                match msg {
                    Some(Ok(Message::Text(text))) => {
                        if is_typing_event_for_session(&text, &session_id)
                            && (text.starts_with(USER_TYPING_START_PREFIX) || text.starts_with(USER_TYPING_STOP_PREFIX))
                        {
                            let _ = tx.send(text.clone());
                            let _ = state.supporter_notifier.send(text);
                        }
                    }
                    Some(Ok(Message::Close(_))) | None => break,
                    _ => {}
                }
            }
        }
    }
}

/// WebSocket for supporter dashboard — pushes notifications when any session updates
pub async fn supporter_ws(
    ws: WebSocketUpgrade,
    Path(_supporter_id): Path<String>,
    State(state): State<Arc<AppState>>,
) -> impl IntoResponse {
    ws.on_upgrade(move |socket| handle_supporter_socket(socket, state))
}

fn supporter_typing_session_id(text: &str) -> Option<&str> {
    text.strip_prefix(SUPPORTER_TYPING_START_PREFIX)
        .or_else(|| text.strip_prefix(SUPPORTER_TYPING_STOP_PREFIX))
}

async fn handle_supporter_socket(mut socket: WebSocket, state: Arc<AppState>) {
    let mut rx = state.supporter_notifier.subscribe();

    loop {
        tokio::select! {
            Ok(event) = rx.recv() => {
                if socket.send(Message::Text(event)).await.is_err() {
                    break;
                }
            }
            msg = socket.recv() => {
                match msg {
                    Some(Ok(Message::Text(text))) => {
                        if let Some(session_id) = supporter_typing_session_id(&text) {
                            if !session_id.is_empty() && state.sessions.contains_key(session_id) {
                                let tx = state.get_or_create_notifier(session_id);
                                let _ = tx.send(text);
                            }
                        }
                    }
                    Some(Ok(Message::Close(_))) | None => break,
                    _ => {}
                }
            }
        }
    }
}