//! Solana support for Icepick.
//!
//! # Command Line Operations
//!
//! The first thing you'll want is a wallet. For this, you can run the command
//! `icepick sol generate-wallet`. The output of the command will be a request to derive a key,
//! which Icepick will interpret upon subsequent invocations and replace with the derived key -
//! after appending the valid `bip44` path. You can then pipe the output of the command to
//! `icepick sol get-wallet-address`, which will take the derived key and format it as a Solana
//! address. In full, the commands look like:
//!
//! ```sh
//! icepick sol generate-wallet | icepick sol get-wallet-address
//! ```
//!
//! Next, you'll want to either airdrop some funds into the wallet if on devnet, or transfer some
//! funds using another source if on mainnet.
//!
//!
//! To transfer funds out of the wallet, you'll need a recent blockhash. The command
//! `icepick sol get-blockhash` can be used to get a recent blockhash. This blockhash must be used
//! within 150 blocks, or about 1 minute, or transactions won't be accepted.
//!
//! The blockhash is then used in the next command, `icepick sol transfer`, which also requires an
//! amount, a to-address, and a from-address. The output of this command can then be saved to
//! removable media, then transferred to an offline signer system where `icepick sol sign` is used
//! to sign the transaction, with the signed result also persisted to removable media. Once the
//! signed transaction is back on an online system, the transaction can be broadcasted using
//! `icepick sol broadcast`.
//!
//! ```sh
//! # On an online system
//! icepick sol get-blockhash > sdcard/blockhash.json
//!
//! # On an offline system
//! blockhash=$(jq -r .blob sdcard/blockhash.json)
//! icepick sol transfer $amount $to_address $from_address | icepick sol sign $blockhash > sdcard/transfer.json
//!
//! # On the online system, again
//! icepick sol broadcast --cluster devnet < sdcard/transfer.json
//! ```
//!
//! You may also want to transfer tokens on the Solana blockchain. You'll first need an account for
//! the token (let's use IPDBG, a devnet token I built for testing):
//!
//! ```sh
//! # On an online system
//! # Assume we have the wallet address as $wallet_address
//! icepick sol get-blockhash > sdcard/blockhash.json
//!
//! # On an offline system
//! blockhash=$(jq -r .blob sdcard/blockhash.json)
//! icepick sol get-token-info IPDBG > sdcard/ipdbg.json
//! token_address=$(jq -r .blob.token_address sdcard/ipdbg.json)
//! icepick sol create-token-account $wallet_address $token_address | icepick sol sign $blockhash > sdcard/create-account.json
//!
//! # On an online system
//! icepick sol broadcast --cluster devnet < sdcard/create-account.json
//! ```

use icepick_module::{
    help::{Argument, ArgumentType},
    Module,
};
use serde::{Deserialize, Serialize};
use solana_sdk::signer::Signer;
use std::str::FromStr;

// How does this not exist in solana_sdk.
const LAMPORTS_PER_SOL: u64 = 1_000_000_000;

#[derive(thiserror::Error, Debug)]
pub enum Error {}

#[derive(Serialize, Deserialize, Debug)]
#[serde(rename_all = "kebab-case")]
pub enum Cluster {
    Devnet,
    Testnet,
    #[serde(alias = "mainnet")]
    MainnetBeta,
}

impl std::str::FromStr for Cluster {
    type Err = &'static str;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "devnet" => Ok(Self::Devnet),
            "testnet" => Ok(Self::Testnet),
            "mainnet" => Ok(Self::MainnetBeta),
            "mainnet-beta" => Ok(Self::MainnetBeta),
            _ => Err("Invalid value"),
        }
    }
}

impl std::fmt::Display for Cluster {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Cluster::Devnet => f.write_str("devnet"),
            Cluster::Testnet => f.write_str("testnet"),
            Cluster::MainnetBeta => f.write_str("mainnet-beta"),
        }
    }
}

