use crate::error::{PyraSolanaError, Result};
use crate::rpc_client::SolanaRpcClient;
use pyra_instructions::constants::PYRA_ADDRESS_TABLE;
use solana_compute_budget_interface::ComputeBudgetInstruction;
use solana_hash::Hash;
use solana_instruction::Instruction;
use solana_message::{v0, AddressLookupTableAccount, VersionedMessage};
use solana_pubkey::Pubkey;
use solana_signature::Signature;
use solana_transaction::versioned::VersionedTransaction;

/// Minimum compute units to request.
const MIN_COMPUTE_UNITS: u32 = 50_000;

/// Default compute units when simulation fails entirely.
/// Matches the fill-bot's fallback — high enough for most Drift operations.
const DEFAULT_COMPUTE_UNITS: u32 = 800_000;

/// Solana's maximum compute unit limit, used for simulation so that
/// programs exceeding the default 200k CU budget can still report
/// actual consumption.
const SIMULATION_COMPUTE_UNIT_LIMIT: u32 = 1_400_000;

/// Buffer multiplier applied to simulated compute unit consumption.
const COMPUTE_UNIT_BUFFER_MULTIPLIER: f64 = 1.5;

/// Base transaction fee in lamports (signature fee).
const BASE_FEE_LAMPORTS: u64 = 5_000;

/// The compute budget program ID.
const COMPUTE_BUDGET_PROGRAM_ID: Pubkey = solana_compute_budget_interface::ID;

/// Check if instructions already contain compute budget instructions.
fn has_compute_budget_instructions(instructions: &[Instruction]) -> bool {
    instructions
        .iter()
        .any(|ix| ix.program_id == COMPUTE_BUDGET_PROGRAM_ID)
}

/// Build a V0 versioned transaction with automatic CU budget, priority fee, and ALT resolution.
///
/// Prepends compute budget instructions if not already present, resolves the Pyra ALT,
/// and fetches the latest blockhash.
pub async fn build_v0_transaction(
    rpc: &SolanaRpcClient,
    instructions: &[Instruction],
    payer: &Pubkey,
) -> Result<VersionedTransaction> {
    let address_lookup_tables = match rpc.pyra_address_lookup_table() {
        Some(alt) => vec![alt.clone()],
        None => resolve_address_lookup_tables(rpc, &[PYRA_ADDRESS_TABLE]).await?,
    };

    let final_instructions = if has_compute_budget_instructions(instructions) {
        instructions.to_vec()
    } else {
        add_compute_budget(rpc, instructions, payer, &address_lookup_tables).await?
    };

    let blockhash = rpc.get_latest_blockhash().await?;

    build_v0_transaction_raw(
        &final_instructions,
        payer,
        &address_lookup_tables,
        blockhash,
    )
}

/// Build a V0 versioned transaction with a heap frame request.
///
/// Same as [`build_v0_transaction`] but also adds a `RequestHeapFrame` instruction
/// alongside the CU budget. Use for swap transactions that need extra heap
/// (e.g. `heap_size = 64 * 1024`).
///
/// `heap_size` is in bytes, must be a multiple of 1024, between 32KB and 256KB.
pub async fn build_v0_transaction_with_heap(
    rpc: &SolanaRpcClient,
    instructions: &[Instruction],
    payer: &Pubkey,
    heap_size: u32,
) -> Result<VersionedTransaction> {
    let address_lookup_tables = match rpc.pyra_address_lookup_table() {
        Some(alt) => vec![alt.clone()],
        None => resolve_address_lookup_tables(rpc, &[PYRA_ADDRESS_TABLE]).await?,
    };

    let final_instructions = if has_compute_budget_instructions(instructions) {
        instructions.to_vec()
    } else {
        let mut budget =
            add_compute_budget(rpc, instructions, payer, &address_lookup_tables).await?;
        // Insert heap frame after the CU limit and price instructions
        budget.insert(2, ComputeBudgetInstruction::request_heap_frame(heap_size));
        budget
    };

    let blockhash = rpc.get_latest_blockhash().await?;

    build_v0_transaction_raw(
        &final_instructions,
        payer,
        &address_lookup_tables,
        blockhash,
    )
}

