naia-server 0.25.0

use std::{
    collections::{HashMap, HashSet},
    hash::Hash,
};

use naia_shared::{
    AuthorityError, EntityAuthAccessor, EntityAuthStatus, GlobalEntity, HostAuthHandler, HostType,
};

use crate::UserKey;

#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)]
pub enum AuthOwner {
    None,
    Server,
    Client(UserKey),
}

impl AuthOwner {
    pub fn from_user_key(user_key: Option<&UserKey>) -> Self {
        match user_key {
            Some(user_key) => AuthOwner::Client(*user_key),
            None => AuthOwner::Server,
        }
    }
}

pub struct ServerAuthHandler {
    host_auth_handler: HostAuthHandler,
    entity_auth_map: HashMap<GlobalEntity, AuthOwner>,
    user_to_entity_map: HashMap<UserKey, HashSet<GlobalEntity>>,
}

impl ServerAuthHandler {
    pub fn new() -> Self {
        Self {
            host_auth_handler: HostAuthHandler::new(),
            entity_auth_map: HashMap::new(),
            user_to_entity_map: HashMap::new(),
        }
    }

    pub fn get_accessor(&self, entity: &GlobalEntity) -> EntityAuthAccessor {
        self.host_auth_handler.get_accessor(entity)
    }

    pub fn register_entity(&mut self, entity: &GlobalEntity) {
        self.host_auth_handler
            .register_entity(HostType::Server, entity);
        self.entity_auth_map.insert(*entity, AuthOwner::None);
    }

    pub fn deregister_entity(&mut self, entity: &GlobalEntity) {
        self.host_auth_handler.deregister_entity(entity);
        self.entity_auth_map.remove(entity);
    }

    pub(crate) fn authority_status(&self, entity: &GlobalEntity) -> Option<EntityAuthStatus> {
        self.host_auth_handler
            .auth_status(entity)
            .map(|host_status| host_status.status())
    }

    /// True iff some user (or the server) currently holds authority on `entity`.
    /// Returns false for entities with no current holder (Available) and
    /// for non-delegated entities (which have no auth_owner record at all).
    pub(crate) fn entity_has_holder(&self, entity: &GlobalEntity) -> bool {
        match self.entity_auth_map.get(entity) {
            Some(AuthOwner::None) | None => false,
            Some(AuthOwner::Server) | Some(AuthOwner::Client(_)) => true,
        }
    }

    pub(crate) fn client_request_authority(
        &mut self,
        entity: &GlobalEntity,
        requester: &AuthOwner,
    ) -> Result<(), AuthorityError> {
        let Some(owner) = self.entity_auth_map.get_mut(entity) else {
            return Err(AuthorityError::NotDelegated);
        };
        if *owner == AuthOwner::None {
            match requester {
                AuthOwner::Server => {
                    *owner = AuthOwner::Server;
                    // If the Server is requesting Authority, grant the Server local Authority
                    self.host_auth_handler
                        .set_auth_status(entity, EntityAuthStatus::Granted);
                }
                AuthOwner::Client(user_key) => {
                    *owner = AuthOwner::Client(*user_key);
                    self.user_to_entity_map
                        .entry(*user_key)
                        .or_default()
                        .insert(*entity);
                    // If a Client is requesting Authority, restrict the Server's local Authority
                    self.host_auth_handler
                        .set_auth_status(entity, EntityAuthStatus::Denied);
                }
                AuthOwner::None => {}
            }

            Ok(())
        } else {
            Err(AuthorityError::NotAvailable)
        }
    }

    pub(crate) fn client_release_authority(
        &mut self,
        entity: &GlobalEntity,
        releaser: &AuthOwner,
    ) -> Result<(), AuthorityError> {
        let Some(owner) = self.entity_auth_map.get_mut(entity) else {
            return Err(AuthorityError::NotDelegated);
        };

        // Server-initiated release is sovereign: it can clear ANY current
        // holder back to Available. This implements the contract
        // [entity-authority-09]: "server release_authority clears holder,
        // all clients observe Available". Clients can only release their
        // own hold (the original `owner == releaser` check).
        let server_initiated = matches!(releaser, AuthOwner::Server);
        if owner == releaser || server_initiated {
            let previous_owner = *owner;
            *owner = AuthOwner::None;
            self.release_all_authority(entity, previous_owner);

            Ok(())
        } else {
            Err(AuthorityError::NotHolder)
        }
    }

