ghost-eye/src/components/wallet/accounts.rs

use std::path::PathBuf;
use std::fs::File;
use std::io::{Write, BufRead, BufReader};

use color_eyre::Result;
use crossterm::event::{KeyCode, KeyEvent};
use ratatui::layout::{Constraint, Margin};
use ratatui::{
    text::Text,
    layout::{Alignment, Rect}, 
    widgets::{
        Block, Cell, Row, Table, TableState, Scrollbar, Padding,
        ScrollbarOrientation, ScrollbarState,
    }, 
    Frame
};
use subxt::{
    tx::PairSigner,
    ext::sp_core::{
        Pair as PairT, 
        sr25519::Pair,
        crypto::{Ss58Codec, Ss58AddressFormat, AccountId32},
    },
};
use tokio::sync::mpsc::UnboundedSender;
use std::sync::mpsc::Sender;

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

struct AccountInfo {
    name: String,
    address: String,
    account_id: [u8; 32],
    seed: String,
    pair_signer: PairSigner<CasperConfig, Pair>,
}

pub struct Accounts {
    is_active: bool,
    action_tx: Option<UnboundedSender<Action>>,
    network_tx: Option<Sender<Action>>,
    wallet_keys_file: PathBuf,
    palette: StylePalette,
    scroll_state: ScrollbarState,
    table_state: TableState,
    wallet_keys: Vec<AccountInfo>,
    balances: std::collections::HashMap<[u8; 32], SystemAccount>,
}

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

impl Accounts {
    pub fn new() -> Self {
        Self {
            is_active: false,
            wallet_keys_file: Default::default(),
            action_tx: None,
            network_tx: None,
            scroll_state: ScrollbarState::new(0),
            table_state: TableState::new(),
            wallet_keys: Vec::new(),
            balances: Default::default(),
            palette: StylePalette::default(),
        }
    }

    fn log_event(&mut self, message: String, level: ActionLevel) {
        if let Some(action_tx) = &self.action_tx {
            let _ = action_tx.send(Action::WalletLog(message, level));
        }
    }

    fn send_balance_request(&mut self, account_id: [u8; 32], remove: bool) {
        if let Some(action_tx) = &self.network_tx {
            let _ = action_tx.send(Action::BalanceRequest(account_id, remove));
        }
    }

    fn set_balance_active(&mut self, index: usize) {
        if let Some(action_tx) = &self.action_tx {
            let account_id = self.wallet_keys[index].account_id;
            let _ = action_tx.send(Action::BalanceSetActive(
                    self.balances.get(&account_id).cloned()));
        }
    }

    fn update_account_name(&mut self) {
        if let Some(index) = self.table_state.selected() {
            if let Some(action_tx) = &self.action_tx {
                let _ = action_tx.send(Action::RenameAccount(
                        self.wallet_keys[index].name.clone()
                ));
            }
        }
    }

    fn create_new_account(&mut self, name: String) {
        let (pair, seed) = Pair::generate();
        let secret_seed = hex::encode(seed);
        let account_id = pair.public().0;
        let pair_signer = PairSigner::<CasperConfig, Pair>::new(pair);
        let address = AccountId32::from(seed.clone())
            .to_ss58check_with_version(Ss58AddressFormat::custom(1996));

        self.log_event(
            format!("new wallet '{}' with public address {} created", 
                &name, &address),
                ActionLevel::Info);

        self.send_balance_request(account_id, false);
        self.wallet_keys.push(AccountInfo {
            name, 
            address, 
            account_id,
            seed: secret_seed, 
            pair_signer,
        });
        self.last_row();
        self.save_to_file();
    }

    fn rename_account(&mut self, new_name: String) {
        if let Some(index) = self.table_state.selected() {
            let old_name = self.wallet_keys[index].name.clone();

            self.log_event(format!("wallet '{}' renamed to {}", 
                    &new_name, &old_name),
                    ActionLevel::Info);
            self.wallet_keys[index].name = new_name;
            self.save_to_file();

        }
    }

     fn swap_up(&mut self) {
        if let Some(src_index) = self.table_state.selected() {
            let dst_index = src_index.saturating_sub(1);
            if src_index > dst_index {
                self.wallet_keys.swap(src_index, dst_index);
                self.previous_row();
                self.save_to_file();
            }
        }
    }

    fn swap_down(&mut self) {
        if let Some(src_index) = self.table_state.selected() {
            let dst_index = src_index + 1;
            if dst_index < self.wallet_keys.len() {
                self.wallet_keys.swap(src_index, dst_index);
                self.next_row();
                self.save_to_file();
            }
        }
    }

