use std::collections::HashMap;
use std::sync::LazyLock;

use schemars::JsonSchema;
use segment::common::anonymize::Anonymize;
use semver::Version;
use serde::{Deserialize, Serialize};

use crate::shards::shard::PeerId;

/// Service version, starting from which `ManualRecovery` state is supported.
pub static MANUAL_RECOVERY_SHARD_STATE_VERSION: LazyLock<Version> =
    LazyLock::new(|| Version::parse("1.16.4-dev").expect("valid version string"));

/// Represents a replica set state
#[derive(Debug, Deserialize, Serialize, Default, PartialEq, Eq, Clone)]
pub struct ReplicaSetState {
    pub is_local: bool,
    pub this_peer_id: PeerId,
    peers: HashMap<PeerId, ReplicaState>,
}

impl ReplicaSetState {
    pub fn get_peer_state(&self, peer_id: PeerId) -> Option<ReplicaState> {
        self.peers.get(&peer_id).copied()
    }

    /// Returns previous state if any
    pub fn set_peer_state(&mut self, peer_id: PeerId, state: ReplicaState) -> Option<ReplicaState> {
        self.peers.insert(peer_id, state)
    }

    pub fn remove_peer_state(&mut self, peer_id: PeerId) -> Option<ReplicaState> {
        self.peers.remove(&peer_id)
    }

    pub fn peers(&self) -> &HashMap<PeerId, ReplicaState> {
        &self.peers
    }

    pub fn check_peers_state_all<F>(&self, check: F) -> bool
    where
        F: Fn(ReplicaState) -> bool,
    {
        self.peers.values().all(|state| check(*state))
    }

    pub fn active_peers(&self) -> Vec<PeerId> {
        self.peers
            .iter()
            .filter_map(|(peer_id, state)| {
                // We consider `ReshardingScaleDown` to be `Active`!
                state.is_active().then_some(*peer_id)
            })
            .collect()
    }

    pub fn readable_peers(&self) -> Vec<PeerId> {
        self.peers
            .iter()
            .filter_map(|(peer_id, state)| state.is_readable().then_some(*peer_id))
            .collect()
    }

    pub fn active_or_resharding_peers(&self) -> impl Iterator<Item = PeerId> + '_ {
        self.peers.iter().filter_map(|(peer_id, state)| {
            matches!(
                state,
                ReplicaState::Active | ReplicaState::Resharding | ReplicaState::ReshardingScaleDown
            )
            .then_some(*peer_id)
        })
    }

    pub fn set_peers(&mut self, peers: HashMap<PeerId, ReplicaState>) {
        self.peers = peers;
    }

    /// Change current `this_peer_id` to a new one.
    pub fn switch_peer_id(&mut self, new_peer_id: PeerId) {
        let old_peer_id = self.this_peer_id;
        self.this_peer_id = new_peer_id;

        self.peers
            .remove(&old_peer_id)
            .and_then(|replica_state| self.peers.insert(new_peer_id, replica_state));
    }

    /// Remove all remote peers from the replica set state and activate local peer.
    pub fn force_local_active(&mut self) {
        self.peers.clear();
        self.peers.insert(self.this_peer_id, ReplicaState::Active);
    }
}

/// State of the single shard within a replica set.
#[derive(
    Debug, Deserialize, Serialize, JsonSchema, Default, PartialEq, Eq, Hash, Clone, Copy, Anonymize,
)]
pub enum ReplicaState {
    // Active and sound
    #[default]
    Active,
    // Failed for some reason
    Dead,
    // The shard is partially loaded and is currently receiving data from other shards
    Partial,
    // Collection is being created
    Initializing,
    // A shard which receives data, but is not used for search
    // Useful for backup shards
    Listener,
    // Deprecated since Qdrant 1.9.0, used in Qdrant 1.7.0 and 1.8.0
    //
    // Snapshot shard transfer is in progress, updates aren't sent to the shard
    // Normally rejects updates. Since 1.8 it allows updates if force is true.
    PartialSnapshot,
    // Shard is undergoing recovery by an external node
    // Normally rejects updates, accepts updates if force is true
    Recovery,
    // Points are being migrated to this shard as part of resharding up
    Resharding,
    // Points are being migrated to this shard as part of resharding down
    ReshardingScaleDown,
    // Active for readers, Partial for writers
    ActiveRead,
    // State for manually creation/recovery of a shard.
    // Usually when snapshot is uploaded.
    // This state is equivalent to `Partial`, except:
    // - it can't receive updates
    // - it is not treated as broken on startup
    ManualRecovery,
}

