use solana_program::pubkey::Pubkey;
use spl_associated_token_account::get_associated_token_address;
use steel::*;

use crate::state::{Proposal, Vote};
use oil_api::{
    consts::MINT_ADDRESS,
    state::{miner_pda, stake_pda_with_id, treasury_pda, Miner, Stake, Treasury},
};

/// Voting power calculation context
/// Contains all the account data needed to calculate voting power
pub struct VotingPowerContext {
    /// Wallet OIL token account data (optional)
    pub wallet_token_account: Option<Vec<u8>>,
    /// Stake account data (one per stake ID)
    pub stake_accounts: Vec<Option<Vec<u8>>>,
    /// Miner account data (optional)
    pub miner_account: Option<Vec<u8>>,
    /// Treasury account data (optional, needed for refined oil calculation)
    pub treasury_account: Option<Vec<u8>>,
}

impl VotingPowerContext {
    /// Calculate voting power from the context
    /// Formula: wallet_balance (1x) + stake_power (multiplied) + unrefined_oil (20x)
    pub fn calculate_voting_power(&self) -> Result<u64, String> {
        let mut total_voting_power = 0u64;
        
        // 1. Get wallet balance (1x multiplier)
        if let Some(ref account_data) = self.wallet_token_account {
            if !account_data.is_empty() && account_data.len() >= 72 {
                // Token account balance is at offset 64 (mint 32 + owner 32 + amount 8)
                let balance_bytes = &account_data[64..72];
                let wallet_balance = u64::from_le_bytes([
                    balance_bytes[0], balance_bytes[1], balance_bytes[2], balance_bytes[3],
                    balance_bytes[4], balance_bytes[5], balance_bytes[6], balance_bytes[7],
                ]);
                total_voting_power = total_voting_power.saturating_add(wallet_balance);
            }
        }
        
        // 2. Get stake power (already multiplied based on lock duration)
        for stake_account_opt in &self.stake_accounts {
            if let Some(ref account_data) = stake_account_opt {
                if !account_data.is_empty() {
                    if let Ok(stake) = Stake::try_from_bytes(account_data) {
                        // Calculate stake score (balance * multiplier)
                        let stake_score = stake.score();
                        total_voting_power = total_voting_power.saturating_add(stake_score);
                    }
                }
            }
        }
        
        // 3. Get unrefined oil (20x multiplier)
        // Use calculated refined oil values (same as /user/state endpoint) to match client calculation
        if let Some(ref miner_data) = self.miner_account {
            if !miner_data.is_empty() {
                if let Ok(miner) = Miner::try_from_bytes(miner_data) {
                    use steel::Numeric;
                    
                    // Calculate refined_oil using current treasury data (read fresh from chain, don't use cache)
                    // This ensures we use the same treasury factors that the client sees
                    let mut calculated_block_refined_oil = miner.block_refined_oil;
                    let mut calculated_auction_refined_oil = miner.auction_refined_oil;
                    
                    // Read treasury account fresh from chain (same as client does)
                    if let Some(ref treasury_data) = self.treasury_account {
                        if !treasury_data.is_empty() {
                            if let Ok(treasury) = Treasury::try_from_bytes(treasury_data) {
                                // Calculate block_refined_oil with fresh treasury factor
                                if treasury.block_rewards_factor > miner.block_rewards_factor {
                                    let accumulated_rewards = treasury.block_rewards_factor - miner.block_rewards_factor;
                                    let personal_rewards = accumulated_rewards * Numeric::from_u64(miner.block_rewards_oil);
                                    calculated_block_refined_oil = miner.block_refined_oil + personal_rewards.to_u64();
                                }
                                
                                // Calculate auction_refined_oil with fresh treasury factor
                                if treasury.auction_rewards_factor > miner.auction_rewards_factor {
                                    let accumulated_rewards = treasury.auction_rewards_factor - miner.auction_rewards_factor;
                                    let personal_rewards = accumulated_rewards * Numeric::from_u64(miner.auction_rewards_oil);
                                    calculated_auction_refined_oil = miner.auction_refined_oil + personal_rewards.to_u64();
                                }
                            }
                        }
                    }
                    
                    // Calculate unrefined oil using calculated refined oil values (matches client)
                    let block_unrefined = miner.block_rewards_oil.saturating_sub(calculated_block_refined_oil);
                    let auction_unrefined = miner.auction_rewards_oil.saturating_sub(calculated_auction_refined_oil);
                    let total_unrefined = block_unrefined.saturating_add(auction_unrefined);
                    
                    // Apply 20x multiplier
                    let unrefined_power = total_unrefined.saturating_mul(20);
                    total_voting_power = total_voting_power.saturating_add(unrefined_power);
                }
            }
        }
        
        Ok(total_voting_power)
    }
    
    /// Get the addresses needed to fetch accounts for voting power calculation
    pub fn get_required_addresses(authority: &Pubkey, stake_ids: &[u64]) -> Vec<Pubkey> {
        let mut addresses = Vec::new();
        
        // Wallet token account
        let mint_address = MINT_ADDRESS;
        addresses.push(get_associated_token_address(authority, &mint_address));
        
        // Stake accounts
        for &stake_id in stake_ids {
            addresses.push(stake_pda_with_id(*authority, stake_id).0);
        }
        
        // Miner account
        addresses.push(miner_pda(*authority).0);
        
        // Treasury account
        addresses.push(treasury_pda().0);
        
        addresses
    }
}

/// Parse a Vote account from raw account data
pub fn parse_vote_account(data: &[u8]) -> Result<Vote, String> {
    if data.is_empty() || data.len() < std::mem::size_of::<Vote>() {
        return Err("Vote account data too short".to_string());
    }
    
    Vote::try_from_bytes(data)
        .copied()
        .map_err(|e| format!("Failed to parse Vote account: {:?}", e))
}

/// Parse a Proposal account from raw account data
pub fn parse_proposal_account(data: &[u8]) -> Result<Proposal, String> {
    if data.is_empty() || data.len() < std::mem::size_of::<Proposal>() {
        return Err("Proposal account data too short".to_string());
    }
    
    Proposal::try_from_bytes(data)
        .copied()
        .map_err(|e| format!("Failed to parse Proposal account: {:?}", e))
}
