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

use super::{ShardTransfer, ShardTransferKey, ShardTransferMethod};
use crate::operations::types::{CollectionError, CollectionResult};
use crate::shards::replica_set::replica_set_state::ReplicaState;
use crate::shards::shard::PeerId;
use crate::shards::shard_holder::shard_mapping::ShardKeyMapping;

pub fn validate_transfer_exists(
    transfer_key: &ShardTransferKey,
    current_transfers: &HashSet<ShardTransfer>,
) -> CollectionResult<()> {
    if !current_transfers.iter().any(|t| &t.key() == transfer_key) {
        return Err(CollectionError::bad_request(format!(
            "There is no transfer for shard {} from {} to {}",
            transfer_key.shard_id, transfer_key.from, transfer_key.to,
        )));
    }

    Ok(())
}

pub fn get_transfer(
    transfer_key: &ShardTransferKey,
    current_transfers: &HashSet<ShardTransfer>,
) -> Option<ShardTransfer> {
    current_transfers
        .iter()
        .find(|t| &t.key() == transfer_key)
        .cloned()
}

/// Confirms that the transfer does not conflict with any other active transfers
///
/// returns `None` if there is no conflicts, otherwise returns conflicting transfer
pub fn check_transfer_conflicts<'a, I>(
    transfer: &ShardTransfer,
    current_transfers: I,
) -> Option<ShardTransfer>
where
    I: Iterator<Item = &'a ShardTransfer>,
{
    let res = current_transfers
        .filter(|t| t.shard_id == transfer.shard_id)
        .find(|t| {
            t.from == transfer.from
                || t.to == transfer.from
                || t.from == transfer.to
                || t.to == transfer.to
        });
    res.cloned()
}

/// Same as `check_transfer_conflicts` but doesn't allow transfers to/from the same peer
/// more than once for the whole collection
pub fn check_transfer_conflicts_strict<'a, I>(
    transfer: &ShardTransfer,
    mut current_transfers: I,
) -> Option<ShardTransfer>
where
    I: Iterator<Item = &'a ShardTransfer>,
{
    let res = current_transfers.find(|t| {
        t.from == transfer.from
            || t.to == transfer.from
            || t.from == transfer.to
            || t.to == transfer.to
    });
    res.cloned()
}

