/* SPDX-License-Identifier: GPL-2.0 OR BSD-3-Clause */
/*
 * Copyright (C) 2024-2025 Intel Corporation
 */

#ifndef __iwl_mld_sta_h__
#define __iwl_mld_sta_h__

#include <net/mac80211.h>

#include "mld.h"
#include "tx.h"

/**
 * struct iwl_mld_rxq_dup_data - Duplication detection data, per STA & Rx queue
 * @last_seq: last sequence per tid.
 * @last_sub_frame_idx: the index of the last subframe in an A-MSDU. This value
 *	will be zero if the packet is not part of an A-MSDU.
 */
struct iwl_mld_rxq_dup_data {
	__le16 last_seq[IWL_MAX_TID_COUNT + 1];
	u8 last_sub_frame_idx[IWL_MAX_TID_COUNT + 1];
} ____cacheline_aligned_in_smp;

/**
 * struct iwl_mld_link_sta - link-level station
 *
 * This represents the link-level sta - the driver level equivalent to the
 * ieee80211_link_sta
 *
 * @last_rate_n_flags: rate_n_flags from the last &iwl_tlc_update_notif
 * @signal_avg: the signal average coming from the firmware
 * @in_fw: whether the link STA is uploaded to the FW (false during restart)
 * @rcu_head: RCU head for freeing this object
 * @fw_id: the FW id of this link sta.
 */
struct iwl_mld_link_sta {
	/* Add here fields that need clean up on restart */
	struct_group(zeroed_on_hw_restart,
		u32 last_rate_n_flags;
		bool in_fw;
		s8 signal_avg;
	);
	/* And here fields that survive a fw restart */
	struct rcu_head rcu_head;
	u32 fw_id;
};

#define iwl_mld_link_sta_dereference_check(mld_sta, link_id)		\
	rcu_dereference_check((mld_sta)->link[link_id],			\
			      lockdep_is_held(&mld_sta->mld->wiphy->mtx))

#define for_each_mld_link_sta(mld_sta, link_sta, link_id)		\
	for (link_id = 0; link_id < ARRAY_SIZE((mld_sta)->link);	\
	     link_id++)							\
		if ((link_sta =						\
			iwl_mld_link_sta_dereference_check(mld_sta, link_id)))

#define IWL_NUM_DEFAULT_KEYS 4

/* struct iwl_mld_ptk_pn - Holds Packet Number (PN) per TID.
 * @rcu_head: RCU head for freeing this data.
 * @pn: Array storing PN for each TID.
 */
struct iwl_mld_ptk_pn {
	struct rcu_head rcu_head;
	struct {
		u8 pn[IWL_MAX_TID_COUNT][IEEE80211_CCMP_PN_LEN];
	} ____cacheline_aligned_in_smp q[];
};

/**
 * struct iwl_mld_per_link_mpdu_counter - per-link TX/RX MPDU counters
 *
 * @tx: Number of TX MPDUs.
 * @rx: Number of RX MPDUs.
 */
struct iwl_mld_per_link_mpdu_counter {
	u32 tx;
	u32 rx;
};

/**
 * struct iwl_mld_per_q_mpdu_counter - per-queue MPDU counter
 *
 * @lock: Needed to protect the counters when modified from statistics.
 * @per_link: per-link counters.
 * @window_start_time: timestamp of the counting-window start
 */
struct iwl_mld_per_q_mpdu_counter {
	spinlock_t lock;
	struct iwl_mld_per_link_mpdu_counter per_link[IWL_FW_MAX_LINK_ID + 1];
	unsigned long window_start_time;
} ____cacheline_aligned_in_smp;

/**
 * struct iwl_mld_sta - representation of a station in the driver.
 *
 * This represent the MLD-level sta, and will not be added to the FW.
 * Embedded in ieee80211_sta.
 *
 * @vif: pointer the vif object.
 * @sta_state: station state according to enum %ieee80211_sta_state
 * @sta_type: type of this station. See &enum iwl_fw_sta_type
 * @mld: a pointer to the iwl_mld object
 * @dup_data: per queue duplicate packet detection data
 * @data_tx_ant: stores the last TX antenna index; used for setting
 *	TX rate_n_flags for injected data frames (toggles on every TX failure).
 * @tid_to_baid: a simple map of TID to Block-Ack fw id
 * @deflink: This holds the default link STA information, for non MLO STA all
 *	link specific STA information is accessed through @deflink or through
 *	link[0] which points to address of @deflink. For MLO Link STA
 *	the first added link STA will point to deflink.
 * @link: reference to Link Sta entries. For Non MLO STA, except 1st link,
 *	i.e link[0] all links would be assigned to NULL by default and
 *	would access link information via @deflink or link[0]. For MLO
 *	STA, first link STA being added will point its link pointer to
 *	@deflink address and remaining would be allocated and the address
 *	would be assigned to link[link_id] where link_id is the id assigned
 *	by the AP.
 * @ptk_pn: Array of pointers to PTK PN data, used to track the Packet Number
 *	per key index and per queue (TID).
 * @mpdu_counters: RX/TX MPDUs counters for each queue.
 */
struct iwl_mld_sta {
	/* Add here fields that need clean up on restart */
	struct_group(zeroed_on_hw_restart,
		enum ieee80211_sta_state sta_state;
		enum iwl_fw_sta_type sta_type;
	);
	/* And here fields that survive a fw restart */
	struct iwl_mld *mld;
	struct ieee80211_vif *vif;
	struct iwl_mld_rxq_dup_data *dup_data;
	u8 tid_to_baid[IWL_MAX_TID_COUNT];
	u8 data_tx_ant;

