use crate::types::{Deposit, TellerOperationId, Withdraw};
use near_sdk::{
    json_types::U128,
    near,
    store::{IterableMap, IterableSet, LookupMap},
    AccountId,
};

/// Storage keys for teller collections
#[derive(near_sdk::BorshStorageKey, borsh::BorshSerialize)]
enum TellerStorageKey {
    PendingDeposits,
    PendingWithdraws,
    UserDepositIds,
    UserWithdrawIds,
    UserDepositIdSet { account_hash: Vec<u8> },
    UserWithdrawIdSet { account_hash: Vec<u8> },
}

/// State for managing pending operations
#[near(serializers=[borsh])]
pub struct TellerState {
    /// Operation ID -> Deposit mapping
    pub pending_deposits: IterableMap<TellerOperationId, Deposit>,
    /// Operation ID -> Withdraw mapping
    pub pending_withdraws: IterableMap<TellerOperationId, Withdraw>,
    /// AccountId -> IterableSet<operation_id> for deposits
    pub user_deposit_ids: LookupMap<AccountId, IterableSet<TellerOperationId>>,
    /// AccountId -> IterableSet<operation_id> for withdrawals
    pub user_withdraw_ids: LookupMap<AccountId, IterableSet<TellerOperationId>>,
    /// Next operation ID to assign
    pub next_operation_id: TellerOperationId,
}

impl TellerState {
    pub fn new() -> Self {
        Self {
            pending_deposits: IterableMap::new(TellerStorageKey::PendingDeposits),
            pending_withdraws: IterableMap::new(TellerStorageKey::PendingWithdraws),
            user_deposit_ids: LookupMap::new(TellerStorageKey::UserDepositIds),
            user_withdraw_ids: LookupMap::new(TellerStorageKey::UserWithdrawIds),
            next_operation_id: 1,
        }
    }

    /// Helper to add operation to user tracking
    fn add_to_user_tracking(
        user_lookup: &mut LookupMap<AccountId, IterableSet<TellerOperationId>>,
        user_id: &AccountId,
        operation_id: TellerOperationId,
        is_deposit: bool,
    ) {
        let user_operations = user_lookup.entry(user_id.clone()).or_insert_with(|| {
            // Create new IterableSet with unique storage key based on account
            let account_hash = near_sdk::env::sha256(user_id.as_bytes());
            if is_deposit {
                IterableSet::new(TellerStorageKey::UserDepositIdSet { account_hash })
            } else {
                IterableSet::new(TellerStorageKey::UserWithdrawIdSet { account_hash })
            }
        });
        user_operations.insert(operation_id);
    }

    /// Add a deposit to pending queue
    pub fn add_deposit(&mut self, deposit: Deposit) -> u64 {
        let operation_id = self.next_operation_id;
        self.next_operation_id += 1;

        // Add to pending deposits
        let owner_id = deposit.owner_id.clone();
        self.pending_deposits.insert(operation_id, deposit);

        // Track user deposits
        Self::add_to_user_tracking(&mut self.user_deposit_ids, &owner_id, operation_id, true);

        operation_id
    }

    /// Add a withdrawal to pending queue
    pub fn add_withdraw(&mut self, withdraw: Withdraw) -> u64 {
        let operation_id = self.next_operation_id;
        self.next_operation_id += 1;

        // Add to pending withdrawals
        let owner_id = withdraw.owner_id.clone();
        self.pending_withdraws.insert(operation_id, withdraw);

        // Track user withdrawals
        Self::add_to_user_tracking(&mut self.user_withdraw_ids, &owner_id, operation_id, false);

        operation_id
    }

    /// Get a deposit by operation ID
    pub fn get_deposit(&self, operation_id: u64) -> Option<Deposit> {
        self.pending_deposits.get(&operation_id).cloned()
    }

    /// Get a withdrawal by operation ID
    pub fn get_withdraw(&self, operation_id: u64) -> Option<Withdraw> {
        self.pending_withdraws.get(&operation_id).cloned()
    }

    /// Helper to remove operation from user tracking
    fn remove_from_user_tracking(
        user_lookup: &mut LookupMap<AccountId, IterableSet<u64>>,
        owner_id: &AccountId,
        operation_id: u64,
    ) {
        // Check if owner exists and remove operation, tracking if set becomes empty
        let is_empty = if let Some(user_operations) = user_lookup.get_mut(owner_id) {
            user_operations.remove(&operation_id);
            user_operations.is_empty()
        } else {
            return; // Owner doesn't exist, nothing to remove
        };

        // Remove the owner entry entirely if no operations remain
        if is_empty {
            user_lookup.remove(owner_id);
        }
    }

    /// Remove a deposit from pending queue
    pub fn remove_deposit(&mut self, operation_id: u64, owner_id: &AccountId) -> Option<Deposit> {
        let deposit = self.pending_deposits.remove(&operation_id);

        if deposit.is_some() {
            Self::remove_from_user_tracking(&mut self.user_deposit_ids, owner_id, operation_id);
        }

        deposit
    }