/// Build a legacy transaction with automatic CU budget and priority fee.
///
/// Prepends compute budget instructions if not already present and fetches the latest blockhash.
/// Does not use ALTs — all addresses remain visible.
pub async fn build_legacy_transaction(
    rpc: &SolanaRpcClient,
    instructions: &[Instruction],
    payer: &Pubkey,
) -> Result<VersionedTransaction> {
    let final_instructions = if has_compute_budget_instructions(instructions) {
        instructions.to_vec()
    } else {
        add_compute_budget(rpc, instructions, payer, &[]).await?
    };

    let blockhash = rpc.get_latest_blockhash().await?;

    build_legacy_transaction_raw(&final_instructions, payer, blockhash)
}

/// Build a legacy transaction with a heap frame request.
///
/// Same as [`build_legacy_transaction`] but also adds a `RequestHeapFrame` instruction
/// alongside the CU budget. Use for swap transactions that need extra heap
/// (e.g. `heap_size = 64 * 1024`).
///
/// `heap_size` is in bytes, must be a multiple of 1024, between 32KB and 256KB.
pub async fn build_legacy_transaction_with_heap(
    rpc: &SolanaRpcClient,
    instructions: &[Instruction],
    payer: &Pubkey,
    heap_size: u32,
) -> Result<VersionedTransaction> {
    let final_instructions = if has_compute_budget_instructions(instructions) {
        instructions.to_vec()
    } else {
        let mut budget = add_compute_budget(rpc, instructions, payer, &[]).await?;
        budget.insert(2, ComputeBudgetInstruction::request_heap_frame(heap_size));
        budget
    };

    let blockhash = rpc.get_latest_blockhash().await?;

    build_legacy_transaction_raw(&final_instructions, payer, blockhash)
}

/// Build a V0 versioned transaction from instructions, payer, ALTs, and blockhash.
///
/// Returns an unsigned transaction with default signatures. Callers are responsible
/// for signing before submission.
pub fn build_v0_transaction_raw(
    instructions: &[Instruction],
    payer: &Pubkey,
    address_lookup_tables: &[AddressLookupTableAccount],
    blockhash: Hash,
) -> Result<VersionedTransaction> {
    let message = v0::Message::try_compile(payer, instructions, address_lookup_tables, blockhash)
        .map_err(|e| {
        PyraSolanaError::Transaction(format!("Failed to compile V0 message: {e}"))
    })?;

    let versioned_message = VersionedMessage::V0(message);
    let num_signers = usize::from(versioned_message.header().num_required_signatures);

    Ok(VersionedTransaction {
        signatures: vec![Signature::default(); num_signers],
        message: versioned_message,
    })
}

/// Build a legacy transaction from instructions, payer, and blockhash.
///
/// Legacy transactions keep all addresses visible (no ALT compression), which is
/// required for Privy's policy evaluator to inspect every account.
pub fn build_legacy_transaction_raw(
    instructions: &[Instruction],
    payer: &Pubkey,
    blockhash: Hash,
) -> Result<VersionedTransaction> {
    let message =
        solana_message::legacy::Message::new_with_blockhash(instructions, Some(payer), &blockhash);
    let versioned_message = VersionedMessage::Legacy(message);
    let num_signers = usize::from(versioned_message.header().num_required_signatures);

    Ok(VersionedTransaction {
        signatures: vec![Signature::default(); num_signers],
        message: versioned_message,
    })
}