/// Confirms that the transfer makes sense with the current state cluster
///
/// Checks:
/// 1. If `from` and `to` exists
/// 2. If `from` have local shard and it is active
/// 3. If there is no active transfers which involve `from` or `to`
/// 4. If a target shard is only set for resharding transfers
///
/// For resharding transfers this also checks:
/// 1. If the source and target shards are different
/// 2. If the source and target shards share the same shard key
///
/// If validation fails, return `BadRequest` error.
pub fn validate_transfer(
    transfer: &ShardTransfer,
    all_peers: &HashSet<PeerId>,
    source_replicas: Option<&HashMap<PeerId, ReplicaState>>,
    destination_replicas: Option<&HashMap<PeerId, ReplicaState>>,
    current_transfers: &HashSet<ShardTransfer>,
    shards_key_mapping: &ShardKeyMapping,
) -> CollectionResult<()> {
    let Some(source_replicas) = source_replicas else {
        return Err(CollectionError::bad_request(format!(
            "Shard {} does not exist",
            transfer.shard_id,
        )));
    };

    if !all_peers.contains(&transfer.from) {
        return Err(CollectionError::bad_request(format!(
            "Peer {} does not exist",
            transfer.from,
        )));
    }

    if !all_peers.contains(&transfer.to) {
        return Err(CollectionError::bad_request(format!(
            "Peer {} does not exist",
            transfer.to,
        )));
    }

    // We allow transfers *from* `ReshardingScaleDown` replicas, because they contain a *superset*
    // of points in a regular replica
    let is_active = matches!(
        source_replicas.get(&transfer.from),
        Some(ReplicaState::Active | ReplicaState::ReshardingScaleDown),
    );

    if !is_active {
        return Err(CollectionError::bad_request(format!(
            "Shard {} is not active on peer {}",
            transfer.shard_id, transfer.from,
        )));
    }

    // If transfer with this key already exist, there are two possible cases:
    // - either we apply identical, but *conflicting* operation
    // - or we re-apply *the same* operation after a crash
    //
    // We can distinguish between the two, because *last step* of `start_resharding`
    // sets destination replica state to `Partial`.
    //
    // If destination replica *is* in `Partial` state, we should reject conflicting operation.
    // If destination replica is *not* in `Partial` state, we should re-apply existing operation.
    if get_transfer(&transfer.key(), current_transfers).is_some() {
        // Resharding/filtered transfers have separate destination shard
        let destination_replicas = destination_replicas.unwrap_or(source_replicas);

        let is_applied = destination_replicas
            .get(&transfer.to)
            .is_some_and(|state| state.is_partial_or_recovery());

        if is_applied {
            return Err(CollectionError::bad_request(format!(
                "Shard {} is already involved in transfer {} -> {}",
                transfer.shard_id, transfer.from, transfer.to,
            )));
        }
    }

    // Exclude this key from conflict check, because we already checked for identical transfer
    // conflict above
    let other_transfers = current_transfers
        .iter()
        .filter(|other| transfer.key() != other.key());

    if let Some(existing_transfer) = check_transfer_conflicts(transfer, other_transfers) {
        return Err(CollectionError::bad_request(format!(
            "Shard {} is already involved in transfer {} -> {}",
            transfer.shard_id, existing_transfer.from, existing_transfer.to,
        )));
    }

    if transfer.method == Some(ShardTransferMethod::ReshardingStreamRecords) {
        let Some(destination_replicas) = destination_replicas else {
            return Err(CollectionError::bad_request(format!(
                "Destination shard {} does not exist",
                transfer.shard_id,
            )));
        };

        let Some(to_shard_id) = transfer.to_shard_id else {
            return Err(CollectionError::bad_request(
                "Target shard is not set for resharding transfer",
            ));
        };

        if transfer.shard_id == to_shard_id {
            return Err(CollectionError::bad_request(format!(
                "Source and target shard must be different for resharding transfer, both are {to_shard_id}",
            )));
        }

        if let Some(ReplicaState::Dead) = destination_replicas.get(&transfer.to) {
            return Err(CollectionError::bad_request(format!(
                "Resharding shard transfer can't be started, \
                 because destination shard {}/{to_shard_id} is dead",
                transfer.to,
            )));
        }

        // Both shard IDs must share the same shard key
        let source_shard_key = shards_key_mapping
            .iter()
            .find(|(_, shard_ids)| shard_ids.contains(&to_shard_id))
            .map(|(key, _)| key);
        let target_shard_key = shards_key_mapping
            .iter()
            .find(|(_, shard_ids)| shard_ids.contains(&to_shard_id))
            .map(|(key, _)| key);
        if source_shard_key != target_shard_key {
            return Err(CollectionError::bad_request(format!(
                "Source and target shard must have the same shard key, but they have {source_shard_key:?} and {target_shard_key:?}",
            )));
        }
    } else if transfer.filter.is_some() {
        let Some(destination_replicas) = destination_replicas else {
            return Err(CollectionError::bad_request(format!(
                "Destination shard {} does not exist",
                transfer.shard_id,
            )));
        };

        let Some(to_shard_id) = transfer.to_shard_id else {
            return Err(CollectionError::bad_request(
                "Target shard is not set for filtered points transfer",
            ));
        };

        if transfer.shard_id == to_shard_id {
            return Err(CollectionError::bad_request(format!(
                "Source and target shard must be different for filtered points transfer, both are {to_shard_id}",
            )));
        }

        if let Some(ReplicaState::Dead) = destination_replicas.get(&transfer.to) {
            return Err(CollectionError::bad_request(format!(
                "Filtered shard transfer can't be started, \
                     because destination shard {}/{to_shard_id} is dead",
                transfer.to,
            )));
        }
    } else if let Some(to_shard_id) = transfer.to_shard_id {
        return Err(CollectionError::bad_request(format!(
            "Target shard {to_shard_id} can only be set for {:?} or filtered streaming records transfers",
            ShardTransferMethod::ReshardingStreamRecords,
        )));
    }

    Ok(())
}