    /// Remove a withdrawal from pending queue
    pub fn remove_withdraw(&mut self, operation_id: u64, owner_id: &AccountId) -> Option<Withdraw> {
        let withdraw = self.pending_withdraws.remove(&operation_id);

        if withdraw.is_some() {
            Self::remove_from_user_tracking(&mut self.user_withdraw_ids, owner_id, operation_id);
        }

        withdraw
    }

    /// Get pending deposit IDs for a user
    pub fn get_user_deposit_ids(&self, account_id: &AccountId) -> Vec<u64> {
        self.user_deposit_ids
            .get(account_id)
            .map(|set| set.iter().cloned().collect())
            .unwrap_or_default()
    }

    /// Get pending withdrawal IDs for a user
    pub fn get_user_withdraw_ids(&self, account_id: &AccountId) -> Vec<u64> {
        self.user_withdraw_ids
            .get(account_id)
            .map(|set| set.iter().cloned().collect())
            .unwrap_or_default()
    }

    /// Get total number of pending deposits
    pub fn deposits_count(&self) -> u64 {
        self.pending_deposits.len() as u64
    }

    /// Get total number of pending withdrawals
    pub fn withdraws_count(&self) -> u64 {
        self.pending_withdraws.len() as u64
    }

    /// Iterator over all pending withdrawals
    pub fn withdraws_iter(&self) -> impl Iterator<Item = (u64, Withdraw)> + '_ {
        self.pending_withdraws
            .iter()
            .map(|(id, withdraw)| (*id, withdraw.clone()))
    }

    /// Get all pending deposits as vector (for processing)
    pub fn get_all_pending_deposits(&self) -> Vec<(u64, Deposit)> {
        self.pending_deposits
            .iter()
            .map(|(id, deposit)| (*id, deposit.clone()))
            .collect()
    }

    /// Get all pending withdrawals as vector (for processing)
    pub fn get_all_pending_withdraws(&self) -> Vec<(u64, Withdraw)> {
        self.pending_withdraws
            .iter()
            .map(|(id, withdraw)| (*id, withdraw.clone()))
            .collect()
    }

    /// Get all unconfirmed withdrawals (ready for confirmation)
    pub fn get_unconfirmed_withdraws(&self) -> Vec<(u64, Withdraw)> {
        self.pending_withdraws
            .iter()
            .filter(|(_, withdraw)| !withdraw.confirmed)
            .map(|(id, withdraw)| (*id, withdraw.clone()))
            .collect()
    }

    /// Get all confirmed withdrawals (ready for processing)
    pub fn get_confirmed_withdraws(&self) -> Vec<(u64, Withdraw)> {
        self.pending_withdraws
            .iter()
            .filter(|(_, withdraw)| withdraw.confirmed)
            .map(|(id, withdraw)| (*id, withdraw.clone()))
            .collect()
    }

    /// Check if a withdrawal is confirmed
    pub fn is_withdraw_confirmed(&self, operation_id: u64) -> bool {
        self.pending_withdraws
            .get(&operation_id)
            .map(|w| w.confirmed)
            .unwrap_or(false)
    }

    /// Get confirmed share price for a withdrawal
    pub fn get_confirmed_share_price(&self, operation_id: u64) -> Option<U128> {
        self.pending_withdraws
            .get(&operation_id)
            .and_then(|w| w.confirmed_share_price)
    }

    /// Restore a withdrawal to pending queue with specific operation ID
    /// Returns true if restored successfully, false if operation ID already exists
    pub fn restore_withdraw(&mut self, operation_id: u64, withdraw: Withdraw) -> bool {
        // Check if operation ID already exists - this indicates state corruption
        if self.pending_withdraws.get(&operation_id).is_some() {
            return false;
        }

        // Add to pending withdrawals with specific ID
        let owner_id = withdraw.owner_id.clone();
        self.pending_withdraws.insert(operation_id, withdraw);

        // Track user withdrawals
        Self::add_to_user_tracking(&mut self.user_withdraw_ids, &owner_id, operation_id, false);

        true
    }

    /// Restore a deposit to pending queue with specific operation ID
    /// Returns true if restored successfully, false if operation ID already exists
    pub fn restore_deposit(&mut self, operation_id: u64, deposit: Deposit) -> bool {
        // Check if operation ID already exists - this indicates state corruption
        if self.pending_deposits.get(&operation_id).is_some() {
            return false;
        }

        // Add to pending deposits with specific ID
        let owner_id = deposit.owner_id.clone();
        self.pending_deposits.insert(operation_id, deposit);

        // Track user deposits
        Self::add_to_user_tracking(&mut self.user_deposit_ids, &owner_id, operation_id, true);

        true
    }
}

#[cfg(test)]
mod tests;