#[derive(Serialize, Deserialize, Debug)]
pub struct GetBlockhash {
    cluster: Option<Cluster>,
}

#[derive(Serialize, Deserialize, Debug)]
pub struct GenerateWallet {
    account: Option<String>,
}

#[derive(Serialize, Deserialize, Debug)]
pub struct GetWalletAddress {}

#[derive(Serialize, Deserialize, Debug)]
pub struct GetTokenInfo {
    token: String,
}

#[derive(Serialize, Deserialize, Debug)]
pub struct Transfer {
    amount: String,
    to_address: String,
    from_account: Option<String>,
    from_address: String,
    fee: Option<String>,
    fee_payer: Option<String>,
    fee_payer_address: Option<String>,
}

#[derive(Serialize, Deserialize, Debug)]
pub struct CreateTokenAccount {
    funder_address: Option<String>,
    wallet_address: String,
    token_address: String,
}

#[derive(Serialize, Deserialize, Debug)]
pub struct TransferToken {
    amount: String,
    token_address: String,
    to_address: String,
    from_account: Option<String>,
    from_address: String,
    decimals: String,
    fee: Option<String>,
    fee_payer: Option<String>,
    fee_payer_address: Option<String>,
}

#[derive(Serialize, Deserialize, Debug)]
pub struct Sign {
    blockhash: String,
}

#[derive(Serialize, Deserialize, Debug)]
pub struct Broadcast {
    cluster: Option<Cluster>,
}

#[derive(Serialize, Deserialize, Debug)]
pub struct Request {
    // NOTE: Can't use the proper XPrv type from Keyfork because Solana's a big stinky
    // and adds in its own derivation constructs that cause type conflicts.
    derived_keys: Option<Vec<[u8; 32]>>,

    // NOTE: This is an opaque type that can be deserialized inside an Operation
    blob: Option<serde_json::Value>,

    #[serde(flatten)]
    operation: Operation,
}

#[derive(Serialize, Deserialize, Debug)]
#[serde(tag = "operation", content = "values")]
pub enum Operation {
    GetBlockhash(GetBlockhash),
    GenerateWallet(GenerateWallet),
    GetWalletAddress(GetWalletAddress),
    GetTokenInfo(GetTokenInfo),
    Transfer(Transfer),
    CreateTokenAccount(CreateTokenAccount),
    TransferToken(TransferToken),
    Sign(Sign),
    Broadcast(Broadcast),
}

pub struct Solana;

impl Solana {
    fn keypair_from_bytes(given_bytes: [u8; 32]) -> solana_sdk::signer::keypair::Keypair {
        use ed25519_dalek::{PublicKey, SecretKey};
        let secret_key = SecretKey::from_bytes(&given_bytes).expect("key should be 32 bytes");
        let mut bytes = [0u8; 64];
        bytes[..32].clone_from_slice(&given_bytes);
        bytes[32..].clone_from_slice(PublicKey::from(&secret_key).as_bytes());
        solana_sdk::signer::keypair::Keypair::from_bytes(&bytes)
            .expect("solana sdk should expect 64 bytes")
    }
}

impl Module for Solana {
    type Error = Error;

    type Request = Request;

