use anchor_lang::prelude::*;
use anchor_spl::token::{self, Mint, MintTo, Token, TokenAccount};
use solana_farm_sdk::{id::zero, program::account, program::protocol::raydium};
use solana_program::program_error::ProgramError;
use vipers::{assert_keys_eq, invariant, Validate};

use crate::{gen_vault_signer_seeds, UserInfo, Vault, VaultStakeEvent};

#[derive(Accounts)]
pub struct Stake<'info> {
    // Vault Information
    #[account(mut)]
    pub vault: Box<Account<'info, Vault>>,
    // The LP Token Mint of the vault.
    #[account(mut)]
    pub vault_lp_token_mint: Box<Account<'info, Mint>>,
    #[account(mut)]
    pub lp_funding_account: Box<Account<'info, TokenAccount>>,
    // User information
    // The following should be passed in by the user of this strategy.
    pub user_account: Signer<'info>,

    // The user info account keeps track of the use of this vault.
    #[account(mut)]
    pub user_info_account: Box<Account<'info, UserInfo>>,

    #[account(mut)]
    pub vault_lp_token_account: Box<Account<'info, TokenAccount>>,

    // The Raydium Farm Account
    #[account(mut)]
    /// CHECK: Raydium verifies
    pub farm: UncheckedAccount<'info>,

    /// Raydium Farm Metadata
    /// CHECK: Raydium verifies
    pub farm_authority: UncheckedAccount<'info>,

    // For Farming Metadata.
    #[account(mut)]
    /// CHECK: Raydium verifies
    pub stake_info_account: UncheckedAccount<'info>,

    // The following are farm reward vaults.
    // They are owned by raydium not delta.
    #[account(mut)]
    pub farm_lp_token_account: Box<Account<'info, TokenAccount>>,
    #[account(mut)]
    pub farm_reward_token_a_account: Box<Account<'info, TokenAccount>>,
    #[account(mut)]
    /// CHECK: Raydium verifies
    pub farm_reward_token_b_account: UncheckedAccount<'info>,

    // The following are custody accounts owned by the vault.
    // they are used for harvesting.
    #[account(mut)]
    pub pool_lp_custody_account: Box<Account<'info, TokenAccount>>,
    #[account(mut)]
    pub farm_token_a_reward_custody_account: Box<Account<'info, TokenAccount>>,
    #[account(mut)]
    pub farm_token_b_reward_custody_account: Box<Account<'info, TokenAccount>>,

    // Raydium Farm Program.
    // Make sure this is the same as in the vault.
    /// CHECK: Raydium verifies
    pub farm_program: UncheckedAccount<'info>,

    // SPL Token Program
    pub spl_token_program: Program<'info, Token>,

    // Clock Program
    pub clock_program: Sysvar<'info, Clock>,
}