    fn delete_row(&mut self) {
        if let Some(index) = self.table_state.selected() {
            if self.wallet_keys.len() > 1 {
                let wallet = self.wallet_keys.remove(index);
                self.send_balance_request(wallet.account_id, true);
                self.previous_row();
                self.save_to_file();

                self.log_event(format!("wallet `{}` with public address {} removed",
                        &wallet.name, &wallet.address),
                        ActionLevel::Warn);
            }
        }
    }

    fn save_to_file(&mut self) {
        let mut file = File::create(&self.wallet_keys_file).unwrap();
        for wallet in self.wallet_keys.iter() {
            writeln!(file, "{}:0x{}", wallet.name, &wallet.seed).unwrap();
        }
    }

    fn read_or_create(&mut self, file_path: &PathBuf) -> Result<()> {
        assert!(self.wallet_keys.len() == 0, "wallet_keys already exists");
        match File::open(file_path) {
            Ok(file) => {
                let reader = BufReader::new(file);
                for line in reader.lines() {
                    let line = line?.replace("\n", "");
                    let line_split_at = line.find(":").unwrap_or(line.len());
                    let (wallet_name, wallet_key) = line.split_at(line_split_at);
                    let wallet_key = &wallet_key[3..];

                    let seed: [u8; 32] = hex::decode(wallet_key)
                        .expect("stored seed is valid hex string; qed")
                        .as_slice()
                        .try_into()
                        .expect("stored seed is valid length; qed");

                    let pair = Pair::from_seed(&seed);
                    let account_id = pair.public().0;
                    let address = AccountId32::from(account_id)
                        .to_ss58check_with_version(Ss58AddressFormat::custom(1996));
                    let pair_signer = PairSigner::<CasperConfig, Pair>::new(pair);

                    self.send_balance_request(account_id, false);
                    self.wallet_keys.push(AccountInfo {
                            name: wallet_name.to_string(), 
                            account_id,
                            address, 
                            seed: wallet_key.to_string(), 
                            pair_signer,
                    });
                }
                self.log_event(format!("read {} wallets from disk", 
                        self.wallet_keys.len()),
                        ActionLevel::Info);
            },
            Err(_) => {
                let (pair, seed) = match std::fs::read_to_string("/etc/ghost/wallet-key") {
                    Ok(content) => {
                        let content = content.replace("\n", "");
                        let content = &content[2..];
                        let seed: [u8; 32] = hex::decode(content)
                            .expect("stored seed is valid hex string; qed")
                            .as_slice()
                            .try_into()
                            .expect("stored seed is valid length; qed");

                        self.log_event(
                            "wallet read from the `/etc/ghost/wallet-key`".to_string(),
                            ActionLevel::Warn);

                        let pair = Pair::from_seed(&seed);
                        (pair, seed)
                    }
                    Err(_) => {
                        self.log_event(
                            "no wallets found on disk, new wallet created".to_string(),
                            ActionLevel::Warn);
                        let (pair, seed) = Pair::generate();
                        (pair, seed)
                    }
                };

                let secret_seed = hex::encode(seed);
                let account_id = pair.public().0;
                let address = AccountId32::from(pair.public().0)
                    .to_ss58check_with_version(Ss58AddressFormat::custom(1996));
                let pair_signer = PairSigner::<CasperConfig, Pair>::new(pair);

                let mut new_file = File::create(file_path)?;
                writeln!(new_file, "ghostie:0x{}", &secret_seed)?;

                self.send_balance_request(account_id, false);
                self.wallet_keys.push(AccountInfo {
                        name: "ghostie".to_string(), 
                        address, 
                        account_id,
                        seed: secret_seed, 
                        pair_signer,
                });
            }
        };
        self.table_state.select(Some(0));
        self.scroll_state = self.scroll_state.content_length(self.wallet_keys.len());
        self.set_balance_active(0);
        Ok(())
    }

    fn first_row(&mut self) {
        if self.wallet_keys.len() > 0 {
            self.table_state.select(Some(0));
            self.scroll_state = self.scroll_state.position(0);
            self.set_balance_active(0);
        }
    }

    fn next_row(&mut self) {
        let i = match self.table_state.selected() {
            Some(i) => {
                if i >= self.wallet_keys.len() - 1 {
                    i
                } else {
                    i + 1
                }
            },
            None => 0,
        };
        self.table_state.select(Some(i));
        self.scroll_state = self.scroll_state.position(i);
        self.set_balance_active(i);
    }

    fn last_row(&mut self) {
        if self.wallet_keys.len() > 0 {
            let last = self.wallet_keys.len() - 1;
            self.table_state.select(Some(last));
            self.scroll_state = self.scroll_state.position(last);
            self.set_balance_active(last);
        }
    }