	struct iwl_mld_link_sta deflink;
	struct iwl_mld_link_sta __rcu *link[IEEE80211_MLD_MAX_NUM_LINKS];
	struct iwl_mld_ptk_pn __rcu *ptk_pn[IWL_NUM_DEFAULT_KEYS];
	struct iwl_mld_per_q_mpdu_counter *mpdu_counters;
};

static inline struct iwl_mld_sta *
iwl_mld_sta_from_mac80211(struct ieee80211_sta *sta)
{
	return (void *)sta->drv_priv;
}

static inline void
iwl_mld_cleanup_sta(void *data, struct ieee80211_sta *sta)
{
	struct iwl_mld_sta *mld_sta = iwl_mld_sta_from_mac80211(sta);
	struct iwl_mld_link_sta *mld_link_sta;
	u8 link_id;

	for (int i = 0; i < ARRAY_SIZE(sta->txq); i++)
		CLEANUP_STRUCT(iwl_mld_txq_from_mac80211(sta->txq[i]));

	for_each_mld_link_sta(mld_sta, mld_link_sta, link_id) {
		CLEANUP_STRUCT(mld_link_sta);

		if (!ieee80211_vif_is_mld(mld_sta->vif)) {
			/* not an MLD STA; only has the deflink with ID zero */
			WARN_ON(link_id);
			continue;
		}

		if (mld_sta->vif->active_links & BIT(link_id))
			continue;

		/* Should not happen as link removal should always succeed */
		WARN_ON(1);
		RCU_INIT_POINTER(mld_sta->link[link_id], NULL);
		RCU_INIT_POINTER(mld_sta->mld->fw_id_to_link_sta[mld_link_sta->fw_id],
				 NULL);
		if (mld_link_sta != &mld_sta->deflink)
			kfree_rcu(mld_link_sta, rcu_head);
	}

	CLEANUP_STRUCT(mld_sta);
}

static inline struct iwl_mld_link_sta *
iwl_mld_link_sta_from_mac80211(struct ieee80211_link_sta *link_sta)
{
	struct iwl_mld_sta *mld_sta = iwl_mld_sta_from_mac80211(link_sta->sta);

	return iwl_mld_link_sta_dereference_check(mld_sta, link_sta->link_id);
}

int iwl_mld_add_sta(struct iwl_mld *mld, struct ieee80211_sta *sta,
		    struct ieee80211_vif *vif, enum iwl_fw_sta_type type);
void iwl_mld_remove_sta(struct iwl_mld *mld, struct ieee80211_sta *sta);
int iwl_mld_fw_sta_id_from_link_sta(struct iwl_mld *mld,
				    struct ieee80211_link_sta *link_sta);
u32 iwl_mld_fw_sta_id_mask(struct iwl_mld *mld, struct ieee80211_sta *sta);
int iwl_mld_update_all_link_stations(struct iwl_mld *mld,
				     struct ieee80211_sta *sta);
void iwl_mld_flush_sta_txqs(struct iwl_mld *mld, struct ieee80211_sta *sta);
void iwl_mld_wait_sta_txqs_empty(struct iwl_mld *mld,
				 struct ieee80211_sta *sta);
void iwl_mld_count_mpdu_rx(struct ieee80211_link_sta *link_sta, int queue,
			   u32 count);
void iwl_mld_count_mpdu_tx(struct ieee80211_link_sta *link_sta, u32 count);

/**
 * struct iwl_mld_int_sta - representation of an internal station
 * (a station that exist in FW and in driver, but not in mac80211)
 *
 * @sta_id: the index of the station in the fw
 * @queue_id: the if of the queue used by the station
 * @sta_type: station type. One of &iwl_fw_sta_type
 */
struct iwl_mld_int_sta {
	u8 sta_id;
	u32 queue_id;
	enum iwl_fw_sta_type sta_type;
};

static inline void
iwl_mld_init_internal_sta(struct iwl_mld_int_sta *internal_sta)
{
	internal_sta->sta_id = IWL_INVALID_STA;
	internal_sta->queue_id = IWL_MLD_INVALID_QUEUE;
}

static inline void
iwl_mld_free_internal_sta(struct iwl_mld *mld,
			  struct iwl_mld_int_sta *internal_sta)
{
	if (WARN_ON(internal_sta->sta_id == IWL_INVALID_STA))
		return;

	RCU_INIT_POINTER(mld->fw_id_to_link_sta[internal_sta->sta_id], NULL);
	iwl_mld_init_internal_sta(internal_sta);
}

int iwl_mld_add_bcast_sta(struct iwl_mld *mld,
			  struct ieee80211_vif *vif,
			  struct ieee80211_bss_conf *link);

int iwl_mld_add_mcast_sta(struct iwl_mld *mld,
			  struct ieee80211_vif *vif,
			  struct ieee80211_bss_conf *link);

int iwl_mld_add_aux_sta(struct iwl_mld *mld,
			struct iwl_mld_int_sta *internal_sta);

void iwl_mld_remove_bcast_sta(struct iwl_mld *mld,
			      struct ieee80211_vif *vif,
			      struct ieee80211_bss_conf *link);

void iwl_mld_remove_mcast_sta(struct iwl_mld *mld,
			      struct ieee80211_vif *vif,
			      struct ieee80211_bss_conf *link);

void iwl_mld_remove_aux_sta(struct iwl_mld *mld,
			    struct ieee80211_vif *vif,
			    struct ieee80211_bss_conf *link);

int iwl_mld_update_link_stas(struct iwl_mld *mld,
			     struct ieee80211_vif *vif,
			     struct ieee80211_sta *sta,
			     u16 old_links, u16 new_links);
#endif /* __iwl_mld_sta_h__ */