use super::GatewayState;
use crate::protocol::WsOutbound;
use crate::protocol::serialize_outbound;
use axum::extract::ws::{Message as WsMessage, WebSocket, WebSocketUpgrade};
use axum::extract::{Extension, Query, State};
use axum::response::Response;
use futures_util::{SinkExt, StreamExt};
use serde::Deserialize;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use tokio::sync::mpsc;
use tokio::time::{Duration, timeout};

#[derive(Debug, Default, Deserialize)]
pub struct WsQuery {
    client_id: Option<String>,
}

pub async fn ws_handler(
    ws: WebSocketUpgrade,
    Query(query): Query<WsQuery>,
    State(state): State<Arc<GatewayState>>,
    Extension(identity): Extension<super::auth::AuthIdentity>,
) -> Response {
    ws.on_upgrade(|socket| async move {
        handle_socket(socket, state, valid_client_id(query.client_id), identity).await;
    })
}

fn valid_client_id(client_id: Option<String>) -> Option<String> {
    client_id.filter(|value| {
        !value.is_empty()
            && value.len() <= 64
            && value
                .bytes()
                .all(|byte| byte.is_ascii_alphanumeric() || byte == b'-' || byte == b'_')
    })
}

async fn handle_socket(
    ws: WebSocket,
    state: Arc<GatewayState>,
    client_id: Option<String>,
    identity: super::auth::AuthIdentity,
) {
    let _connection_guard = ConnectionGuard::new(state.ws_connections.clone());

    // Create channel for sending outbound messages to this client
    let (sender, mut receiver) = mpsc::channel::<WsOutbound>(100);

    // Get CLI chat channel
    let cli_chat_channel = state.cli_chat_channel();

    // Register client with CliChatChannel and get initial session id
    let (session_id, client) = cli_chat_channel
        .register_client(sender.clone(), client_id)
        .await;
    // Send session established message
    let _ = sender
        .send(WsOutbound::SessionEstablished {
            session_id: session_id.clone(),
            capabilities: {
                let mut capabilities = vec!["turn_snapshots_v1".to_string()];
                if state.uploads.enabled() {
                    capabilities.push("file_transfer_v1".to_string());
                }
                capabilities
            },
        })
        .await;

    tracing::info!(session_id = %session_id, "CLI session established");

    let (mut ws_sender, mut ws_receiver) = ws.split();

    // Task: forward from receiver to WebSocket
    let mut writer_task = tokio::spawn(async move {
        while let Some(msg) = receiver.recv().await {
            if let Ok(text) = serialize_outbound(&msg)
                && ws_sender.send(WsMessage::Text(text.into())).await.is_err()
            {
                break;
            }
        }
    });

    // Main loop: receive WebSocket messages and forward to CliChatChannel
    let cancellation = state.connection_shutdown.clone();
    let mut writer_finished = false;
    let mut auth_check = tokio::time::interval(Duration::from_secs(5));
    auth_check.tick().await;
    loop {
        tokio::select! {
            _ = cancellation.cancelled() => break,
            _ = auth_check.tick() => {
                if !state.auth.identity_is_active(&identity).await {
                    tracing::info!(session_id = %session_id, "WebSocket authorization was revoked");
                    break;
                }
            }
            result = &mut writer_task => {
                writer_finished = true;
                if let Err(error) = result {
                    tracing::warn!(session_id = %session_id, error = %error, "WebSocket writer task failed");
                }
                break;
            }
            msg = ws_receiver.next() => {
                match msg {
                    Some(Ok(WsMessage::Text(text))) => {
                        cli_chat_channel.handle_inbound(client.clone(), &text).await;
                    }
                    Some(Ok(WsMessage::Close(_))) | Some(Err(_)) | None => {
                        tracing::debug!(session_id = %session_id, "WebSocket closed");
                        break;
                    }
                    _ => {}
                }
            }
        }
    }

    cli_chat_channel.unregister_client(&client).await;
    drop(client);
    drop(sender);
    if !writer_finished
        && timeout(Duration::from_secs(2), &mut writer_task)
            .await
            .is_err()
    {
        writer_task.abort();
        let _ = writer_task.await;
    }
    tracing::info!(session_id = %session_id, "CLI session ended");
}