    fn describe_operations() -> Vec<icepick_module::help::Operation> {
        let cluster = Argument {
            name: "cluster".to_string(),
            description: "The cluster to interact with (mainnet, testnet, devnet)".to_string(),
            r#type: ArgumentType::Optional,
        };
        let account = Argument {
            name: "from-account".to_string(),
            description: "The derivation account used for the transaction.".to_string(),
            r#type: ArgumentType::Optional,
        };
        let fee = Argument {
            name: "fee".to_string(),
            description: "A custom fee for the transaction".to_string(),
            r#type: ArgumentType::Optional,
        };
        let fee_payer_address = Argument {
            name: "fee-payer-address".to_string(),
            description: "The address used to pay the fee.".to_string(),
            r#type: ArgumentType::Optional,
        };
        let fee_payer = Argument {
            name: "fee-payer".to_string(),
            description: "The derivation account used to pay the fee.".to_string(),
            r#type: ArgumentType::Optional,
        };
        let from_address = Argument {
            name: "from-address".to_string(),
            description: concat!(
                "The address to send SOL from; will be used to verify ",
                "the derivation account."
            )
            .to_string(),
            r#type: ArgumentType::Required,
        };
        vec![
            icepick_module::help::Operation {
                name: "get-blockhash".to_string(),
                description: "Get the latest blockhash".to_string(),
                arguments: vec![cluster.clone()],
            },
            icepick_module::help::Operation {
                name: "generate-wallet".to_string(),
                description: "Generate the derivation index for a wallet.".to_string(),
                arguments: vec![Argument {
                    name: "account".to_string(),
                    description: "The derivation account used for generating the wallet."
                        .to_string(),
                    r#type: ArgumentType::Optional,
                }],
            },
            icepick_module::help::Operation {
                name: "get-wallet-address".to_string(),
                description: "Get the address for a given wallet.".to_string(),
                arguments: vec![],
            },
            icepick_module::help::Operation {
                name: "get-token-info".to_string(),
                description: "Get the address for a given token.".to_string(),
                arguments: vec![Argument {
                    name: "token".to_string(),
                    description: "The token to look up".to_string(),
                    r#type: ArgumentType::Required,
                }],
            },
            icepick_module::help::Operation {
                name: "transfer".to_string(),
                description: "Transfer SOL from a Keyfork wallet to an external wallet."
                    .to_string(),
                arguments: vec![
                    Argument {
                        name: "amount".to_string(),
                        description: "The amount of SOL to transfer.".to_string(),
                        r#type: ArgumentType::Required,
                    },
                    account.clone(),
                    Argument {
                        name: "to-address".to_string(),
                        description: "The address to send SOL to.".to_string(),
                        r#type: ArgumentType::Required,
                    },
                    from_address.clone(),
                    fee.clone(),
                    fee_payer.clone(),
                    fee_payer_address.clone(),
                ],
            },
            icepick_module::help::Operation {
                name: "stake".to_string(),
                description: "Stake SOL to earn rewards.".to_string(),
                arguments: vec![
                    Argument {
                        name: "amount".to_string(),
                        description: "The amount of SOL to stake.".to_string(),
                        r#type: ArgumentType::Required,
                    },
                    account.clone(),
                    from_address.clone(),
                    fee.clone(),
                    fee_payer.clone(),
                    fee_payer_address.clone(),
                ],
            },
            // kinda BS that you have to make an account for a token, but ok.
            icepick_module::help::Operation {
                name: "create-token-account".to_string(),
                description: "Create an account for a given token".to_string(),
                arguments: vec![
                    Argument {
                        name: "wallet-address".to_string(),
                        description: "The address of the token.".to_string(),
                        r#type: ArgumentType::Required,
                    },
                    Argument {
                        name: "token-address".to_string(),
                        description: "The address of the token.".to_string(),
                        r#type: ArgumentType::Required,
                    },
                    Argument {
                        name: "funder-address".to_string(),
                        description: "The address of the funder (signer).".to_string(),
                        r#type: ArgumentType::Optional,
                    },
                ],
            },
            icepick_module::help::Operation {
                name: "transfer-token".to_string(),
                description: "Transfer tokens from a Keyfork wallet to an external wallet."
                    .to_string(),
                arguments: vec![
                    Argument {
                        name: "amount".to_string(),
                        description: "The amount of tokens to transfer.".to_string(),
                        r#type: ArgumentType::Required,
                    },
                    Argument {
                        name: "token-address".to_string(),
                        description: "The address of the token.".to_string(),
                        r#type: ArgumentType::Required,
                    },
                    account.clone(),
                    Argument {
                        name: "to-address".to_string(),
                        description: "The address to send the tokens to.".to_string(),
                        r#type: ArgumentType::Required,
                    },
                    Argument {
                        name: "from-address".to_string(),
                        description: "The address to send the tokens from; will be used to verify the derivation account.".to_string(),
                        r#type: ArgumentType::Required,
                    },
                    Argument {
                        name: "decimals".to_string(),
                        description: "The decimals of the token.".to_string(),
                        r#type: ArgumentType::Required,
                    },
                    fee.clone(),
                    fee_payer.clone(),
                    fee_payer_address.clone(),
                ],
            },
            icepick_module::help::Operation {
                name: "sign".to_string(),
                description: "Sign a previously-generated transaction.".to_string(),
                arguments: vec![Argument {
            name: "blockhash".to_string(),
            description: "A recent blockhash".to_string(),
            r#type: ArgumentType::Required,
        }],
            },
            icepick_module::help::Operation {
                name: "broadcast".to_string(),
                description: "Broadcast a signed transaction".to_string(),
                arguments: vec![cluster.clone()],
            },
        ]
    }