    fn previous_row(&mut self) {
        let i = match self.table_state.selected() {
            Some(i) => {
                if i == 0 {
                    0
                } else {
                    i - 1
                }
            },
            None => 0
        };
        self.table_state.select(Some(i));
        self.scroll_state = self.scroll_state.position(i);
        self.set_balance_active(i);
    }

    fn prepare_u128(&self, value: u128, after: usize, ticker: Option<&str>) -> String {
        let value = value as f64 / 10f64.powi(18);
        format!("{:.after$}{}", value, ticker.unwrap_or_default())
    }
}

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

impl Component for Accounts {
    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_hover_style(style.get("hover_style").copied());
            self.palette.with_normal_border_style(style.get("normal_border_style").copied());
            self.palette.with_hover_border_style(style.get("hover_border_style").copied());
            self.palette.with_normal_title_style(style.get("normal_title_style").copied());
            self.palette.with_hover_title_style(style.get("hover_title_style").copied());
            self.palette.with_highlight_style(style.get("highlight_style").copied());
            self.palette.with_scrollbar_style(style.get("scrollbar_style").copied());
        }

        let mut wallet_keys_file = config.config.data_dir;
        wallet_keys_file.push("wallet-keys");
        self.read_or_create(&wallet_keys_file)?;
        self.wallet_keys_file = wallet_keys_file;
        Ok(())
    }

    fn update(&mut self, action: Action) -> Result<Option<Action>> {
        match action {
            Action::NewAccount(name) => self.create_new_account(name),
            Action::UpdateAccountName(new_name) => self.rename_account(new_name),
            Action::BalanceResponse(account_id, balance) => {
                if self.wallet_keys.iter().any(|wallet| wallet.account_id == account_id) {
                    let _ = self.balances.insert(account_id, balance);
                    if let Some(index) = self.table_state.selected() {
                        if self.wallet_keys[index].account_id == account_id {
                            self.set_balance_active(index);
                        }
                    }
                }
            },
            _ => {}
        };
        Ok(None)
    }

    fn handle_key_event(&mut self, key: KeyEvent) -> Result<Option<Action>> {
        match key.code {
           KeyCode::Up | KeyCode::Char('k') if self.is_active => self.previous_row(), 
           KeyCode::Down | KeyCode::Char('j') if self.is_active => self.next_row(), 
           KeyCode::Char('K') if self.is_active => self.swap_up(),
           KeyCode::Char('J') if self.is_active => self.swap_down(),
           KeyCode::Char('g') if self.is_active => self.first_row(),
           KeyCode::Char('G') if self.is_active => self.last_row(),
           KeyCode::Char('D') if self.is_active => self.delete_row(),
           KeyCode::Char('R') if self.is_active => self.update_account_name(),
           _ => {},
        };
        Ok(None)
    }

    fn draw(&mut self, frame: &mut Frame, area: Rect) -> Result<()> {
        let [_, place, _] = super::account_layout(area);
        let (border_style, border_type) = self.palette.create_border_style(self.is_active);

        let table = Table::new(
            self.wallet_keys
                .iter()
                .map(|info| {
                    let balance = self.balances
                        .get(&info.account_id)
                        .map(|b| b.free)
                        .unwrap_or_default();
                    Row::new(vec![
                        Cell::from(Text::from(info.name.clone()).alignment(Alignment::Left)),
                        Cell::from(Text::from(info.address.clone()).alignment(Alignment::Center)), 
                        Cell::from(Text::from(self.prepare_u128(balance, 2, Some(" CSPR"))).alignment(Alignment::Right)),
                    ])
                }),
            [
                Constraint::Min(5),
                Constraint::Max(51),
                Constraint::Min(11),
            ],
        )
        .highlight_style(self.palette.create_highlight_style())
        .column_spacing(1)
        .block(Block::bordered()
            .border_style(border_style)
            .border_type(border_type)
            .padding(Padding::right(2))
            .title_alignment(Alignment::Right)
            .title_style(self.palette.create_title_style(false))
            .title("My Accounts"));

        let scrollbar = Scrollbar::default()
            .orientation(ScrollbarOrientation::VerticalRight)
            .begin_symbol(None)
            .end_symbol(None)
            .style(self.palette.create_scrollbar_style());

        frame.render_stateful_widget(table, place, &mut self.table_state);
        frame.render_stateful_widget(
            scrollbar,
            place.inner(Margin { vertical: 1, horizontal: 1 }),
            &mut self.scroll_state,
        );

        Ok(())
    }
}