use std::fmt;

use schemars::JsonSchema;
use segment::types::ShardKey;
use serde::{Deserialize, Serialize};
use uuid::Uuid;

use super::shard::{PeerId, ShardId};
use crate::operations::cluster_ops::ReshardingDirection;

#[derive(Clone, Debug, Eq, PartialEq, Deserialize, Serialize)]
pub struct ReshardState {
    pub uuid: Uuid,
    pub peer_id: PeerId,
    pub shard_id: ShardId,
    pub shard_key: Option<ShardKey>,
    pub direction: ReshardingDirection,
    pub stage: ReshardingStage,
}

impl ReshardState {
    pub fn new(
        uuid: Uuid,
        direction: ReshardingDirection,
        peer_id: PeerId,
        shard_id: ShardId,
        shard_key: Option<ShardKey>,
    ) -> Self {
        Self {
            uuid,
            direction,
            peer_id,
            shard_id,
            shard_key,
            stage: ReshardingStage::MigratingPoints,
        }
    }

    pub fn matches(&self, key: &ReshardKey) -> bool {
        let ReshardKey {
            uuid,
            direction,
            peer_id,
            shard_id,
            shard_key,
        } = key;

        self.uuid == *uuid
            && self.direction == *direction
            && self.peer_id == *peer_id
            && self.shard_id == *shard_id
            && self.shard_key == *shard_key
    }

    pub fn key(&self) -> ReshardKey {
        let ReshardState {
            uuid,
            peer_id,
            shard_id,
            shard_key,
            direction,
            stage: _,
        } = self;

        ReshardKey {
            uuid: *uuid,
            direction: *direction,
            peer_id: *peer_id,
            shard_id: *shard_id,
            shard_key: shard_key.clone(),
        }
    }
}

/// Resharding stages
///
/// # Warning
///
/// This enum is ordered!
#[derive(Copy, Clone, Debug, Default, Eq, PartialEq, Ord, PartialOrd, Deserialize, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum ReshardingStage {
    #[default]
    MigratingPoints,
    ReadHashRingCommitted,
    WriteHashRingCommitted,
}

/// Unique identifier of a resharding task
#[derive(Clone, Debug, Eq, PartialEq, Hash, Deserialize, Serialize, JsonSchema)]
pub struct ReshardKey {
    #[schemars(skip)]
    pub uuid: Uuid,
    #[serde(default)]
    pub direction: ReshardingDirection,
    pub peer_id: PeerId,
    pub shard_id: ShardId,
    pub shard_key: Option<ShardKey>,
}

impl ReshardKey {
    /// Pin the auto-sharding invariant that resharding always targets the
    /// last shard id. `shard_number` must equal `shard_id` or `shard_id + 1`;
    /// any other value would silently corrupt the count via the id-derived
    /// `shard_number` update.
    pub(crate) fn debug_assert_targets_last_shard(&self, shard_number: u32) {
        let shard_id = self.shard_id;
        debug_assert!(
            shard_id
                .checked_add(1)
                .is_some_and(|next| shard_number == shard_id || shard_number == next),
            "auto-sharding resharding must target the last shard id; \
             shard_number={shard_number} shard_id={shard_id}",
        );
    }
}

impl fmt::Display for ReshardKey {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}/{}/{:?}", self.peer_id, self.shard_id, self.shard_key)
    }
}