struct ConnectionGuard {
    counter: Arc<AtomicUsize>,
}

impl ConnectionGuard {
    fn new(counter: Arc<AtomicUsize>) -> Self {
        counter.fetch_add(1, Ordering::Relaxed);
        Self { counter }
    }
}

impl Drop for ConnectionGuard {
    fn drop(&mut self) {
        self.counter.fetch_sub(1, Ordering::Relaxed);
    }
}

#[derive(Debug, Default, Deserialize)]
pub struct WsLogsQuery {
    level: Option<String>,
    search: Option<String>,
}

pub async fn ws_logs_handler(
    ws: WebSocketUpgrade,
    Query(query): Query<WsLogsQuery>,
    Extension(_identity): Extension<super::auth::AuthIdentity>,
) -> Response {
    ws.on_upgrade(|socket| async move {
        handle_logs_socket(socket, query).await;
    })
}

async fn handle_logs_socket(ws: WebSocket, query: WsLogsQuery) {
    let Some(tx) = crate::logging::log_sender() else {
        return;
    };
    let mut rx = tx.subscribe();
    let (mut ws_sender, mut ws_receiver) = ws.split();

    let min_level = query.level.as_deref().map(parse_min_level).unwrap_or(0);
    let search = query
        .search
        .filter(|s| !s.is_empty())
        .map(|s| s.to_ascii_lowercase());

    loop {
        tokio::select! {
            result = rx.recv() => {
                match result {
                    Ok(event) => {
                        if level_rank(&event.level) < min_level {
                            continue;
                        }
                        if let Some(needle) = &search
                            && !event.message.to_ascii_lowercase().contains(needle)
                            && !event.target.to_ascii_lowercase().contains(needle)
                        {
                            continue;
                        }
                        let json = serde_json::json!({
                            "ts": event.ts,
                            "level": event.level,
                            "target": event.target,
                            "message": event.message,
                        });
                        if ws_sender
                            .send(WsMessage::Text(json.to_string().into()))
                            .await
                            .is_err()
                        {
                            break;
                        }
                    }
                    Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
                    Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
                }
            }
            msg = ws_receiver.next() => {
                match msg {
                    Some(Ok(WsMessage::Close(_))) | None => break,
                    _ => {}
                }
            }
        }
    }
}

fn parse_min_level(level: &str) -> u8 {
    match level.to_ascii_uppercase().as_str() {
        "DEBUG" => 1,
        "INFO" => 2,
        "WARN" => 3,
        "ERROR" => 4,
        _ => 0,
    }
}

fn level_rank(level: &str) -> u8 {
    match level {
        "TRACE" => 0,
        "DEBUG" => 1,
        "INFO" => 2,
        "WARN" => 3,
        "ERROR" => 4,
        _ => 2,
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn client_id_is_strictly_bounded() {
        assert_eq!(
            valid_client_id(Some("client_123-abc".to_string())).as_deref(),
            Some("client_123-abc")
        );
        assert!(valid_client_id(Some("bad/query".to_string())).is_none());
        assert!(valid_client_id(Some("x".repeat(65))).is_none());
        assert!(valid_client_id(Some(String::new())).is_none());
    }

    #[test]
    fn connection_guard_tracks_its_scope() {
        let connections = Arc::new(AtomicUsize::new(0));
        {
            let _guard = ConnectionGuard::new(connections.clone());
            assert_eq!(connections.load(Ordering::Relaxed), 1);
        }
        assert_eq!(connections.load(Ordering::Relaxed), 0);
    }
}