impl<'info> Stake<'info> {
    fn stake(&mut self, lp_tokens_received: u64) -> Result<()> {
        account::transfer_tokens(
            &self.lp_funding_account.to_account_info(),
            &self.pool_lp_custody_account.to_account_info(),
            &self.user_account,
            lp_tokens_received,
        )?;

        let initial_lp_token_custody_balance = self.pool_lp_custody_account.amount;

        // // Stake LP tokens
        let seed = gen_vault_signer_seeds!(self.vault);
        let signer_seeds = &[&seed[..]];

        let dual_rewards = self.farm_reward_token_b_account.key() != zero::id();
        let initial_token_a_reward_balance = self.farm_token_a_reward_custody_account.amount;
        let initial_token_b_reward_balance = if dual_rewards {
            self.farm_token_b_reward_custody_account.amount
        } else {
            0
        };

        msg!("Stake LP tokens");
        let stake_balance =
            raydium::get_stake_account_balance(&self.stake_info_account.to_account_info())?;

        msg!("starting to stake...");

        raydium::stake_with_seeds(
            &[
                self.vault.to_account_info(),
                self.stake_info_account.to_account_info(),
                // first the farm lp and reward custody accounts - all should be owned by the vault.
                self.pool_lp_custody_account.to_account_info(),
                self.farm_token_a_reward_custody_account.to_account_info(),
                self.farm_token_b_reward_custody_account.to_account_info(),
                self.farm_program.to_account_info(),
                // then comes the farm lp and reward accounts.
                // these live in the farm and are owned by raydium.
                self.farm_lp_token_account.to_account_info(),
                self.farm_reward_token_a_account.to_account_info(),
                self.farm_reward_token_b_account.to_account_info(),
                // solana information
                self.clock_program.to_account_info(),
                self.spl_token_program.to_account_info(),
                // farm information.
                self.farm.to_account_info(),
                self.farm_authority.to_account_info(),
            ],
            signer_seeds,
            lp_tokens_received,
        )?;

        msg!("Finished staking.");
        invariant!(
            initial_lp_token_custody_balance == self.pool_lp_custody_account.amount,
            "Error: Stake instruction didn't result in expected amount of LP tokens spent"
        );

        // update Vault stats
        let token_a_rewards = account::get_balance_increase(
            &self.farm_token_a_reward_custody_account.to_account_info(),
            initial_token_a_reward_balance,
        )?;
        let token_b_rewards = if dual_rewards {
            account::get_balance_increase(
                &self.farm_token_b_reward_custody_account.to_account_info(),
                initial_token_b_reward_balance,
            )?
        } else {
            0
        };
        msg!(
            "Update Vault stats. token_a_rewards: {}, token_b_rewards: {}",
            token_a_rewards,
            token_b_rewards
        );
        self.vault.add_rewards(token_a_rewards, token_b_rewards)?;

        // compute Vault tokens to mint
        let vault_token_supply_amount = self.vault_lp_token_mint.supply;
        let vt_to_mint = if vault_token_supply_amount == 0 || stake_balance == 0 {
            lp_tokens_received
        } else {
            account::to_token_amount(
                lp_tokens_received as f64 / stake_balance as f64
                    * self.vault_lp_token_mint.supply as f64,
                0,
            )?
        };

        // mint vault tokens to user
        msg!(
            "Mint Vault tokens to the user. vt_to_mint: {}, vt_supply_amount: {}, stake_balance: {}",
            vt_to_mint, vault_token_supply_amount,
            stake_balance
        );
        if vt_to_mint == 0 {
            msg!("Error: Add liquidity instruction didn't result in Vault tokens mint");
            return Err(ProgramError::Custom(170).into());
        }

        let mint_to_accounts = MintTo {
            mint: self.vault_lp_token_mint.to_account_info(),
            to: self.vault_lp_token_account.to_account_info(),
            authority: self.vault.to_account_info(),
        };
        let cpi_ctx = CpiContext::new_with_signer(
            self.spl_token_program.to_account_info(),
            mint_to_accounts,
            signer_seeds,
        );
        token::mint_to(cpi_ctx, vt_to_mint)?;
        msg!("Completed add liquidity and stake.");

        Ok(())
    }
}

pub fn handler(ctx: Context<Stake>, amount: u64) -> Result<()> {
    ctx.accounts.stake(amount)?;
    emit!(VaultStakeEvent {
        vault: ctx.accounts.vault.key(),
        amount,
    });
    Ok(())
}

impl<'info> Validate<'info> for Stake<'info> {
    fn validate(&self) -> Result<()> {
        // make sure the user info account is for this vault.
        assert_keys_eq!(self.user_info_account.vault, self.vault);

        assert_keys_eq!(self.vault_lp_token_account.owner, self.user_account,);

        // make sure user is same user as user info account
        assert_keys_eq!(self.user_info_account.user_account, self.user_account);

        // make sure that the vault token account is the same mint as the vault lp mint.
        assert_keys_eq!(
            self.vault.vault_lp_token_mint,
            self.vault_lp_token_account.mint,
            "Vault lp token account mint is incorrect."
        );

        assert_keys_eq!(self.lp_funding_account.mint, self.vault.pool_lp_token_mint);

        // make sure farm metadata is same as specified in vault
        assert_keys_eq!(self.farm, self.vault.farm);
        assert_keys_eq!(self.farm_authority, self.vault.farm_authority);

        // TODO - not specifying stake info account because it can be different on devnet and on mainnet.

        // Farm Vaults
        assert_keys_eq!(self.farm_lp_token_account, self.vault.farm_lp_token_account);
        assert_keys_eq!(
            self.farm_reward_token_a_account,
            self.vault.farm_reward_token_a_account
        );
        assert_keys_eq!(
            self.farm_reward_token_b_account,
            self.vault.farm_reward_token_b_account
        );

        // custody accounts

        assert_keys_eq!(
            self.pool_lp_custody_account,
            self.vault.pool_lp_custody_account
        );
        assert_keys_eq!(
            self.farm_token_a_reward_custody_account,
            self.vault.farm_token_a_reward_custody_account
        );
        assert_keys_eq!(
            self.farm_token_b_reward_custody_account,
            self.vault.farm_token_b_reward_custody_account
        );

        assert_keys_eq!(self.farm_program, self.vault.farm_program_id);

        Ok(())
    }
}