use crate::error::{PyraSolanaError, Result, RpcAttemptError};
use futures::stream::{FuturesUnordered, StreamExt};
use solana_account::Account;
use solana_account_decoder_client_types::token::UiTokenAmount;
use solana_address_lookup_table_interface::state::AddressLookupTable;
use solana_client::{
    nonblocking::rpc_client::RpcClient,
    rpc_config::{RpcSendTransactionConfig, RpcSimulateTransactionConfig},
    rpc_request::TokenAccountsFilter,
    rpc_response::{RpcKeyedAccount, RpcPrioritizationFee, RpcSimulateTransactionResult},
};
use solana_commitment_config::CommitmentConfig;
use solana_hash::Hash;
use solana_message::AddressLookupTableAccount;
use solana_pubkey::Pubkey;
use solana_transaction::versioned::VersionedTransaction;
use std::sync::Arc;
use std::time::Duration;

/// Per-RPC timeout for fallback reads. If a single RPC is slow but keeps the
/// connection open, we move on to the next one rather than blocking indefinitely.
const FALLBACK_RPC_TIMEOUT: Duration = Duration::from_secs(5);

/// Transaction confirmation status.
#[derive(Debug, Clone)]
pub struct SignatureStatus {
    pub confirmed: bool,
    pub finalized: bool,
    pub err: Option<String>,
}

/// Multi-endpoint Solana RPC client with failover and Pyra ALT caching.
///
/// Provides two failover strategies for any RPC operation:
/// - [`shotgun_send`](Self::shotgun_send): fires at all RPCs concurrently, returns first success.
///   Use for transaction sends (safe because Solana txs are idempotent).
/// - [`fallback_read`](Self::fallback_read): tries RPCs sequentially with a per-RPC timeout.
///   Use for reads and queries.
///
/// Both strategies are public — call them directly with a closure to use any
/// `RpcClient` method without needing a wrapper in this crate. For one-off
/// calls that don't need failover, use [`primary_rpc`](Self::primary_rpc).
///
/// Fetches and caches the Pyra address lookup table on creation so that
/// `build_v0_transaction` doesn't need an RPC round-trip per call.
#[derive(Clone)]
pub struct SolanaRpcClient {
    rpcs: Vec<Arc<RpcClient>>,
    pyra_alt: Option<AddressLookupTableAccount>,
}

impl SolanaRpcClient {
    /// Create a client and fetch the Pyra ALT for caching.
    pub async fn new(urls: &[String]) -> Result<Self> {
        let mut client = Self::new_without_alt_cache(urls)?;

        let alt_pubkey = pyra_instructions::constants::PYRA_ADDRESS_TABLE;
        let alt = client.get_address_lookup_table(&alt_pubkey).await?;
        client.pyra_alt = Some(alt);

        Ok(client)
    }

    /// Create a client without fetching the Pyra ALT.
    ///
    /// Use this when you don't need `build_v0_transaction` or when
    /// a live RPC is not available at construction time.
    pub fn new_without_alt_cache(urls: &[String]) -> Result<Self> {
        if urls.is_empty() {
            return Err(PyraSolanaError::Rpc(
                "At least one Solana RPC URL is required".into(),
            ));
        }

        let rpcs: Vec<Arc<RpcClient>> = urls
            .iter()
            .map(|url| {
                Arc::new(RpcClient::new_with_commitment(
                    url.clone(),
                    CommitmentConfig::confirmed(),
                ))
            })
            .collect();

        Ok(Self {
            rpcs,
            pyra_alt: None,
        })
    }

    /// Returns the cached Pyra address lookup table, if fetched at creation.
    pub fn pyra_address_lookup_table(&self) -> Option<&AddressLookupTableAccount> {
        self.pyra_alt.as_ref()
    }

    /// Returns the primary (first) underlying `RpcClient`.
    ///
    /// Use this for one-off calls that don't need failover, or when you need
    /// access to `RpcClient` methods not wrapped by this crate.
    pub fn primary_rpc(&self) -> &RpcClient {
        // Safe: constructor rejects empty URL lists.
        &self.rpcs[0]
    }