/// Estimate compute units by simulating the transaction.
///
/// Prepends a max CU limit instruction to the simulation tx so that programs
/// exceeding Solana's default 200k CU budget can still report actual consumption.
/// Applies a 1.5x buffer to the simulated consumption and clamps the result
/// to at least [`MIN_COMPUTE_UNITS`].
///
/// Returns [`DEFAULT_COMPUTE_UNITS`] if simulation fails for any reason,
/// so callers can still build and send the transaction.
pub async fn estimate_compute_units(
    rpc: &SolanaRpcClient,
    instructions: &[Instruction],
    payer: &Pubkey,
    address_lookup_tables: &[AddressLookupTableAccount],
) -> Result<u32> {
    // Prepend a max CU limit so simulation doesn't fail on >200k CU programs
    let mut sim_instructions = Vec::with_capacity(
        instructions
            .len()
            .checked_add(1)
            .ok_or(PyraSolanaError::MathOverflow)?,
    );
    sim_instructions.push(ComputeBudgetInstruction::set_compute_unit_limit(
        SIMULATION_COMPUTE_UNIT_LIMIT,
    ));
    sim_instructions.extend_from_slice(instructions);

    let blockhash = Hash::default();
    let tx = build_v0_transaction_raw(&sim_instructions, payer, address_lookup_tables, blockhash)?;

    let result = match rpc.simulate_transaction(&tx).await {
        Ok(r) => r,
        Err(_) => return Ok(DEFAULT_COMPUTE_UNITS),
    };

    match result.units_consumed {
        Some(units) if result.err.is_none() => {
            #[allow(clippy::allow_attributes, clippy::allow_attributes_without_reason, clippy::arithmetic_side_effects, reason = "f64 mul cannot overflow")]
            let buffered = units as f64 * COMPUTE_UNIT_BUFFER_MULTIPLIER;
            let estimated = buffered as u32;
            Ok(estimated.max(MIN_COMPUTE_UNITS))
        }
        _ => Ok(DEFAULT_COMPUTE_UNITS),
    }
}

/// Estimate priority fee using IQR filtering of recent prioritization fees.
///
/// Collects writable accounts from the instructions, fetches recent prioritization
/// fees for those accounts, filters outliers using the interquartile range method,
/// and returns the mean of the remaining fees.
pub async fn estimate_priority_fee(
    rpc: &SolanaRpcClient,
    instructions: &[Instruction],
) -> Result<u64> {
    let mut writable_accounts: Vec<Pubkey> = instructions
        .iter()
        .flat_map(|ix| {
            ix.accounts
                .iter()
                .filter(|a| a.is_writable)
                .map(|a| a.pubkey)
        })
        .collect();
    writable_accounts.sort_unstable();
    writable_accounts.dedup();

    let fees = rpc
        .get_recent_prioritization_fees(&writable_accounts)
        .await?;

    let mut values: Vec<u64> = fees
        .iter()
        .map(|f| f.prioritization_fee)
        .filter(|&f| f > 0)
        .collect();

    iqr_filtered_mean(&mut values)
}

/// Compute the IQR-filtered mean of a set of fee values.
///
/// Sorts the input, removes outliers outside 1.5x IQR of Q1/Q3, and returns
/// the mean of the remaining values. Returns 0 if the input is empty or all
/// values are filtered out.
pub fn iqr_filtered_mean(values: &mut [u64]) -> Result<u64> {
    if values.is_empty() {
        return Ok(0);
    }

    values.sort_unstable();

    let len = values.len();
    let q1_idx = len.checked_div(4).ok_or(PyraSolanaError::MathOverflow)?;
    let q3_idx = len
        .checked_mul(3)
        .ok_or(PyraSolanaError::MathOverflow)?
        .checked_div(4)
        .ok_or(PyraSolanaError::MathOverflow)?;
    let q1 = values[q1_idx];
    let q3 = values[q3_idx];
    let iqr = q3.saturating_sub(q1);
    let iqr_bound = iqr.saturating_mul(3).saturating_div(2);
    let lower = q1.saturating_sub(iqr_bound);
    let upper = q3.saturating_add(iqr_bound);

    let filtered: Vec<u64> = values
        .iter()
        .copied()
        .filter(|&f| f >= lower && f <= upper)
        .collect();

    if filtered.is_empty() {
        return Ok(0);
    }

    let sum: u64 = filtered
        .iter()
        .try_fold(0u64, |acc, &f| acc.checked_add(f))
        .ok_or(PyraSolanaError::MathOverflow)?;
    let count = u64::try_from(filtered.len()).map_err(|_| PyraSolanaError::MathOverflow)?;
    let avg = sum
        .checked_div(count)
        .ok_or(PyraSolanaError::MathOverflow)?;

    Ok(avg)
}

