use hive_console_sdk::agent::usage_agent::RequestDetails;
use hive_router_plan_executor::headers::response::ResponseHeaderSink;
use http::Method;
use ntex::channel::oneshot;
use ntex::http::{header::HeaderName, header::HeaderValue, HeaderMap};
use ntex::router::Path;
use ntex::service::{fn_factory_with_config, fn_service, fn_shutdown, Service};
use ntex::web::{self, ws, Error, HttpRequest, HttpResponse};
use ntex::{chain, rt};
use sonic_rs::{JsonContainerTrait, JsonValueTrait, Value};
use std::cell::RefCell;
use std::collections::HashMap;
use std::io;
use std::rc::Rc;
use std::sync::Arc;
use std::time::Instant;
use tokio::sync::mpsc;
use tracing::{debug, error, trace, warn, Instrument};

use hive_router_internal::telemetry::metrics::catalog::values::SubscriptionEndReason;
use hive_router_internal::telemetry::metrics::subscription_metrics::{
    ActiveClientConnectionGuard, SubscriptionTransport,
};
use hive_router_internal::telemetry::traces::spans::graphql::GraphQLOperationSpan;
use hive_router_plan_executor::executors::graphql_transport_ws::{
    ClientMessage, CloseCode, ConnectionInitPayload, ServerMessage, WS_SUBPROTOCOL,
};
use hive_router_plan_executor::executors::websocket_common::{
    handshake_timeout, heartbeat, parse_frame_to_text, FrameNotParsedToText, WsState,
};
use hive_router_plan_executor::hooks::on_graphql_params::GraphQLParams;
use hive_router_plan_executor::plugin_context::{
    PluginContext, PluginRequestState, RouterHttpRequest,
};
use hive_router_plan_executor::request_context::{RequestContextExt, SharedRequestContext};
use hive_router_plan_executor::response::graphql_error::{GraphQLError, GraphQLErrorExtensions};
use hive_router_query_planner::state::supergraph_state::OperationKind;

use crate::jwt::errors::JwtError;
use crate::pipeline::active_subscriptions::SubscriptionEvent;
use crate::pipeline::error::PipelineError;
use crate::pipeline::execute_planned_request;
use crate::pipeline::header::{ResponseMode, SingleContentType, StreamContentType};
use crate::pipeline::{
    hash_graphql_extensions, hash_graphql_variables, inbound_request_fingerprint,
    normalize::normalize_request_with_cache, parser::parse_operation_with_cache, usage_reporting,
    validation::validate_operation_with_cache,
};
use crate::schema_state::SchemaState;
use crate::shared_state::{RouterSharedState, SharedRouterResponse};
use crate::telemetry::HeaderExtractor;
use hive_router_internal::telemetry::logging::request_id::WithRequestIdentifiers;
use hive_router_internal::telemetry::logging::scope::RequestLogScope;
use hive_router_internal::telemetry::logging::summary::{self, WithRequestSummary};
use hive_router_internal::telemetry::logging::targets;

type WsStateRef = Rc<RefCell<WsState<tokio::sync::mpsc::Sender<()>>>>;

pub async fn ws_index(
    req: HttpRequest,
    schema_state: web::types::State<Arc<SchemaState>>,
    shared_state: web::types::State<Arc<RouterSharedState>>,
) -> Result<HttpResponse, Error> {
    let schema_state = schema_state.get_ref().clone();
    let shared_state = shared_state.get_ref().clone();

    let accepted_subprotocol = ws::subprotocols(&req)
        .find(|p| *p == WS_SUBPROTOCOL)
        .map(|_| WS_SUBPROTOCOL);

    let plugin_context = req.extensions().get::<Arc<PluginContext>>().cloned();
    let request_context = match req.read_request_context() {
        Ok(ctx) => ctx,
        Err(_) => return Ok(HttpResponse::BadRequest().finish()),
    };

    ws::start(
        req.clone(),
        accepted_subprotocol,
        fn_factory_with_config(move |sink: ws::WsSink| {
            let schema_state = schema_state.clone();
            let shared_state = shared_state.clone();
            let plugin_context = plugin_context.clone();
            let request_context = request_context.clone();
            let req = req.clone();
            async move {
                ws_service(
                    accepted_subprotocol.is_some(),
                    sink,
                    schema_state,
                    shared_state,
                    req,
                    plugin_context,
                    request_context,
                )
                .await
            }
        }),
    )
    .await
}