    /// Returns all underlying `RpcClient` instances.
    pub fn all_rpcs(&self) -> &[Arc<RpcClient>] {
        &self.rpcs
    }

    /// Fire `op` against ALL RPCs concurrently, return the first `Ok`.
    /// If all fail, return the last error.
    ///
    /// Use for idempotent operations like sending transactions.
    ///
    /// ```rust,ignore
    /// let sig = rpc.shotgun_send(|client| {
    ///     let tx = tx.clone();
    ///     async move {
    ///         client.send_transaction(&tx).await
    ///             .map_err(|e| PyraSolanaError::Rpc(e.to_string()))
    ///     }
    /// }).await?;
    /// ```
    pub async fn shotgun_send<F, Fut, T>(&self, op: F) -> Result<T>
    where
        F: Fn(Arc<RpcClient>) -> Fut,
        Fut: std::future::Future<Output = Result<T>> + Send + 'static,
        T: Send + 'static,
    {
        let mut futs: FuturesUnordered<_> = self
            .rpcs
            .iter()
            .enumerate()
            .map(|(i, rpc)| {
                let fut = op(Arc::clone(rpc));
                async move { (i, fut.await) }
            })
            .collect();

        let mut attempts = Vec::new();
        while let Some((i, result)) = futs.next().await {
            match result {
                Ok(value) => return Ok(value),
                Err(e) => attempts.push(RpcAttemptError {
                    endpoint_index: i,
                    error: e.to_string(),
                }),
            }
        }

        Err(PyraSolanaError::AllEndpointsFailed { attempts })
    }

    /// Try `op` against RPCs in order, return on first `Ok`.
    /// Each attempt is bounded by [`FALLBACK_RPC_TIMEOUT`] so a slow-but-connected
    /// RPC doesn't block the entire chain.
    ///
    /// Use for read operations and queries.
    ///
    /// ```rust,ignore
    /// let balance = rpc.fallback_read(|client| {
    ///     let addr = token_account;
    ///     async move {
    ///         client.get_token_account_balance(&addr).await
    ///             .map_err(|e| PyraSolanaError::Rpc(e.to_string()))
    ///     }
    /// }).await?;
    /// ```
    pub async fn fallback_read<F, Fut, T>(&self, op: F) -> Result<T>
    where
        F: Fn(Arc<RpcClient>) -> Fut,
        Fut: std::future::Future<Output = Result<T>>,
    {
        let mut attempts = Vec::new();
        for (i, rpc) in self.rpcs.iter().enumerate() {
            let result = tokio::time::timeout(FALLBACK_RPC_TIMEOUT, op(Arc::clone(rpc))).await;
            match result {
                Ok(Ok(value)) => return Ok(value),
                Ok(Err(e)) => {
                    attempts.push(RpcAttemptError {
                        endpoint_index: i,
                        error: e.to_string(),
                    });
                }
                Err(_elapsed) => {
                    attempts.push(RpcAttemptError {
                        endpoint_index: i,
                        error: format!("timed out after {}s", FALLBACK_RPC_TIMEOUT.as_secs()),
                    });
                }
            }
        }
        Err(PyraSolanaError::AllEndpointsFailed { attempts })
    }

    pub async fn get_latest_blockhash(&self) -> Result<Hash> {
        self.fallback_read(|rpc| async move {
            rpc.get_latest_blockhash()
                .await
                .map_err(|e| PyraSolanaError::Rpc(format!("get_latest_blockhash: {e}")))
        })
        .await
    }

    pub async fn simulate_transaction(
        &self,
        tx: &VersionedTransaction,
    ) -> Result<RpcSimulateTransactionResult> {
        let tx = tx.clone();
        self.fallback_read(|rpc| {
            let tx = tx.clone();
            async move {
                let config = RpcSimulateTransactionConfig {
                    sig_verify: false,
                    replace_recent_blockhash: true,
                    commitment: Some(CommitmentConfig::confirmed()),
                    ..Default::default()
                };
                rpc.simulate_transaction_with_config(&tx, config)
                    .await
                    .map(|r| r.value)
                    .map_err(|e| PyraSolanaError::Rpc(format!("simulate_transaction: {e}")))
            }
        })
        .await
    }

