use anyhow::{anyhow, Result};
use itertools::Itertools;
use solana_client::{
    rpc_client::RpcClient,
    rpc_config::{RpcSendTransactionConfig, RpcSimulateTransactionConfig},
};
use solana_compute_budget_interface::ComputeBudgetInstruction;
use solana_hash::Hash;
use solana_instruction::Instruction;
use solana_keypair::{Keypair, Signature};
use solana_message::{v0::Message, VersionedMessage};
use solana_pubkey::Pubkey;
use solana_signer::Signer;
use solana_transaction::versioned::VersionedTransaction;
use solana_transaction_status_client_types::TransactionStatus;
use std::{
    cmp::min,
    collections::HashMap,
    sync::OnceLock,
    thread::sleep,
    time::{Duration, Instant},
};

use crate::{debug, log::Ansi};

const DEFAULT_COMPUTE_UNIT_PRICE: u64 = 1000;
const MAX_COMPUTE_UNITS: u32 = 1400000;
const MIN_COMPUTE_UNITS_MARGIN: u32 = 10000;
const MAX_PRIORITY_FEE_MICRO_LAMPORTS: u64 = 500000000; // 50 lamports per compute unit
const SEND_TX_INTERVAL: Duration = Duration::from_secs(1);

pub static SKIP_SIMULATION: OnceLock<bool> = OnceLock::new();

pub fn process_instructions(
    rpc: &RpcClient,
    signers: &[Keypair],
    instructions: &[Instruction],
) -> Result<()> {
    let fee_payer = signers.first().map(|x| x.pubkey()).unwrap_or_default();
    let message = build_transaction(rpc, &fee_payer, instructions)?;
    process_transaction(rpc, signers, message)
}

pub fn process_transaction(
    rpc: &RpcClient,
    signers: &[Keypair],
    message: VersionedMessage,
) -> Result<()> {
    if !signers.is_empty() {
        let signature = send_transaction(rpc, signers, message)?;
        println!("{}", "Transaction sent".bold());
        println!(
            "{}",
            signature.link(&format!("https://solscan.io/tx/{}", signature))
        );
    } else {
        println!("{}", "Exporting transaction to be signed later".bold());
        let tx = export_transaction(rpc, &message)?;
        println!("{}", tx);
    }
    Ok(())
}

pub fn send_transaction(
    rpc: &RpcClient,
    signers: &[Keypair],
    message: VersionedMessage,
) -> Result<String> {
    let result = send_transactions(rpc, signers, vec![message]);
    let first_result = result
        .into_iter()
        .next()
        .ok_or(anyhow!("No transaction result returned"))?;

    first_result
}

pub fn send_transactions(
    rpc: &RpcClient,
    signers: &[Keypair],
    messages: Vec<VersionedMessage>,
) -> Vec<Result<String>> {
    let transactions: Vec<VersionedTransaction> = messages
        .into_iter()
        .map(|x| VersionedTransaction {
            signatures: signers
                .iter()
                .map(|s| s.sign_message(&x.serialize()))
                .collect(),
            message: x,
        })
        .collect();

    let mut transactions_to_send: HashMap<Signature, &VersionedTransaction> =
        transactions.iter().map(|x| (x.signatures[0], x)).collect();

    let mut transaction_results: HashMap<Signature, Result<String>> = HashMap::new();

    if !*SKIP_SIMULATION.get().unwrap_or(&false) {
        for (signature, transaction) in transactions_to_send.clone() {
            let result = rpc.simulate_transaction_with_config(
                transaction,
                RpcSimulateTransactionConfig {
                    sig_verify: false,
                    replace_recent_blockhash: true,
                    ..Default::default()
                },
            );

            match result {
                Ok(result) => {
                    if let Some(err) = result.value.err {
                        debug!(
                            "Simulation failed:\n{}",
                            result.value.logs.unwrap_or_default().join("\n")
                        );
                        transaction_results.insert(signature, Err(anyhow!(err)));
                        transactions_to_send.remove(&signature);
                    }
                }
                Err(e) => {
                    transaction_results.insert(signature, Err(anyhow!(e)));
                    transactions_to_send.remove(&signature);
                }
            }
        }
    }

    // Retry sending transactions to cluster until either confirmed or expired (up to 90 seconds)
    let start_time = Instant::now();
    let mut next_tick = Instant::now() + SEND_TX_INTERVAL;
    while start_time.elapsed().as_secs() < 90 && !transactions_to_send.is_empty() {
        for transaction in transactions_to_send.values() {
            let send_result = rpc.send_transaction_with_config(
                *transaction,
                RpcSendTransactionConfig {
                    skip_preflight: true,
                    ..RpcSendTransactionConfig::default()
                },
            );
            if let Err(e) = send_result {
                debug!("Failed to send transaction: {}", e);
            }
        }

        let signatures = transactions_to_send.keys().cloned().collect();
        let statuses = get_signature_statuses_batched(rpc, signatures);

        match statuses {
            Ok(statuses) => {
                for (signature, status) in statuses {
                    if let Some(err) = status.err {
                        transaction_results.insert(signature, Err(anyhow!(err)));
                        transactions_to_send.remove(&signature);
                    } else {
                        transaction_results.insert(signature, Ok(signature.to_string()));
                        transactions_to_send.remove(&signature);
                    }
                }
            }
            Err(e) => {
                debug!("Failed to get signature statuses: {}", e);
            }
        }

        if Instant::now() < next_tick {
            sleep(next_tick - Instant::now());
        }
        next_tick += SEND_TX_INTERVAL;
    }

    let mut results: Vec<Result<String>> = Vec::new();

    for transaction in transactions {
        let signature = transaction.signatures[0];
        if let Some(status) = transaction_results.get(&signature) {
            match status {
                Ok(status) => results.push(Ok(status.clone())),
                Err(err) => results.push(Err(anyhow!("{}", err))),
            }
        } else {
            results.push(Err(anyhow!("Transaction expired {}", signature)));
        }
    }

    results
}