#[allow(clippy::too_many_arguments)]
async fn ws_service(
    has_accepted_subprotocol: bool,
    sink: ws::WsSink,
    schema_state: Arc<SchemaState>,
    shared_state: Arc<RouterSharedState>,
    req: HttpRequest,
    plugin_context: Option<Arc<PluginContext>>,
    request_context: SharedRequestContext,
) -> Result<impl Service<ws::Frame, Response = Option<ws::Message>, Error = io::Error>, web::Error>
{
    if !has_accepted_subprotocol {
        warn!(target: targets::WEBSOCKET_SERVER, "websocket connection rejecting due to unacceptable subprotocol");
        let _ = sink.send(CloseCode::SubprotocolNotAcceptable.into()).await;
        // we dont return an Err here because we want to gracefully close the
        // connection for the client side with a close frame. returning an Err
        // would result in an abrupt termination of the connection
    } else {
        debug!(target: targets::WEBSOCKET_SERVER, "websocket connection accepted");
    }

    let conn_guard: Option<ActiveClientConnectionGuard> = if has_accepted_subprotocol {
        Some(
            shared_state
                .telemetry_context
                .metrics
                .subscriptions
                .active_client_connection(SubscriptionTransport::WebSocket),
        )
    } else {
        None
    };

    let ws_uri: Rc<http::Uri> = Rc::new(
        shared_state
            .router_config
            .websocket_path()
            .expect("websocket path must exist because the websocket handler wouldn't have been mounted otherwise")
            .parse()
            .unwrap_or_else(|_| http::Uri::from_static("/graphql")),
    );
    let ws_path: Rc<Path<http::Uri>> = Rc::new(Path::new((*ws_uri).clone()));

    let (heartbeat_tx, heartbeat_rx) = oneshot::channel();
    let (acknowledged_tx, acknowledged_rx) = oneshot::channel();

    let state: WsStateRef = Rc::new(RefCell::new(WsState::new(acknowledged_tx)));

    rt::spawn(heartbeat(state.clone(), sink.clone(), heartbeat_rx));
    rt::spawn(handshake_timeout(
        state.clone(),
        sink.clone(),
        acknowledged_rx,
        CloseCode::ConnectionInitTimeout,
    ));

    let state_for_service = state.clone();
    let service = fn_service(move |frame| {
        let sink = sink.clone();
        let state = state_for_service.clone();
        let schema_state = schema_state.clone();
        let shared_state = shared_state.clone();
        let plugin_context = plugin_context.clone();
        let request_context = request_context.clone();
        let ws_uri = ws_uri.clone();
        let ws_path = ws_path.clone();
        let req = req.clone();
        async move {
            match parse_frame_to_text(frame, &state) {
                Ok(text) => Ok(handle_text_frame(
                    text,
                    sink,
                    state,
                    &schema_state,
                    &shared_state,
                    &req,
                    plugin_context,
                    &request_context,
                    &ws_uri,
                    &ws_path,
                )
                .await),
                Err(FrameNotParsedToText::Message(msg)) => Ok(Some(msg)),
                Err(FrameNotParsedToText::Closed) => {
                    // we dont need to emit anything here because the conneciton is already closed
                    Ok(None)
                }
                Err(FrameNotParsedToText::None) => Ok(None),
            }
        }
    });

    let on_shutdown = fn_shutdown(async move || {
        // stop heartbeat and handshake timeout tasks on shutdown
        let _ = heartbeat_tx.send(());
        if let Some(tx) = state.borrow_mut().acknowledged_tx.take() {
            let _ = tx.send(());
        }
        // clearing the map will drop all the senders, which will
        // in turn cancel all active subscription streams and perform
        // the cleanup in there
        state.borrow_mut().subscriptions.clear();
        // drop conn_guard here to decrement the connection counter
        drop(conn_guard);
    });

    Ok(chain(service).and_then(on_shutdown))
}

