ghost-eye/src/components/validator/withdraw_popup.rs

use crossterm::event::{KeyCode, KeyEvent, KeyEventKind};
use color_eyre::Result;
use ratatui::{
    layout::{Alignment, Constraint, Flex, Layout, Rect}, 
    widgets::{Block, Clear, Paragraph}, 
    Frame
};
use tokio::sync::mpsc::UnboundedSender;
use std::sync::mpsc::Sender;

use super::{Component, PartialComponent, CurrentTab};
use crate::{
    action::Action, 
    config::Config, 
    palette::StylePalette, 
    types::{ActionLevel, ActionTarget}, 
};

#[derive(Debug)]
pub struct WithdrawPopup {
    is_active: bool,
    action_tx: Option<UnboundedSender<Action>>,
    network_tx: Option<Sender<Action>>,
    stash_account: [u8; 32],
    secret_seed: [u8; 32],
    slashing_spans_length: u32,
    unlocking_is_empty: bool,
    existential_deposit: u128,
    active_balance: u128,
    palette: StylePalette
}

impl Default for WithdrawPopup {
    fn default() -> Self {
        Self::new()
    }
}

impl WithdrawPopup {
    pub fn new() -> Self {
        Self {
            is_active: false,
            stash_account: [0u8; 32],
            secret_seed: [0u8; 32],
            slashing_spans_length: 0u32,
            unlocking_is_empty: false,
            existential_deposit: 0u128,
            active_balance: 0u128,
            action_tx: None,
            network_tx: None,
            palette: StylePalette::default(),
        }
    }

    fn close_popup(&mut self) {
        self.is_active = false;
        if let Some(action_tx) = &self.action_tx {
            let _ = action_tx.send(Action::ClosePopup);
        }
    }

    fn stash_should_be_killed(&self) -> bool {
        self.unlocking_is_empty && self.active_balance < self.existential_deposit
    }

    fn proceed(&mut self) {
        if let Some(network_tx) = &self.network_tx { 
            let spans_needed = if self.stash_should_be_killed() {
                if let Some(action_tx) = &self.action_tx {
                    let _ = action_tx.send(Action::EventLog(
                            "Current stash account will be killed during this transaction".to_string(),
                            ActionLevel::Warn, 
                            ActionTarget::ValidatorLog));
                }
                self.slashing_spans_length
            } else {
                0
            };
            let _ = network_tx.send(Action::WithdrawUnbondedFrom(self.secret_seed, spans_needed));
            if let Some(action_tx) = &self.action_tx {
                let _ = action_tx.send(Action::ClosePopup);
            }
        }
    }
}

impl PartialComponent for WithdrawPopup {
    fn set_active(&mut self, current_tab: CurrentTab) {
        match current_tab {
            CurrentTab::WithdrawPopup => self.is_active = true,
            _ => self.is_active = false,
        };
    }
}

impl Component for WithdrawPopup {
    fn register_network_handler(&mut self, tx: Sender<Action>) -> Result<()> {
        self.network_tx = Some(tx);
        Ok(())
    }

    fn register_action_handler(&mut self, tx: UnboundedSender<Action>) -> Result<()> {
        self.action_tx = Some(tx);
        Ok(())
    }

    fn register_config_handler(&mut self, config: Config) -> Result<()> {
        if let Some(style) = config.styles.get(&crate::app::Mode::Wallet) {
            self.palette.with_normal_style(style.get("normal_style").copied());
            self.palette.with_normal_border_style(style.get("normal_border_style").copied());
            self.palette.with_normal_title_style(style.get("normal_title_style").copied());
            self.palette.with_popup_style(style.get("popup_style").copied());
            self.palette.with_popup_title_style(style.get("popup_title_style").copied());
        }
        Ok(())
    }

    fn handle_key_event(&mut self, key: KeyEvent) -> Result<Option<Action>> {
        if self.is_active && key.kind == KeyEventKind::Press {
            match key.code {
                KeyCode::Enter => self.proceed(),
                KeyCode::Esc => self.close_popup(), 
                _ => {},
            };
        }
        Ok(None)
    }

    fn update(&mut self, action: Action) -> Result<Option<Action>> {
        match action {
            Action::SetExistentialDeposit(existential_deposit) => self.existential_deposit = existential_deposit,
            Action::SetStashAccount(account_id) => self.stash_account = account_id,
            Action::SetStashSecret(secret_seed) => self.secret_seed = secret_seed,
            Action::SetSlashingSpansLength(length, account_id) if self.stash_account == account_id => 
                self.slashing_spans_length = length as u32,
            Action::SetStakedAmountRatio(_, active_balance, account_id) if self.stash_account == account_id => 
                self.active_balance = active_balance.unwrap_or_default(),
            Action::SetValidatorEraUnlocking(unlockings, account_id) if self.stash_account == account_id =>
                self.unlocking_is_empty = unlockings.is_empty(),
            _ => {}
        };
        Ok(None)
    }

    fn draw(&mut self, frame: &mut Frame, area: Rect) -> Result<()> {
        if self.is_active {
            let (border_style, border_type) = self.palette.create_popup_style();
            let text = if self.stash_should_be_killed() {
                " Do you want to kill current stash account?"
            } else {
                " Do you want to withdraw all unbonded funds?"
            };
            let input = Paragraph::new(text)
                .block(Block::bordered()
                    .border_style(border_style)
                    .border_type(border_type)
                    .title_style(self.palette.create_popup_title_style())
                    .title_alignment(Alignment::Right)
                    .title("Enter to proceed / Esc to close"));
            let v = Layout::vertical([Constraint::Max(3)]).flex(Flex::Center);
            let h = Layout::horizontal([Constraint::Max(50)]).flex(Flex::Center);
            let [area] = v.areas(area);
            let [area] = h.areas(area);

            frame.render_widget(Clear, area);
            frame.render_widget(input, area);
        }
        Ok(())
    }
}