    pub async fn get_recent_prioritization_fees(
        &self,
        accounts: &[Pubkey],
    ) -> Result<Vec<RpcPrioritizationFee>> {
        let accounts = accounts.to_vec();
        self.fallback_read(|rpc| {
            let accounts = accounts.clone();
            async move {
                rpc.get_recent_prioritization_fees(&accounts)
                    .await
                    .map_err(|e| {
                        PyraSolanaError::Rpc(format!("get_recent_prioritization_fees: {e}"))
                    })
            }
        })
        .await
    }

    pub async fn send_transaction(
        &self,
        tx: &VersionedTransaction,
    ) -> Result<solana_signature::Signature> {
        let tx = tx.clone();
        self.shotgun_send(|rpc| {
            let tx = tx.clone();
            async move {
                rpc.send_transaction(&tx)
                    .await
                    .map_err(|e| PyraSolanaError::Rpc(format!("send_transaction: {e}")))
            }
        })
        .await
    }

    pub async fn send_transaction_with_config(
        &self,
        tx: &VersionedTransaction,
        skip_preflight: bool,
    ) -> Result<solana_signature::Signature> {
        let tx = tx.clone();
        self.shotgun_send(|rpc| {
            let tx = tx.clone();
            async move {
                let config = RpcSendTransactionConfig {
                    skip_preflight,
                    ..Default::default()
                };
                rpc.send_transaction_with_config(&tx, config)
                    .await
                    .map_err(|e| PyraSolanaError::Rpc(format!("send_transaction: {e}")))
            }
        })
        .await
    }

    pub async fn get_account_data(&self, pubkey: &Pubkey) -> Result<Option<Vec<u8>>> {
        let pubkey = *pubkey;
        self.fallback_read(|rpc| async move {
            match rpc.get_account(&pubkey).await {
                Ok(account) => Ok(Some(account.data)),
                Err(e) => {
                    let err_str = e.to_string();
                    if err_str.contains("AccountNotFound")
                        || err_str.contains("could not find account")
                    {
                        Ok(None)
                    } else {
                        Err(PyraSolanaError::Rpc(format!(
                            "get_account_data for {pubkey}: {e}"
                        )))
                    }
                }
            }
        })
        .await
    }

    pub async fn get_address_lookup_table(
        &self,
        alt_pubkey: &Pubkey,
    ) -> Result<AddressLookupTableAccount> {
        let alt_pubkey = *alt_pubkey;
        self.fallback_read(|rpc| async move {
            let account = rpc.get_account(&alt_pubkey).await.map_err(|e| {
                PyraSolanaError::Rpc(format!("get_account for ALT {alt_pubkey}: {e}"))
            })?;

            let lookup_table = AddressLookupTable::deserialize(&account.data).map_err(|e| {
                PyraSolanaError::Rpc(format!("Failed to deserialize ALT {alt_pubkey}: {e}"))
            })?;

            Ok(AddressLookupTableAccount {
                key: alt_pubkey,
                addresses: lookup_table.addresses.to_vec(),
            })
        })
        .await
    }

    pub async fn get_signature_statuses(
        &self,
        signatures: &[solana_signature::Signature],
    ) -> Result<Vec<Option<SignatureStatus>>> {
        let signatures = signatures.to_vec();
        self.fallback_read(|rpc| {
            let signatures = signatures.clone();
            async move {
                let response = rpc
                    .get_signature_statuses(&signatures)
                    .await
                    .map_err(|e| PyraSolanaError::Rpc(format!("get_signature_statuses: {e}")))?;

                Ok(response
                    .value
                    .into_iter()
                    .map(|opt| {
                        opt.map(|status| {
                            let confirmed =
                                status.satisfies_commitment(CommitmentConfig::confirmed());
                            let finalized =
                                status.satisfies_commitment(CommitmentConfig::finalized());
                            let err = status.err.map(|e| format!("{e}"));
                            SignatureStatus {
                                confirmed,
                                finalized,
                                err,
                            }
                        })
                    })
                    .collect())
            }
        })
        .await
    }