/// Ensure a subscription is removed from active subscriptions when dropped (server-side).
struct SubscriptionGuard {
    state: WsStateRef,
    id: String,
}

impl Drop for SubscriptionGuard {
    fn drop(&mut self) {
        self.state.borrow_mut().subscriptions.remove(&self.id);
        trace!(target: targets::WEBSOCKET_SERVER, id = %self.id, "Subscription removed from active subscriptions");
    }
}

lazy_static::lazy_static! {
    static ref WS_URI_PATH: http::Uri = http::Uri::from_static("/graphql");
}

#[allow(clippy::too_many_arguments)]
async fn handle_text_frame(
    text: String,
    sink: ws::WsSink,
    state: WsStateRef,
    schema_state: &Arc<SchemaState>,
    shared_state: &Arc<RouterSharedState>,
    req: &HttpRequest,
    plugin_context: Option<Arc<PluginContext>>,
    request_context: &SharedRequestContext,
    ws_uri: &http::Uri,
    ws_path: &Path<http::Uri>,
) -> Option<ws::Message> {
    // IMPORTANT: we intentionally check if a supergraph is selected this "late" because
    // browsers cannot interpret a rejected HTTP Upgrade request and will show a cryptic
    // error to clients. instead, we accept the connection but close it immediately if no
    // supergraph is available, which results in a more understandable error
    let supergraph = match schema_state.select_supergraph(req) {
        Ok(supergraph) => supergraph,
        Err(err) => {
            error!(target: targets::WEBSOCKET_SERVER, err = ?err, "Supergraph runtime error");
            return Some(
                CloseCode::InternalServerError(Some("Supergraph runtime error".to_string())).into(),
            );
        }
    };
    let supergraph = supergraph.as_ref();

    // TODO: cover response header aggregation for WS
    let response_header_sink = ResponseHeaderSink::default();
    let client_msg: ClientMessage = match sonic_rs::from_str(&text) {
        Ok(msg) => msg,
        Err(e) => {
            error!(target: targets::WEBSOCKET_SERVER, error = ?e, "Failed to parse client message to JSON");
            return Some(CloseCode::BadRequest("Invalid message received from client").into());
        }
    };

    trace!(target: targets::WEBSOCKET_SERVER, type = client_msg.as_ref(), "Received client message");

    match client_msg {
        ClientMessage::ConnectionInit { payload } => {
            if state.borrow().handshake_received {
                warn!(target: targets::WEBSOCKET_SERVER, "Received multiple connection init messages, rejecting connection");

                return Some(CloseCode::TooManyInitialisationRequests.into());
            }
            state.borrow_mut().handshake_received = true;
            state.borrow_mut().init_payload = payload;
            state.borrow_mut().complete_handshake();

            let _ = sink.send(ServerMessage::ack()).await;

            debug!(target: targets::WEBSOCKET_SERVER, "websocket connection acknowledged");

            let header_map =
                parse_headers_from_connection_init_payload(state.borrow().init_payload.as_ref());
            if !header_map.is_empty() {
                trace!(target: targets::WEBSOCKET_SERVER, "Connection init message contains headers in the payload");
            } else {
                trace!(target: targets::WEBSOCKET_SERVER, "Connection init message does not contain headers in the payload");
            }

            None
        }
        ClientMessage::Subscribe { id, payload } => {
            if let Some(msg) = state.borrow().check_acknowledged() {
                return Some(msg);
            }

            if state.borrow().subscriptions.contains_key(&id) {
                return Some(CloseCode::SubscriberAlreadyExists(id).into());
            }

            let started_at = Instant::now();
            let operation_span = GraphQLOperationSpan::new();
            let span_clone = operation_span.clone();

            let result = async {
                let config = &shared_state.router_config.websocket;

                let connection_init_headers = if config.headers.accepts_connection_headers() {
                    let state_borrow = state.borrow();
                    parse_headers_from_connection_init_payload(state_borrow.init_payload.as_ref())
                } else {
                    HeaderMap::new()
                };

                let extensions_headers = if config.headers.accepts_operation_headers() {
                    parse_headers_from_extensions(payload.extensions.as_ref())
                } else {
                    HeaderMap::new()
                };

                // merge, extensions have precedence
                let mut headers = connection_init_headers;
                for (key, value) in extensions_headers.iter() {
                    headers.insert(key.clone(), value.clone());
                }

                let parent_ctx = shared_state
                    .telemetry_context
                    .extract_context(&HeaderExtractor(&headers));
                let request_identifiers =
                  Arc::new(shared_state
                        .telemetry_context
                        .logging_correlation_extractor
                        .extract(&headers, &parent_ctx));
                let headers = Arc::new(headers);

                // store the merged headers back to init_payload if configured to do so
                if config.headers.persist {
                    if let Some(ref mut init_payload) = state.borrow_mut().init_payload {
                        for (key, value) in headers.iter() {
                            if let Ok(val_str) = value.to_str() {
                                init_payload
                                    .fields
                                    .insert(key.to_string(), Value::from(val_str));
                            }
                        }
                    }
                }

                let inner_res = async {
                  let _summary = summary::SummaryOnDrop::new(started_at);
                  let Some(supergraph) = supergraph else {
                      // IMPORTANT: we dont do this earlier because router might be loading the
                      // supergraph as we speak, so we simply reject the operation request instead
                      error!(target: targets::WEBSOCKET_SERVER, "No supergraph available yet, unable to process client subscribe message");

                      return Some(ServerMessage::error(
                          &id,
                          &[GraphQLError::from_message_and_extensions(
                              "No supergraph available yet".to_string(),
                              GraphQLErrorExtensions::new_from_code("SERVICE_UNAVAILABLE"),
                          )],
                      ));
                  };

                  if supergraph.snapshot.is_retired() {
                      // the supergraph selected at upgrade time has since been retired (configured
                      // reload, or the owning plugin dropped/replaced its variant) - reject new
                      // operations on this connection rather than create work from a stale schema.
                      return Some(ServerMessage::error(
                          &id,
                          &[GraphQLError::from_message_and_extensions(
                              "Supergraph used by this connection has been retired".to_string(),
                              GraphQLErrorExtensions::new_from_code("SERVICE_UNAVAILABLE"),
                          )],
                      ));
                  }

                  let payload = GraphQLParams {
                      query: Some(payload.query),
                      operation_name: payload.operation_name,
                      variables: payload.variables.unwrap_or_default(),
                      extensions: payload.extensions,
                  };

                  // synthetic router http request for plugins - there's no real http request
                  // in the ws subscribe flow, so we assemble one from the ws path and merged headers.
                  // of course there is the http upgrade request, but that one is useless for the plugin system
                  let plugin_req_state = if let (Some(plugins), Some(ref plugin_context)) = (
                      shared_state.plugins.as_ref(),
                      plugin_context,
                  ) {
                      Some(PluginRequestState {
                          plugins: plugins.clone(),
                          router_http_request: RouterHttpRequest {
                              uri: ws_uri,
                              method: &Method::POST,
                              version: http::Version::HTTP_11,
                              headers: headers.as_ref(),
                              path: ws_uri.path(),
                              query_string: ws_uri.query().unwrap_or(""),
                              match_info: ws_path,
                          },
                          context: plugin_context.clone(),
                          request_context: request_context.clone(),
                      })
                  } else {
                      None
                  };

                  let client_name = headers
                      .get(
                          shared_state
                              .router_config
                              .telemetry
                              .client_identification
                              .name_header.get_header_ref(),
                      )
                      .and_then(|v| v.to_str().ok());
                  let client_version = headers
                      .get(
                          shared_state
                              .router_config
                              .telemetry
                              .client_identification
                              .version_header.get_header_ref(),
                      )
                      .and_then(|v| v.to_str().ok());

                  let parser_result =
                      match parse_operation_with_cache(shared_state, &payload, &plugin_req_state).await {
                          Ok(result) => result,
                          Err(err) => return Some(err.into_server_message(&id, shared_state)),
                      };

                  let parser_payload = match parser_result {
                      crate::pipeline::parser::ParseResult::Payload(payload) => payload,
                      crate::pipeline::parser::ParseResult::EarlyResponse(_) => {
                          return Some(ServerMessage::error(
                              &id,
                              &[GraphQLError::from_message_and_code(
                                  "Unexpected early response during parse",
                                  "INTERNAL_SERVER_ERROR",
                              )],
                          ));
                      }
                  };

                  operation_span.record_details(
                      &parser_payload.minified_document,
                      (&parser_payload).into(),
                      client_name,
                      client_version,
                      &parser_payload.hive_operation_hash,
                  );

                  match validate_operation_with_cache(
                      supergraph,
                      schema_state,
                      shared_state,
                      &parser_payload,
                      &plugin_req_state,
                  )
                  .await
                  {
                      Ok(Some(_)) => {
                          return Some(ServerMessage::error(
                              &id,
                              &[GraphQLError::from_message_and_code(
                                  "Unexpected early response during validation",
                                  "INTERNAL_SERVER_ERROR",
                              )],
                          ));
                      }
                      Ok(None) => {}
                      Err(err) => return Some(err.into_server_message(&id, shared_state)),
                  }

                  let normalize_payload = match normalize_request_with_cache(
                      &supergraph.snapshot,
                      &supergraph.runtime,
                      schema_state,
                      &payload,
                      &parser_payload,
                  )
                  .await
                  {
                      Ok(payload) => payload,
                      Err(err) => return Some(err.into_server_message(&id, shared_state)),
                  };

                  let is_subscription = matches!(
                      normalize_payload.operation_for_plan.operation_kind,
                      Some(OperationKind::Subscription)
                  );

                  if is_subscription && !shared_state.router_config.subscriptions.enabled {
                      return Some(PipelineError::SubscriptionsNotSupported.into_server_message(&id, shared_state));
                  }

                  let request_dedupe_enabled =
                      shared_state.router_config.traffic_shaping.router.dedupe.enabled;

                  let fingerprint = if request_dedupe_enabled
                      && matches!(
                          normalize_payload.operation_for_plan.operation_kind,
                          // same deduplication applies for queries and subscriptions
                          None | Some(OperationKind::Query) | Some(OperationKind::Subscription)
                      ) {
                      let variables_hash = hash_graphql_variables(&payload.variables);
                      let extensions_hash = payload
                          .extensions
                          .as_ref()
                          .map_or(0, hash_graphql_extensions);
                      Some(inbound_request_fingerprint(
                          &Method::POST,
                          ws_uri.path(),
                          headers.as_ref(),
                          &shared_state.in_flight_requests_header_policy,
                          supergraph.snapshot.cache_id,
                          normalize_payload.normalized_operation_hash,
                          variables_hash,
                          extensions_hash,
                      ))
                  } else {
                      None
                  };

                  // synthetic request details for plan executor
                  let response_mode = ResponseMode::Dual(
                      SingleContentType::default(),
                      StreamContentType::default(),
                  );
                  let method = Method::POST;

                  debug!(
                      target: targets::GRAPHQL_EXECUTION,
                      client_name,
                      client_version,
                      operation_name = parser_payload.operation_name.as_deref(),
                      operation = payload.query.as_deref(),
                      "graphql request started",
                  );

                  summary::record(|s| {
                      s.set_client_info(client_name, client_version);
                      s.set_operation_name(parser_payload.operation_name.as_deref());
                      s.set_operation_type(parser_payload.operation_type.as_str());
                      s.set_operation_hash(Some(parser_payload.hive_operation_hash.as_str()));
                      s.set_response_mode(response_mode.as_str());
                  });

                  let exec = |guard| execute_planned_request(
                      &method,
                      ws_uri,
                      headers.as_ref().clone(),
                      // TODO: WebSocket subscriptions do not yet expose route path params
                      Default::default(),
                      payload,
                      &normalize_payload,
                      supergraph,
                      shared_state,
                      schema_state,
                      operation_span,
                      plugin_req_state,
                      request_context,
                      &response_mode,
                      guard,
                      response_header_sink.clone()
                  );

                  let shared_response = if let Some(fp) = fingerprint {
                      let result = if is_subscription {
                          shared_state
                              .in_flight_requests
                              .claim(fp)
                              .get_or_try_init_with_guard(|guard| exec(Some(guard)))
                              .await
                      } else {
                          shared_state
                              .in_flight_requests
                              .claim(fp)
                              .get_or_try_init(|| exec(None))
                              .await
                      };
                      let (shared_response, _role) = match result {
                          Ok(result) => result,
                          Err(PipelineError::JwtError(err)) => {
                              let _ = sink.send(err.clone().into_server_message(&id, shared_state)).await;
                              // we report error as graphql error, but we also close the
                              // connection since we're dealing with auth so let's be safe
                              return Some(err.into_close_message());
                          },
                          Err(err) => return Some(err.into_server_message(&id, shared_state)),
                      };
                      Arc::unwrap_or_clone(shared_response)
                  } else {
                      match exec(None).await {
                          Ok(result) => result,
                          Err(PipelineError::JwtError(err)) => {
                              let _ = sink.send(err.clone().into_server_message(&id, shared_state)).await;
                              // we report error as graphql error, but we also close the
                              // connection since we're dealing with auth so let's be safe
                              return Some(err.into_close_message());
                          },
                          Err(err) => return Some(err.into_server_message(&id, shared_state)),
                      }
                  };

                  if let Some(hive_usage_agent) = &shared_state.hive_usage_agent {
                      let mut headers_vec = Vec::with_capacity(headers.len());
                      for (name, value) in headers.iter() {
                          if let Ok(val_str) = value.to_str() {
                              headers_vec.push((name.to_string(), val_str.to_string()));
                          }
                      }

                      let request_details = RequestDetails {
                          method: Method::POST,
                          url: (*ws_uri).clone(),
                          headers: headers_vec,
                      };
                      usage_reporting::collect_usage_report(
                          supergraph.snapshot.supergraph_schema.clone(),
                          started_at.elapsed(),
                          client_name,
                          client_version,
                          normalize_payload.operation_for_plan.name.as_deref(),
                          normalize_payload.operation_for_plan.operation_kind.as_ref(),
                          &parser_payload.minified_document,
                          hive_usage_agent,
                          shared_response.error_count(),
                          Some(request_details),
                          // The graphql-transport-ws Subscribe payload carries no document id
                          // and this path never invokes the document id resolver.
                          None,
                      )
                      .await;
                  }

                  debug!(target: targets::GRAPHQL_EXECUTION, error_count = shared_response.error_count(), "graphql request completed");

                  summary::record(|s| s.set_error_count(shared_response.error_count() as u32));

                  match shared_response {
                      SharedRouterResponse::Single(response) => {
                          let _ = sink.send(ServerMessage::next(&id, &response.body)).await;
                          Some(ServerMessage::complete(&id))
                      }
                      SharedRouterResponse::Stream(response) => {
                          let (cancel_tx, mut cancel_rx) = mpsc::channel::<()>(1);

                          state
                              .borrow_mut()
                              .subscriptions
                              .insert(id.clone(), cancel_tx);

                          let guard = SubscriptionGuard {
                              state: state.clone(),
                              id: id.clone(),
                          };

                          let mut receiver = response
                              .receiver
                              .unwrap_or_else(|| response.body.subscribe());

                          let client_op_guard = shared_state.telemetry_context.metrics.subscriptions.active_client_operation(SubscriptionTransport::WebSocket);
                          let metrics = shared_state.telemetry_context.metrics.clone();

                          trace!(target: targets::WEBSOCKET_SERVER, id = %id, "Subscription started");

                          let sink = sink.clone();
                          let id_for_loop = id.clone();
                          // must be spawned - blocking the frame handler would prevent
                          // ClientMessage::Complete from being received and processed,
                          // making cancellation impossible
                          let log_scope = RequestLogScope::capture();
                          rt::spawn(log_scope.scope(async move {
                              let _guard = guard;
                              let mut client_op_guard = client_op_guard;
                              let mut cancelled = false;

                              loop {
                                  tokio::select! {
                                      maybe_item = receiver.recv() => {
                                          match maybe_item {
                                              Ok(SubscriptionEvent::Raw(data)) => {
                                                  let _ = sink.send(ServerMessage::next(&id_for_loop, &data)).await;
                                                  metrics.subscriptions.record_client_sent(SubscriptionTransport::WebSocket);
                                              }
                                              Ok(SubscriptionEvent::Error(errors)) => {
                                                  client_op_guard.set_end_reason(SubscriptionEndReason::Error);
                                                  let _ = sink.send(ServerMessage::error(&id_for_loop, &errors)).await;
                                                  break;
                                              }
                                              Err(tokio::sync::broadcast::error::RecvError::Lagged(n)) => {
                                                  // NOTE: not warn to avoid log spam when receiver starts
                                                  // lagging. users should rely on the lagged_messages
                                                  // metric to detect slow consumers and tune accordingly
                                                  debug!(target: targets::WEBSOCKET_SERVER, id = %id_for_loop, lagged = n, "Broadcast receiver lagged, dropping message");
                                                  metrics.subscriptions.record_client_lag(SubscriptionTransport::WebSocket, n);
                                                  continue;
                                              }
                                              Err(tokio::sync::broadcast::error::RecvError::Closed) => {
                                                  client_op_guard.set_end_reason(SubscriptionEndReason::Completed);
                                                  break;
                                              }
                                          }
                                      }
                                      _ = cancel_rx.recv() => {
                                          cancelled = true;
                                          break;
                                      }
                                  }
                              }

                              if cancelled {
                                  trace!(target: targets::WEBSOCKET_SERVER, id = %id_for_loop, "Subscription cancelled");
                              } else {
                                  trace!(target: targets::WEBSOCKET_SERVER, id = %id_for_loop, "Subscription completed");
                                  let _ = sink.send(ServerMessage::complete(&id_for_loop)).await;
                              }
                          }));

                          None
                      }
                  }
                }.with_request_id(request_identifiers.clone()).with_request_summary().await;

                inner_res
            }
            .instrument(span_clone)
            .await;

            result
        }
        ClientMessage::Complete { id } => {
            if let Some(msg) = state.borrow().check_acknowledged() {
                return Some(msg);
            }

            if let Some(cancel_tx) = state.borrow_mut().subscriptions.remove(&id) {
                trace!(target: targets::WEBSOCKET_SERVER, id = %id, "Client requested subscription cancellation");
                let _ = cancel_tx.try_send(());
            }
            None
        }
        ClientMessage::Ping {} => {
            // respond with pong always, regardless of acknowledged state
            // the client should be able to use subprotocol pings/pongs to check liveness
            Some(ServerMessage::pong())
        }
        ClientMessage::Pong {} => None,
    }
}