impl ReplicaState {
    /// Check if replica state is active
    /// Used to define if this replica can be used as a source of truth.
    pub fn is_active(self) -> bool {
        match self {
            ReplicaState::Active => true,
            ReplicaState::ReshardingScaleDown => true,

            ReplicaState::Dead
            | ReplicaState::Partial
            | ReplicaState::ManualRecovery
            | ReplicaState::Initializing
            | ReplicaState::Listener
            | ReplicaState::PartialSnapshot
            | ReplicaState::Recovery
            | ReplicaState::Resharding
            | ReplicaState::ActiveRead => false,
        }
    }

    /// Check that replica has full dataset, so it can be used for read operations.
    pub fn is_readable(self) -> bool {
        match self {
            ReplicaState::Active => true,
            ReplicaState::ReshardingScaleDown => true,
            ReplicaState::ActiveRead => true,
            // False from here on
            ReplicaState::Dead => false,
            ReplicaState::Partial => false,
            ReplicaState::ManualRecovery => false,
            ReplicaState::Initializing => false,
            ReplicaState::Listener => false,
            ReplicaState::PartialSnapshot => false,
            ReplicaState::Recovery => false,
            ReplicaState::Resharding => false,
        }
    }

    pub fn is_updatable(self) -> bool {
        match self {
            ReplicaState::Active => true,
            ReplicaState::Partial => true,
            ReplicaState::Initializing => true,
            ReplicaState::Listener => true,
            ReplicaState::Recovery | ReplicaState::PartialSnapshot => false,
            ReplicaState::Resharding | ReplicaState::ReshardingScaleDown => true,
            ReplicaState::Dead => false,
            ReplicaState::ActiveRead => true,
            ReplicaState::ManualRecovery => false,
        }
    }

    /// Check if this peer can be used as a source of truth within a shard_id.
    /// For instance:
    /// - It can be the only receiver of updates
    /// - It can be a primary replica for ordered writes
    pub fn can_be_source_of_truth(self) -> bool {
        match self {
            ReplicaState::Active => true,
            ReplicaState::ActiveRead => true, // Can be only one replica per shard_id
            ReplicaState::Resharding => true, // Can be only one replica per shard_id
            ReplicaState::ReshardingScaleDown => true, // Acts like Active, until resharding is committed
            // false from here on
            ReplicaState::Partial => false,
            ReplicaState::ManualRecovery => false,
            ReplicaState::Initializing => false,
            ReplicaState::Listener => false,
            ReplicaState::PartialSnapshot => false,
            ReplicaState::Recovery => false,
            ReplicaState::Dead => false,
        }
    }

    /// Check whether the replica state is active or listener or resharding.
    /// Healthy state means that replica does not require **automatic** recovery.
    pub fn is_healthy(self) -> bool {
        match self {
            ReplicaState::Active
            | ReplicaState::Listener
            | ReplicaState::Resharding
            | ReplicaState::ManualRecovery
            | ReplicaState::ReshardingScaleDown => true,

            ReplicaState::Dead
            | ReplicaState::Initializing
            | ReplicaState::Partial
            | ReplicaState::PartialSnapshot
            | ReplicaState::Recovery
            | ReplicaState::ActiveRead => false,
        }
    }

    /// Check if the replica state requires automatic recovery to be scheduled.
    pub fn requires_recovery(self) -> bool {
        match self {
            ReplicaState::Dead => true,
            ReplicaState::Active
            | ReplicaState::Partial
            | ReplicaState::Initializing
            | ReplicaState::Listener
            | ReplicaState::PartialSnapshot
            | ReplicaState::Recovery
            | ReplicaState::ManualRecovery
            | ReplicaState::Resharding
            | ReplicaState::ReshardingScaleDown
            | ReplicaState::ActiveRead => false,
        }
    }

    /// Check whether the replica state is partial or partial-like.
    ///
    /// In other words: is the state related to shard transfers?
    //
    // TODO(resharding): What's the best way to handle `ReshardingScaleDown` properly!?
    pub fn is_partial_or_recovery(self) -> bool {
        match self {
            ReplicaState::Partial
            | ReplicaState::ManualRecovery
            | ReplicaState::PartialSnapshot
            | ReplicaState::Recovery
            | ReplicaState::Resharding
            | ReplicaState::ReshardingScaleDown
            | ReplicaState::ActiveRead => true,

            ReplicaState::Active
            | ReplicaState::Dead
            | ReplicaState::Initializing
            | ReplicaState::Listener => false,
        }
    }

    /// Returns `true` if the replica state is resharding, either up or down.
    pub fn is_resharding(&self) -> bool {
        match self {
            ReplicaState::Resharding | ReplicaState::ReshardingScaleDown => true,

            ReplicaState::Partial
            | ReplicaState::ManualRecovery
            | ReplicaState::PartialSnapshot
            | ReplicaState::Recovery
            | ReplicaState::Active
            | ReplicaState::Dead
            | ReplicaState::Initializing
            | ReplicaState::Listener
            | ReplicaState::ActiveRead => false,
        }
    }
}