    fn handle_request(request: Self::Request) -> Result<serde_json::Value, Self::Error> {
        match request.operation {
            Operation::GetBlockhash(GetBlockhash { cluster }) => {
                let cluster = cluster.unwrap_or(Cluster::MainnetBeta);
                let cluster_url = format!("https://api.{cluster}.solana.com");
                let client = solana_rpc_client::rpc_client::RpcClient::new(cluster_url);
                let response = client.get_latest_blockhash().unwrap();
                Ok(serde_json::json!({
                    "blob": {
                        "blockhash": response.to_string(),
                    },
                }))
            }
            Operation::GenerateWallet(GenerateWallet { account }) => {
                let account = u32::from_str(account.as_deref().unwrap_or("0")).unwrap();
                Ok(serde_json::json!({
                    "blob": {},
                    "derivation_accounts": [(account | 1 << 31)],
                }))
            }
            Operation::GetWalletAddress(_) => {
                // NOTE: panics if doesn't exist
                let key = request.derived_keys.unwrap()[0];
                let keypair = Self::keypair_from_bytes(key);
                let pubkey = keypair.pubkey();
                Ok(serde_json::json!({
                    "blob": {
                        "pubkey": pubkey.to_string(),
                    }
                }))
            }
            Operation::GetTokenInfo(GetTokenInfo { token }) => {
                let values = match token.as_str() {
                    // Only exists on devnet
                    "IPDBG" => Some(("3V6hm5ifSLSWLZ86NpTxo5iVguGq9qCUtry6bn5PtT23", 9u8)),
                    // Only exists on mainnet
                    "PYTH" => Some(("HZ1JovNiVvGrGNiiYvEozEVgZ58xaU3RKwX8eACQBCt3", 6u8)),
                    _ => None,
                };

                Ok(match values {
                    Some((address, decimals)) => serde_json::json!({
                        "blob": {
                            "token_address": address,
                            "token_decimals": decimals,
                        }
                    }),
                    None => serde_json::json!({
                        "blob": {},
                        "error": "key was not found!",
                    }),
                })
            }
            Operation::Transfer(Transfer {
                amount,
                from_account,
                to_address,
                from_address,
                fee: _,
                fee_payer,
                fee_payer_address,
            }) => {
                // TODO:
                // parse address for to_address

                let amount = f64::from_str(&amount).expect("float amount");
                let amount: u64 = (amount * LAMPORTS_PER_SOL as f64) as u64;

                use solana_sdk::pubkey::Pubkey;
                let to_pk = Pubkey::from_str(&to_address).unwrap();
                let from_pk = Pubkey::from_str(&from_address).unwrap();
                let payer_account_and_pk = {
                    // If a fee payer is given, a fee payer address must also be given, since the
                    // address must be known before signing the transaction.
                    match (&fee_payer, &fee_payer_address) {
                        (Some(payer), Some(address)) => {
                            // Use the provided account
                            Some((
                                u32::from_str(payer).unwrap(),
                                Pubkey::from_str(address).unwrap(),
                            ))
                        }
                        (None, None) => {
                            // Use the transaction account
                            None
                        }
                        _ => panic!("Invalid combination of fee_payer and fee_payer_address"),
                    }
                };
                let instruction =
                    solana_sdk::system_instruction::transfer(&from_pk, &to_pk, amount);
                let message = solana_sdk::message::Message::new(
                    &[instruction],
                    payer_account_and_pk.map(|v| v.1).as_ref(),
                );
                let transaction = solana_sdk::transaction::Transaction::new_unsigned(message);
                // TODO: error handling from_str
                let from_account = from_account
                    .and_then(|a| u32::from_str(&a).ok())
                    .unwrap_or(0);
                let mut requested_accounts = vec![];
                requested_accounts.push(from_account | 1 << 31);
                if let Some((account, _)) = &payer_account_and_pk {
                    requested_accounts.push(*account | 1 << 31);
                }
                Ok(serde_json::json!({
                    "blob": {
                        "transaction": transaction,
                    },
                    "derivation_accounts": requested_accounts,
                }))
            }
            Operation::CreateTokenAccount(CreateTokenAccount {
                funder_address,
                wallet_address,
                token_address,
            }) => {
                // TODO: allow changing derivation account of funder_address
                use sata::instruction::create_associated_token_account;
                use solana_sdk::pubkey::Pubkey;
                use spl_associated_token_account as sata;
                use spl_token::ID as TOKEN_ID;
                let funder_address = funder_address.unwrap_or_else(|| wallet_address.clone());
                let funder_pubkey = Pubkey::from_str(&funder_address).unwrap();
                let wallet_pubkey = Pubkey::from_str(&wallet_address).unwrap();
                let token_pubkey = Pubkey::from_str(&token_address).unwrap();
                let instruction = create_associated_token_account(
                    &funder_pubkey,
                    &wallet_pubkey,
                    &token_pubkey,
                    &TOKEN_ID,
                );
                let message =
                    solana_sdk::message::Message::new(&[instruction], Some(&funder_pubkey));
                let transaction = solana_sdk::transaction::Transaction::new_unsigned(message);
                #[allow(clippy::identity_op)]
                Ok(serde_json::json!({
                    "blob": {
                        "transaction": transaction,
                    },
                    "derivation_accounts": [0u32 | 1 << 31],
                }))
            }
            Operation::TransferToken(TransferToken {
                amount,
                token_address,
                to_address,
                from_account,
                from_address,
                decimals,
                fee,
                fee_payer,
                fee_payer_address,
            }) => {
                // TODO: deduplicate code used in Transfer

                let amount = f64::from_str(&amount).expect("float amount");
                let decimals = u8::from_str(&decimals).expect("decimals");
                let amount: u64 = (amount * 10u64.pow(decimals as u32) as f64) as u64;

                use solana_sdk::pubkey::Pubkey;
                use spl_associated_token_account::get_associated_token_address;
                let to_pk = Pubkey::from_str(&to_address).unwrap();
                let from_pk = Pubkey::from_str(&from_address).unwrap();
                let token_pk = Pubkey::from_str(&token_address).unwrap();
                let payer_account_and_pk = {
                    // If a fee payer is given, a fee payer address must also be given, since the
                    // address must be known before signing the transaction.
                    match (&fee_payer, &fee_payer_address) {
                        (Some(payer), Some(address)) => {
                            // Use the provided account
                            Some((
                                u32::from_str(payer).unwrap(),
                                Pubkey::from_str(address).unwrap(),
                            ))
                        }
                        (None, None) => {
                            // Use the transaction account
                            None
                        }
                        _ => panic!("Invalid combination of fee_payer and fee_payer_address"),
                    }
                };
                let token_program_id = spl_token::ID;
                let mut signers = vec![&from_pk];
                if let Some((_, pk)) = payer_account_and_pk.as_ref() {
                    signers.push(pk);
                }

                let from_token_address = get_associated_token_address(&from_pk, &token_pk);
                let to_token_address = get_associated_token_address(&to_pk, &token_pk);
                let mut instruction = spl_token_2022::instruction::transfer_checked(
                    &token_program_id,   // token program id
                    &from_token_address, // source, as token address
                    &token_pk,           // mint
                    &to_token_address,   // destination, as token address
                    &from_pk,            // authority, as source sol address
                    // TODO: signers should be [] when not using multisig
                    // but should contain all signers when multisig
                    &[],      // signers
                    amount,   // amount
                    decimals, // decimals
                )
                .unwrap();
                // TODO: check if this works with payer
                // this is required because the Solana SDK does not set the primary transactional
                // key as writable (the one that would be paying computation fees) in the event a
                // payer is not provided. The transactional account must be writable for the
                // computation fee to be paid.
                if payer_account_and_pk.is_none() {
                    for account in instruction.accounts.iter_mut() {
                        if account.pubkey == from_pk {
                            account.is_writable = true;
                        }
                    }
                }
                let message = solana_sdk::message::Message::new(
                    &[instruction],
                    payer_account_and_pk.map(|v| v.1).as_ref(),
                );
                // message.header.num_readonly_signed_accounts = 0;
                let transaction =
                    solana_sdk::transaction::Transaction::new_unsigned(message.clone());
                /*
                use base64::prelude::*;
                eprintln!("{}", BASE64_STANDARD.encode(transaction.message_data()));
                */

                #[allow(clippy::identity_op)]
                Ok(serde_json::json!({
                    "blob": {
                        "transaction": transaction,
                    },
                    "derivation_accounts": [0u32 | 1 << 31],
                }))
            }
            Operation::Sign(Sign { blockhash }) => {
                let transaction = request
                    .blob
                    .and_then(|b| b.get("transaction").cloned())
                    .expect("was given transaction");
                let mut transaction: solana_sdk::transaction::Transaction =
                    serde_json::from_value(transaction).expect("valid message blob");
                let keys = request
                    .derived_keys
                    .unwrap_or_default()
                    .iter()
                    .map(|k| Self::keypair_from_bytes(*k))
                    .collect::<Vec<_>>();

                let hash = solana_sdk::hash::Hash::from_str(&blockhash).unwrap();
                transaction
                    .try_sign(&keys, hash)
                    .expect("not enough keys provided");
                Ok(serde_json::json!({
                    "blob": {
                        "transaction": transaction,
                    }
                }))
            }
            Operation::Broadcast(Broadcast { cluster }) => {
                let cluster = cluster.unwrap_or(Cluster::MainnetBeta);
                let cluster_url = format!("https://api.{cluster}.solana.com");

                let transaction = request
                    .blob
                    .and_then(|b| b.get("transaction").cloned())
                    .expect("was given transaction");
                let transaction: solana_sdk::transaction::Transaction =
                    serde_json::from_value(transaction).expect("valid message blob");
                transaction.verify().expect("invalid signatures");
                let client = solana_rpc_client::rpc_client::RpcClient::new(cluster_url);
                let _simulated_response = client.simulate_transaction(&transaction).unwrap();
                let response = client.send_and_confirm_transaction(&transaction);
                Ok(match response {
                    Ok(s) => {
                        serde_json::json!({
                            "blob": {
                                "status": "send_and_confirm",
                                "succcess": s.to_string(),
                            }
                        })
                    }
                    Err(e) => {
                        serde_json::json!({
                            "blob": {
                                "status": "send_and_confirm",
                                "error": e.to_string(),
                            }
                        })
                    }
                })
            }
        }
    }
}