/// Parses headers from a sonic_rs::Object into a HeaderMap. Only stringifiable
/// values are included; nulls, arrays, objects and other non-primitive values
/// are ignored.
fn parse_headers_from_object(headers_obj: &sonic_rs::Object) -> HeaderMap {
    let mut header_map = HeaderMap::new();

    for (key, value) in headers_obj.iter() {
        let Ok(name) = HeaderName::try_from(key as &str) else {
            continue;
        };
        let header_value = if let Some(s) = value.as_str() {
            HeaderValue::try_from(s).ok()
        } else if let Some(b) = value.as_bool() {
            Some(HeaderValue::from_static(if b { "true" } else { "false" }))
        } else if let Some(n) = value.as_i64() {
            Some(HeaderValue::from(n))
        } else if let Some(n) = value.as_u64() {
            Some(HeaderValue::from(n))
        } else if let Some(f) = value.as_f64() {
            // f64 has no HeaderValue::From impl, string alloc unavoidable
            HeaderValue::try_from(f.to_string()).ok()
        } else {
            None // ignore nulls, arrays, objects
        };
        if let Some(val) = header_value {
            header_map.insert(name, val);
        }
    }

    header_map
}

fn parse_headers_from_connection_init_payload(
    payload: Option<&ConnectionInitPayload>,
) -> HeaderMap {
    let mut header_map = HeaderMap::new();
    if let Some(payload) = payload {
        // First check if there's a nested "headers" object
        if let Some(headers_prop) = payload.fields.get("headers") {
            if let Some(headers_obj) = headers_prop.as_object() {
                header_map = parse_headers_from_object(headers_obj);
                return header_map;
            }
        }

        // If no nested "headers" object, treat all top-level fields as potential headers
        // Convert the entire fields HashMap to a sonic_rs::Object
        let mut obj = sonic_rs::Object::new();
        for (k, v) in payload.fields.iter() {
            obj.insert(k, v.clone());
        }
        header_map = parse_headers_from_object(&obj);
    }
    header_map
}

