use std::sync::Arc;

use axum::{
    extract::{
        ws::{Message, WebSocket},
        Query, State, WebSocketUpgrade,
    },
    response::Response,
    routing::get,
    Router,
};
use futures::{sink::SinkExt, stream::StreamExt};
use serde::Deserialize;

use crate::models::error::ApiError;
use crate::services::auth as auth_service;
use crate::AppState;

pub fn router() -> Router<Arc<AppState>> {
    Router::new().route("/", get(ws_handler))
}

#[derive(Deserialize)]
struct WsQuery {
    token: String,
}

/// WebSocket upgrade handler
///
/// Authenticates via JWT token in query param, then upgrades to WebSocket.
/// Subscribes to NATS events for the user's license and forwards to client.
async fn ws_handler(
    ws: WebSocketUpgrade,
    Query(query): Query<WsQuery>,
    State(state): State<Arc<AppState>>,
) -> Result<Response, ApiError> {
    // Validate JWT token
    let claims = auth_service::validate_token(&state.config, &query.token)
        .map_err(|_| ApiError::Unauthorized)?;

    let license_id = claims
        .license_id
        .ok_or(ApiError::LicenseInvalid)?;

    // Upgrade to WebSocket
    Ok(ws.on_upgrade(move |socket| handle_socket(socket, license_id, state)))
}

/// Handle WebSocket connection after upgrade
async fn handle_socket(socket: WebSocket, license_id: uuid::Uuid, state: Arc<AppState>) {
    let (mut sender, mut receiver) = socket.split();

    // Check if NATS is available
    let nats = match &state.nats {
        Some(client) => client.clone(),
        None => {
            tracing::warn!("WebSocket connected but NATS unavailable");
            let _ = sender
                .send(Message::Text(
                    serde_json::json!({
                        "type": "error",
                        "message": "Event bus unavailable"
                    })
                    .to_string()
                    .into(),
                ))
                .await;
            return;
        }
    };

    // Subscribe to license-scoped events
    // Pattern: corrosion.{license_id}.> (all events for this license)
    let subject = format!("corrosion.{}.*", license_id);
    let mut sub = match nats.subscribe(subject.clone()).await {
        Ok(s) => s,
        Err(e) => {
            tracing::error!("Failed to subscribe to NATS: {}", e);
            let _ = sender
                .send(Message::Text(
                    serde_json::json!({
                        "type": "error",
                        "message": "Failed to subscribe to events"
                    })
                    .to_string()
                    .into(),
                ))
                .await;
            return;
        }
    };

    tracing::info!(
        "WebSocket connected for license {} (subscribed to {})",
        license_id,
        subject
    );

    // Send welcome message
    let _ = sender
        .send(Message::Text(
            serde_json::json!({
                "type": "connected",
                "license_id": license_id,
                "subscribed_to": subject
            })
            .to_string()
            .into(),
        ))
        .await;

    // Spawn task to forward NATS messages to WebSocket
    let mut send_task = tokio::spawn(async move {
        while let Some(msg) = sub.next().await {
            let payload = String::from_utf8_lossy(&msg.payload).to_string();

            // Parse subject to extract event type
            // Format: corrosion.{license_id}.{event_type}
            let event_type = msg
                .subject
                .split('.')
                .nth(2)
                .unwrap_or("unknown")
                .to_string();

            // Send to WebSocket client
            let ws_msg = serde_json::json!({
                "type": "event",
                "event": event_type,
                "subject": msg.subject.as_str(),
                "data": serde_json::from_str::<serde_json::Value>(&payload).ok(),
                "raw": payload
            });

            if sender
                .send(Message::Text(ws_msg.to_string().into()))
                .await
                .is_err()
            {
                break; // Client disconnected
            }
        }
    });

    // Spawn task to handle incoming WebSocket messages (client → NATS)
    let nats_clone = nats.clone();
    let mut recv_task = tokio::spawn(async move {
        while let Some(Ok(msg)) = receiver.next().await {
            match msg {
                Message::Text(text) => {
                    // Parse incoming message
                    if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(&text) {
                        if let Some(subject) = parsed.get("subject").and_then(|v| v.as_str()) {
                            if let Some(data) = parsed.get("data") {
                                // Publish to NATS
                                let payload = data.to_string();
                                if let Err(e) = nats_clone.publish(subject.to_string(), payload.into()).await {
                                    tracing::error!("Failed to publish to NATS: {}", e);
                                }
                            }
                        }
                    }
                }
                Message::Close(_) => {
                    tracing::info!("WebSocket closed by client (license: {})", license_id);
                    break;
                }
                // Ping/Pong is handled automatically by Axum
                _ => {}
            }
        }
    });

    // Wait for either task to complete (connection closed)
    tokio::select! {
        _ = (&mut send_task) => {
            recv_task.abort();
        }
        _ = (&mut recv_task) => {
            send_task.abort();
        }
    }

    tracing::info!("WebSocket connection closed for license {}", license_id);
}
