use anyhow::Result;
use clap::Subcommand;
use k256::ecdsa::{signature::Signer, Signature, SigningKey};
use sha2::{Digest, Sha256};
use serde_json::json;
use crate::{keystore, rpc::Client};

#[derive(Subcommand)]
pub enum NodeCmd {
    /// Show chain info and latest block
    Info,
    /// Show a specific block by height
    Block {
        height: u64,
    },
    /// Show account state
    Account {
        address: String,
    },
    /// Stake Shell tokens (StakeOp)
    Stake {
        #[arg(long)] amount: u64,
        #[arg(long, default_value = "default")] from: String,
    },
    /// Unstake Shell tokens (StakeOp)
    Unstake {
        #[arg(long)] amount: u64,
        #[arg(long, default_value = "default")] from: String,
    },
    /// Register as a validator (ValidatorRegister)
    Register {
        #[arg(long)] stake: u64,
        #[arg(long, default_value = "default")] from: String,
    },
}

pub async fn run(cmd: NodeCmd, rpc: &Client) -> Result<()> {
    match cmd {
        NodeCmd::Info => {
            let info = rpc.call("nu_chainInfo", vec![]).await?;
            println!("{}", serde_json::to_string_pretty(&info)?);
        }
        NodeCmd::Block { height } => {
            let block = rpc.call("nu_getBlock", vec![json!(height)]).await?;
            if block.is_null() {
                println!("Block {height} not found");
            } else {
                println!("{}", serde_json::to_string_pretty(&block)?);
            }
        }
        NodeCmd::Account { address } => {
            let account = rpc.call("nu_getAccount", vec![json!(address)]).await?;
            let balance_shell: u64 = account["balance"].as_u64().unwrap_or(0);
            let staked_shell:  u64 = account["staked"].as_u64().unwrap_or(0);
            println!("Address : {address}");
            println!("Balance : {} Shell ({:.4} NU)", balance_shell, balance_shell as f64 / 100_000.0);
            println!("Staked  : {} Shell ({:.4} NU)", staked_shell,  staked_shell  as f64 / 100_000.0);
            println!("Nonce   : {}", account["nonce"].as_u64().unwrap_or(0));
            println!("PoN     : {}", account["pon_score"].as_f64().unwrap_or(1.0));
        }
        NodeCmd::Stake { amount, from } => {
            let result = stake_op(rpc, &from, true, amount).await?;
            println!("Staked {} Shell: {}", amount, result["tx_id"].as_str().unwrap_or("?"));
        }
        NodeCmd::Unstake { amount, from } => {
            let result = stake_op(rpc, &from, false, amount).await?;
            println!("Unstaked {} Shell: {}", amount, result["tx_id"].as_str().unwrap_or("?"));
        }
        NodeCmd::Register { stake, from } => {
            let (key, sender) = load_key(&from)?;
            let account = rpc.call("nu_getAccount", vec![json!(sender)]).await?;
            let nonce: u64 = account["nonce"].as_u64().unwrap_or(0);

            let payload = json!({ "ValidatorRegister": { "stake": stake } });
            let result  = send_tx(&key, &sender, nonce, payload, rpc).await?;
            println!("Registered as validator: {}", result["tx_id"].as_str().unwrap_or("?"));
        }
    }
    Ok(())
}

async fn stake_op(rpc: &Client, label: &str, stake: bool, amount: u64) -> Result<serde_json::Value> {
    let (key, sender) = load_key(label)?;
    let account = rpc.call("nu_getAccount", vec![json!(sender)]).await?;
    let nonce: u64 = account["nonce"].as_u64().unwrap_or(0);
    let payload = json!({ "StakeOp": { "stake": stake, "amount": amount } });
    send_tx(&key, &sender, nonce, payload, rpc).await
}

fn load_key(label: &str) -> Result<(SigningKey, String)> {
    let dir       = crate::config::keystore_dir();
    let path      = dir.join(format!("{label}.key"));
    let password  = keystore::prompt_password(false)?;
    let key_bytes = keystore::load_encrypted(&path, &password)?;
    let key       = SigningKey::from_bytes(key_bytes.as_slice().into())?;
    let sender    = hex::encode(key.verifying_key().to_sec1_bytes());
    Ok((key, sender))
}

async fn send_tx(
    key:     &SigningKey,
    sender:  &str,
    nonce:   u64,
    payload: serde_json::Value,
    rpc:     &Client,
) -> Result<serde_json::Value> {
    let payload_bytes  = serde_json::to_vec(&payload)?;
    let tx_id          = hex::encode(Sha256::digest(&payload_bytes));
    let sig: Signature = key.sign(tx_id.as_bytes());
    let sig_bytes      = sig.to_bytes().to_vec();

    let tx = json!({
        "tx_id":   tx_id,
        "sender":  sender,
        "nonce":   nonce,
        "fee":     0,
        "sig":     sig_bytes,
        "payload": payload,
    });
    Ok(rpc.call("nu_sendRawTx", vec![tx]).await?)
}