/// Prepend compute budget instructions (CU limit + priority fee) to the given instructions.
///
/// Simulates the transaction to estimate compute units and fetches recent prioritization
/// fees, then prepends `SetComputeUnitLimit` and `SetComputeUnitPrice` instructions.
pub async fn add_compute_budget(
    rpc: &SolanaRpcClient,
    instructions: &[Instruction],
    payer: &Pubkey,
    address_lookup_tables: &[AddressLookupTableAccount],
) -> Result<Vec<Instruction>> {
    let cu_limit = estimate_compute_units(rpc, instructions, payer, address_lookup_tables).await?;
    let cu_price = estimate_priority_fee(rpc, instructions).await?;

    let mut result = Vec::with_capacity(
        instructions
            .len()
            .checked_add(2)
            .ok_or(PyraSolanaError::MathOverflow)?,
    );
    result.push(ComputeBudgetInstruction::set_compute_unit_limit(cu_limit));
    result.push(ComputeBudgetInstruction::set_compute_unit_price(cu_price));
    result.extend_from_slice(instructions);

    Ok(result)
}

/// Sign a versioned transaction with the given keypairs.
pub fn sign_transaction(
    message: VersionedMessage,
    signers: &[&solana_keypair::Keypair],
) -> Result<VersionedTransaction> {
    let signer_refs: Vec<&dyn solana_signer::Signer> = signers
        .iter()
        .map(|k| *k as &dyn solana_signer::Signer)
        .collect();
    VersionedTransaction::try_new(message, &signer_refs)
        .map_err(|e| PyraSolanaError::Transaction(format!("Failed to sign transaction: {e}")))
}

/// Partially sign a versioned transaction at signature index 0.
///
/// Places a single signature and leaves remaining slots as `Signature::default()` for
/// another party (e.g. Privy) to complete.
pub fn partial_sign_transaction(
    tx: &mut VersionedTransaction,
    signer: &solana_keypair::Keypair,
) -> Result<()> {
    let message_bytes = tx.message.serialize();
    let sig = solana_signer::Signer::try_sign_message(signer, &message_bytes)
        .map_err(|e| PyraSolanaError::Transaction(format!("Failed to partially sign: {e}")))?;
    let slot = tx
        .signatures
        .get_mut(0)
        .ok_or(PyraSolanaError::Transaction(
            "No signature slots in transaction".into(),
        ))?;
    *slot = sig;
    Ok(())
}

/// Resolve a list of ALT pubkeys into full `AddressLookupTableAccount` structs.
pub async fn resolve_address_lookup_tables(
    rpc: &SolanaRpcClient,
    alt_pubkeys: &[Pubkey],
) -> Result<Vec<AddressLookupTableAccount>> {
    let mut tables = Vec::with_capacity(alt_pubkeys.len());
    for pubkey in alt_pubkeys {
        tables.push(rpc.get_address_lookup_table(pubkey).await?);
    }
    Ok(tables)
}

/// Estimate the total transaction fee in lamports.
///
/// Fee = (compute_unit_limit * priority_fee_microlamports / 1_000_000) + base fee (5000 lamports).
pub fn estimate_transaction_fee(compute_unit_limit: u32, priority_fee_microlamports: u64) -> u64 {
    (compute_unit_limit as u64)
        .saturating_mul(priority_fee_microlamports)
        .saturating_div(1_000_000)
        .saturating_add(BASE_FEE_LAMPORTS)
}

/// Estimate the transaction fee and check it does not exceed the given maximum.
///
/// Returns `Ok(estimated_fee)` if within budget, or `Err(GasFeeTooHigh)` if not.
pub fn estimate_transaction_fee_with_gas_limit(
    compute_unit_limit: u32,
    priority_fee_microlamports: u64,
    max_fee_lamports: u64,
) -> Result<u64> {
    let estimated = estimate_transaction_fee(compute_unit_limit, priority_fee_microlamports);
    if estimated > max_fee_lamports {
        return Err(PyraSolanaError::GasFeeTooHigh {
            estimated,
            max: max_fee_lamports,
        });
    }
    Ok(estimated)
}