    pub async fn get_slot(&self) -> Result<u64> {
        self.fallback_read(|rpc| async move {
            rpc.get_slot()
                .await
                .map_err(|e| PyraSolanaError::Rpc(format!("get_slot: {e}")))
        })
        .await
    }

    pub async fn get_balance(&self, pubkey: &Pubkey) -> Result<u64> {
        let pubkey = *pubkey;
        self.fallback_read(|rpc| async move {
            rpc.get_balance(&pubkey)
                .await
                .map_err(|e| PyraSolanaError::Rpc(format!("get_balance for {pubkey}: {e}")))
        })
        .await
    }

    pub async fn get_minimum_balance_for_rent_exemption(&self, data_len: usize) -> Result<u64> {
        self.fallback_read(|rpc| async move {
            rpc.get_minimum_balance_for_rent_exemption(data_len)
                .await
                .map_err(|e| {
                    PyraSolanaError::Rpc(format!(
                        "get_minimum_balance_for_rent_exemption({data_len}): {e}"
                    ))
                })
        })
        .await
    }

    pub async fn get_token_account_balance(&self, pubkey: &Pubkey) -> Result<UiTokenAmount> {
        let pubkey = *pubkey;
        self.fallback_read(|rpc| async move {
            rpc.get_token_account_balance(&pubkey).await.map_err(|e| {
                PyraSolanaError::Rpc(format!("get_token_account_balance for {pubkey}: {e}"))
            })
        })
        .await
    }

    pub async fn get_multiple_accounts(&self, pubkeys: &[Pubkey]) -> Result<Vec<Option<Account>>> {
        let pubkeys = pubkeys.to_vec();
        self.fallback_read(|rpc| {
            let pubkeys = pubkeys.clone();
            async move {
                rpc.get_multiple_accounts(&pubkeys)
                    .await
                    .map_err(|e| PyraSolanaError::Rpc(format!("get_multiple_accounts: {e}")))
            }
        })
        .await
    }

    /// Poll signature status until confirmed (or finalized) or timeout.
    ///
    /// Returns `Ok(SignatureStatus)` once the transaction reaches the requested
    /// commitment level. Returns an error if the transaction fails on-chain
    /// or the timeout is exceeded.
    pub async fn confirm_transaction(
        &self,
        signature: &solana_signature::Signature,
        finalized: bool,
        timeout: Duration,
    ) -> Result<SignatureStatus> {
        let start = tokio::time::Instant::now();
        let poll_interval = Duration::from_millis(500);

        loop {
            let statuses = self.get_signature_statuses(&[*signature]).await?;

            if let Some(Some(status)) = statuses.first() {
                if let Some(ref err) = status.err {
                    return Err(PyraSolanaError::Transaction(format!(
                        "Transaction {signature} failed: {err}"
                    )));
                }

                let reached = if finalized {
                    status.finalized
                } else {
                    status.confirmed
                };

                if reached {
                    return Ok(status.clone());
                }
            }

            if start.elapsed() >= timeout {
                let level = if finalized { "finalized" } else { "confirmed" };
                return Err(PyraSolanaError::Rpc(format!(
                    "Transaction {signature} not {level} after {}s",
                    timeout.as_secs()
                )));
            }

            tokio::time::sleep(poll_interval).await;
        }
    }

    pub async fn get_token_accounts_by_owner(
        &self,
        owner: &Pubkey,
        filter: TokenAccountsFilter,
    ) -> Result<Vec<RpcKeyedAccount>> {
        let owner = *owner;
        self.fallback_read(|rpc| {
            let filter = match &filter {
                TokenAccountsFilter::Mint(mint) => TokenAccountsFilter::Mint(*mint),
                TokenAccountsFilter::ProgramId(id) => TokenAccountsFilter::ProgramId(*id),
            };
            async move {
                rpc.get_token_accounts_by_owner(&owner, filter)
                    .await
                    .map_err(|e| {
                        PyraSolanaError::Rpc(format!(
                            "get_token_accounts_by_owner for {owner}: {e}"
                        ))
                    })
            }
        })
        .await
    }
}