    /// Server-priority give_authority: assign `target_user` as the
    /// authority holder regardless of who currently holds it.
    /// Implements contract [entity-authority-10] ("server priority:
    /// give_authority overrides current holder"). Unlike
    /// `client_request_authority` (which fails with `NotAvailable`
    /// when the slot is held), this is sovereign and never fails on
    /// the holder check.
    ///
    /// Returns the previous owner so callers can fan out the
    /// appropriate auth-status changes (e.g. emit `AuthLost` to the
    /// previous holder).
    pub(crate) fn server_give_authority_to_client(
        &mut self,
        entity: &GlobalEntity,
        target_user: &UserKey,
    ) -> Result<AuthOwner, AuthorityError> {
        let Some(owner) = self.entity_auth_map.get_mut(entity) else {
            return Err(AuthorityError::NotDelegated);
        };
        let previous_owner = *owner;
        // Idempotent: if the target user already holds, return without
        // touching state. Re-issuing SetAuthority would otherwise drive
        // an illegal Granted→Granted transition through the auth channel.
        if previous_owner == AuthOwner::Client(*target_user) {
            return Ok(previous_owner);
        }
        *owner = AuthOwner::Client(*target_user);
        // Clear bookkeeping for the previous holder. release_all_authority
        // sets host_auth_status to Available as a side effect; the per-user
        // SetAuthority messages in the caller will overwrite that with
        // Granted (for target_user) and Denied (for everyone else).
        self.release_all_authority(entity, previous_owner);
        // Track the new client holder.
        self.user_to_entity_map
            .entry(*target_user)
            .or_default()
            .insert(*entity);
        // Restrict the server's local Authority — a client now holds it.
        self.host_auth_handler
            .set_auth_status(entity, EntityAuthStatus::Denied);
        Ok(previous_owner)
    }

    pub(crate) fn server_take_authority(
        &mut self,
        entity: &GlobalEntity,
    ) -> Result<AuthOwner, AuthorityError> {
        let Some(owner) = self.entity_auth_map.get_mut(entity) else {
            return Err(AuthorityError::NotDelegated);
        };

        let previous_owner = *owner;
        if previous_owner == AuthOwner::Server {
            return Ok(previous_owner);
        }
        *owner = AuthOwner::Server;
        // Clean up previous holder (sets Available as side effect; overridden below)
        self.release_all_authority(entity, previous_owner);
        // Server holds authority — transition to Granted
        self.host_auth_handler
            .set_auth_status(entity, EntityAuthStatus::Granted);

        Ok(previous_owner)
    }

    fn release_all_authority(&mut self, entity: &GlobalEntity, owner: AuthOwner) -> bool {
        if owner == AuthOwner::None {
            // no change was made
            return false;
        }

        if let AuthOwner::Client(user_key) = owner {
            let mut remove_user = false;
            if let Some(entities) = self.user_to_entity_map.get_mut(&user_key) {
                entities.remove(entity);
                remove_user = true;
            }
            if remove_user {
                self.user_to_entity_map.remove(&user_key);
            }
        }

        self.host_auth_handler
            .set_auth_status(entity, EntityAuthStatus::Available);

        true
    }

    pub(crate) fn user_all_owned_entities(
        &self,
        user_key: &UserKey,
    ) -> Option<&HashSet<GlobalEntity>> {
        if let Some(entities) = self.user_to_entity_map.get(user_key) {
            return Some(entities);
        }
        None
    }

    /// Check if a user is the authority holder for a specific entity
    pub(crate) fn user_is_authority_holder(
        &self,
        user_key: &UserKey,
        entity: &GlobalEntity,
    ) -> bool {
        self.entity_auth_map
            .get(entity)
            .map(|owner| *owner == AuthOwner::Client(*user_key))
            .unwrap_or(false)
    }
}