pub fn export_transaction(rpc: &RpcClient, message: &VersionedMessage) -> Result<String> {
    let tx = VersionedTransaction {
        signatures: vec![],
        message: message.clone(),
    };
    if !*SKIP_SIMULATION.get().unwrap_or(&false) {
        let result = rpc.simulate_transaction_with_config(
            &tx,
            RpcSimulateTransactionConfig {
                sig_verify: false,
                replace_recent_blockhash: true,
                ..Default::default()
            },
        )?;

        if let Some(err) = result.value.err {
            return Err(anyhow!("Failed to simulate transaction: {}", err));
        }
    }

    let serialized = tx.message.serialize();
    let result = bs58::encode(serialized).into_string();
    Ok(result)
}

pub fn build_transaction(
    rpc: &RpcClient,
    payer: &Pubkey,
    instructions: &[Instruction],
) -> Result<VersionedMessage> {
    let compute_unit_limit = get_compute_budget_limit(rpc, payer, instructions);
    let lock_writable_accounts = get_lock_writable_accounts(instructions);
    let compute_unit_price = get_compute_budget_price(rpc, &lock_writable_accounts);
    let blockhash = rpc.get_latest_blockhash()?;

    build_transaction_sync(
        &blockhash,
        payer,
        Some(compute_unit_limit),
        Some(compute_unit_price),
        None,
        instructions,
    )
}

pub fn build_transaction_sync(
    blockhash: &Hash,
    payer: &Pubkey,
    compute_unit_limit: Option<u32>,
    compute_unit_price: Option<u64>,
    loaded_accounts_data_size: Option<u32>,
    instructions: &[Instruction],
) -> Result<VersionedMessage> {
    debug!(
        "Building transaction with CU limit {:?}, CU price {:?}, accounts loaded data size {:?}",
        compute_unit_limit, compute_unit_price, loaded_accounts_data_size
    );

    let mut instructions = instructions.to_vec();
    if let Some(compute_unit_limit) = compute_unit_limit {
        instructions.push(ComputeBudgetInstruction::set_compute_unit_limit(
            compute_unit_limit,
        ));
    }
    if let Some(compute_unit_price) = compute_unit_price {
        instructions.push(ComputeBudgetInstruction::set_compute_unit_price(
            compute_unit_price,
        ));
    }
    if let Some(loaded_accounts_data_size) = loaded_accounts_data_size {
        instructions.push(
            ComputeBudgetInstruction::set_loaded_accounts_data_size_limit(
                loaded_accounts_data_size,
            ),
        );
    }

    let message =
        VersionedMessage::V0(Message::try_compile(payer, &instructions, &[], *blockhash)?);

    Ok(message)
}

pub fn get_compute_budget_limit(
    rpc: &RpcClient,
    payer: &Pubkey,
    instructions: &[Instruction],
) -> u32 {
    let message = Message::try_compile(payer, instructions, &[], Hash::new_unique());

    let mut compute_units = MAX_COMPUTE_UNITS;

    if let Ok(message) = message {
        let transaction = VersionedTransaction {
            signatures: vec![],
            message: VersionedMessage::V0(message),
        };

        let result = rpc.simulate_transaction_with_config(
            &transaction,
            RpcSimulateTransactionConfig {
                sig_verify: false,
                replace_recent_blockhash: true,
                commitment: None,
                encoding: None,
                accounts: None,
                min_context_slot: None,
                inner_instructions: false,
            },
        );

        if let Ok(result) = result {
            if let Some(units_consumed) = result.value.units_consumed {
                let est_compute_units = units_consumed as u32;
                let compute_units_margin = min(est_compute_units / 10, MIN_COMPUTE_UNITS_MARGIN);
                compute_units = est_compute_units + compute_units_margin;
            }
        }
    }

    compute_units
}

pub fn get_lock_writable_accounts(instructions: &[Instruction]) -> Vec<Pubkey> {
    instructions
        .iter()
        .flat_map(|instruction| instruction.accounts.iter())
        .filter(|account| account.is_writable)
        .map(|account| account.pubkey)
        .collect::<Vec<_>>()
}

pub fn get_compute_budget_price(rpc: &RpcClient, lock_writable_accounts: &[Pubkey]) -> u64 {
    let recent_prioritization_fees = rpc.get_recent_prioritization_fees(lock_writable_accounts);

    let mut compute_unit_price = DEFAULT_COMPUTE_UNIT_PRICE;
    if let Ok(recent_prioritization_fees) = recent_prioritization_fees {
        let sorted_fees = recent_prioritization_fees
            .iter()
            .sorted_by(|a, b| a.prioritization_fee.cmp(&b.prioritization_fee))
            .collect::<Vec<_>>();
        let median_fee = sorted_fees[sorted_fees.len() / 2].prioritization_fee;
        compute_unit_price = median_fee;
    }

    min(compute_unit_price, MAX_PRIORITY_FEE_MICRO_LAMPORTS)
}

fn get_signature_statuses_batched(
    rpc: &RpcClient,
    signatures: Vec<Signature>,
) -> Result<HashMap<Signature, TransactionStatus>> {
    let mut statuses: HashMap<Signature, TransactionStatus> = HashMap::new();
    for batch in signatures.chunks(256) {
        let statuses_batch = rpc.get_signature_statuses(batch)?;
        for (i, status) in statuses_batch.value.into_iter().enumerate() {
            if let Some(status) = status {
                statuses.insert(batch[i], status);
            }
        }
    }
    Ok(statuses)
}