fn parse_headers_from_extensions(extensions: Option<&HashMap<String, Value>>) -> HeaderMap {
    let mut header_map = HeaderMap::new();
    if let Some(ext) = extensions {
        if let Some(headers_value) = ext.get("headers") {
            if let Some(headers_obj) = headers_value.as_object() {
                header_map = parse_headers_from_object(headers_obj);
            }
        }
    }
    header_map
}

// NOTE: no `From` trait because it can into ws message and ws closecode but both are ws::Message
impl PipelineError {
    fn into_server_message(self, id: &str, shared_state: &RouterSharedState) -> ws::Message {
        let code = self.graphql_error_code();
        let message = self.graphql_error_message();

        shared_state
            .telemetry_context
            .metrics
            .graphql
            .record_error(code);

        let graphql_error = GraphQLError::from_message_and_extensions(
            message,
            GraphQLErrorExtensions::new_from_code(code),
        );

        ServerMessage::error(id, &[graphql_error])
    }
}

// NOTE: no `From` trait because it can into ws message and ws closecode but both are ws::Message
impl JwtError {
    fn into_server_message(self, id: &str, shared_state: &RouterSharedState) -> ws::Message {
        let code = self.error_code();

        shared_state
            .telemetry_context
            .metrics
            .graphql
            .record_error(code);

        ServerMessage::error(
            id,
            &[GraphQLError::from_message_and_code(self.to_string(), code)],
        )
    }
    fn into_close_message(self) -> ws::Message {
        CloseCode::Forbidden(self.error_code().to_string()).into()
    }
}

#[cfg(test)]
mod tests {
    use sonic_rs::json;

    use super::*;

    #[test]
    fn should_parse_headers_from_object() {
        let headers_json = json!({
            "authorization": "Bearer token123",
            "x-custom-header": "custom-value",
            "x-number": 42,
            "x-bool": true,
            "x-float": 3.14
        });

        let headers_obj = headers_json.as_object().expect("Failed to get object");

        let headers = parse_headers_from_object(&headers_obj);

        assert_eq!(
            headers.get("authorization").unwrap().to_str().unwrap(),
            "Bearer token123"
        );
        assert_eq!(
            headers.get("x-custom-header").unwrap().to_str().unwrap(),
            "custom-value"
        );
        assert_eq!(headers.get("x-number").unwrap().to_str().unwrap(), "42");
        assert_eq!(headers.get("x-bool").unwrap().to_str().unwrap(), "true");
        assert_eq!(headers.get("x-float").unwrap().to_str().unwrap(), "3.14");
    }

    // TODO: hella tests
}
