use crate::Vault;
use crate::*;
use anchor_lang::prelude::*;
use anchor_spl::token::{Mint, Token, TokenAccount};
use solana_farm_sdk::id::zero;
use solana_farm_sdk::{program::account, program::protocol::raydium};
use vipers::{assert_keys_eq, Validate};

#[derive(Accounts)]
pub struct Crank3<'info> {
    #[account(mut)]
    pub vault: Box<Account<'info, Vault>>,
    /// Manager of the vault.
    pub manager: Signer<'info>,
    #[account(mut)]
    pub pool_token_a_custody_account: Box<Account<'info, TokenAccount>>,
    #[account(mut)]
    pub pool_token_b_custody_account: Box<Account<'info, TokenAccount>>,
    #[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>>,

    // For Farming Metadata.
    #[account(mut)]
    /// CHECK: Raydium will verify
    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)]
    pub farm_reward_token_b_account: Box<Account<'info, TokenAccount>>,

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

    // Raydium Programs.
    /// CHECK: Raydium will verify
    pub pool_program: UncheckedAccount<'info>,
    /// AMM/Pool Metadata
    #[account(mut)]
    pub pool_coin_token_account: Box<Account<'info, TokenAccount>>,

    #[account(mut)]
    pub pool_pc_token_account: Box<Account<'info, TokenAccount>>,
    // SPL Token Program
    pub spl_token_program: Program<'info, Token>,
    #[account(mut)]
    pub amm: Box<Account<'info, AmmInfoV4>>,
    #[account(mut)]
    /// CHECK: Raydium will verify
    pub amm_authority: UncheckedAccount<'info>,
    #[account(mut)]
    /// CHECK: Raydium will verify
    pub amm_open_orders: UncheckedAccount<'info>,
    #[account(mut)]
    /// CHECK: Raydium will verify
    pub amm_target: UncheckedAccount<'info>,
    #[account(mut)]
    /// CHECK: Raydium will verify
    pub serum_market: UncheckedAccount<'info>,

    #[account(mut)]
    pub lp_mint: Account<'info, Mint>,

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

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

    /// Raydium Farm Metadata
    /// The following are accounts from the AMM that need to be passed in
    /// CHECK: Raydium will verify
    pub farm_authority: UncheckedAccount<'info>,

    // Raydium Farm Program.
    /// CHECK: Raydium will verify
    pub farm_program: UncheckedAccount<'info>,
}

impl<'info> Crank3<'info> {
    fn crank3(&mut self) -> Result<()> {
        check_min_crank_interval(&self.vault)?;
        self.vault.update_crank_time()?;
        self.vault.crank_step = 3;
        // read balances
        let token_a_balance =
            account::get_token_balance(&self.pool_token_a_custody_account.to_account_info())?;
        let token_b_balance =
            account::get_token_balance(&self.pool_token_b_custody_account.to_account_info())?;
        let lp_token_balance =
            account::get_token_balance(&self.pool_lp_custody_account.to_account_info())?;
        msg!(
            "Read balances. token_a_balance: {}, token_b_balance: {}",
            token_a_balance,
            token_b_balance
        );
        if token_a_balance < 10 || token_b_balance < 10 {
            msg!("Nothing to do: Not enough tokens to compound");
            return Ok(());
        }

        // compute and check pool ratios
        let (pool_coin_balance, pool_pc_balance) = raydium::get_pool_token_balances(
            &self.pool_coin_token_account.to_account_info(),
            &self.pool_pc_token_account.to_account_info(),
            &self.amm_open_orders.to_account_info(),
            &self.amm.to_account_info(),
        )?;
        let pool_ratio = if pool_coin_balance != 0 {
            pool_pc_balance as f64 / pool_coin_balance as f64
        } else {
            0.0
        };
        let custody_ratio = account::get_token_pair_ratio(
            &self.pool_token_a_custody_account.to_account_info(),
            &self.pool_token_b_custody_account.to_account_info(),
        )?;
        msg!(
            "Compute pool ratios. custody_ratio: {}, pool_ratio: {}",
            custody_ratio,
            pool_ratio
        );
        if custody_ratio == 0.0 || pool_ratio == 0.0 {
            msg!("Pool ratio is zero");
            return Ok(());
        }
        if (custody_ratio - pool_ratio).abs() * 100.0 / pool_ratio > 10.0 {
            msg!("Unbalanced tokens, run Crank2 first");
            return Ok(());
        }

        let seed = gen_vault_signer_seeds!(self.vault);
        let signer_seeds = &[&seed[..]];

        // calculate deposit amounts
        let (max_token_a_deposit_amount, max_token_b_deposit_amount) =
            if custody_ratio >= pool_ratio {
                raydium::get_pool_deposit_amounts(
                    &self.pool_coin_token_account.to_account_info(),
                    &self.pool_pc_token_account.to_account_info(),
                    &self.amm_open_orders.to_account_info(),
                    &self.amm.to_account_info(),
                    token_a_balance,
                    0,
                )?
            } else {
                raydium::get_pool_deposit_amounts(
                    &self.pool_coin_token_account.to_account_info(),
                    &self.pool_pc_token_account.to_account_info(),
                    &self.amm_open_orders.to_account_info(),
                    &self.amm.to_account_info(),
                    0,
                    token_b_balance,
                )?
            };
        // one of the amounts can come out over the balance because ratios didn't reflect
        // deposited volume, while get_pool_deposit_amounts does include it.
        // in this case we just flip the side.
        let (max_token_a_deposit_amount, max_token_b_deposit_amount) =
            if max_token_b_deposit_amount > token_b_balance {
                raydium::get_pool_deposit_amounts(
                    &self.pool_coin_token_account.to_account_info(),
                    &self.pool_pc_token_account.to_account_info(),
                    &self.amm_open_orders.to_account_info(),
                    &self.amm.to_account_info(),
                    0,
                    token_b_balance,
                )?
            } else if max_token_a_deposit_amount > token_a_balance {
                raydium::get_pool_deposit_amounts(
                    &self.pool_coin_token_account.to_account_info(),
                    &self.pool_pc_token_account.to_account_info(),
                    &self.amm_open_orders.to_account_info(),
                    &self.amm.to_account_info(),
                    token_a_balance,
                    0,
                )?
            } else {
                (max_token_a_deposit_amount, max_token_b_deposit_amount)
            };

        msg!("Deposit tokens into the pool. max_token_a_deposit_amount: {}, max_token_b_deposit_amount: {}",
                 max_token_a_deposit_amount,
                 max_token_b_deposit_amount);
        if max_token_a_deposit_amount == 0
            || max_token_b_deposit_amount == 0
            || raydium::estimate_lp_tokens_amount(
                &self.lp_mint.to_account_info(),
                max_token_a_deposit_amount,
                max_token_b_deposit_amount,
                pool_coin_balance,
                pool_pc_balance,
            )? < 2
        {
            msg!("Nothing to do: Tokens balance is not large enough");
            return Ok(());
        }

        raydium::add_liquidity_with_seeds(
            &[
                self.vault.to_account_info(),
                self.pool_token_a_custody_account.to_account_info(),
                self.pool_token_b_custody_account.to_account_info(),
                self.pool_lp_custody_account.to_account_info(),
                self.pool_program.to_account_info(),
                self.pool_coin_token_account.to_account_info(),
                self.pool_pc_token_account.to_account_info(),
                self.lp_mint.to_account_info(),
                self.spl_token_program.to_account_info(),
                self.amm.to_account_info(),
                self.amm_authority.to_account_info(),
                self.amm_open_orders.to_account_info(),
                self.amm_target.to_account_info(),
                self.serum_market.to_account_info(),
            ],
            signer_seeds,
            max_token_a_deposit_amount,
            max_token_b_deposit_amount,
        )?;

        // Check tokens spent and return change back to user
        let tokens_a_spent = account::check_tokens_spent(
            &self.pool_token_a_custody_account.to_account_info(),
            token_a_balance,
            max_token_a_deposit_amount,
        )?;
        let tokens_b_spent = account::check_tokens_spent(
            &self.pool_token_b_custody_account.to_account_info(),
            token_b_balance,
            max_token_b_deposit_amount,
        )?;

        // Stake LP tokens
        let dual_rewards = self.farm_reward_token_b_account.key() != zero::id();
        let lp_tokens_received = account::check_tokens_received(
            &self.pool_lp_custody_account.to_account_info(),
            lp_token_balance,
            1,
        )?;
        msg!(
            "Stake LP tokens. tokens_a_spent: {}, tokens_b_spent: {}, lp_tokens_received: {}",
            tokens_a_spent,
            tokens_b_spent,
            lp_tokens_received
        );
        let token_a_reward_balance = self.farm_token_a_reward_custody_account.amount;
        let token_b_reward_balance = if dual_rewards {
            self.farm_token_b_reward_custody_account.amount
        } else {
            0
        };

        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,
        )?;
        if lp_token_balance
            != account::get_token_balance(&self.pool_lp_custody_account.to_account_info())?
        {
            msg!("Error: Stake instruction didn't result in expected amount of LP tokens spent");
            return Err(ProgramError::Custom(165).into());
        }

        // update Vault stats
        let token_a_rewards = account::get_balance_increase(
            &self.farm_token_a_reward_custody_account.to_account_info(),
            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(),
                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)?;
        self.vault.add_liquidity(tokens_a_spent, tokens_b_spent)?;

        Ok(())
    }
}
pub fn handler(ctx: Context<Crank3>) -> Result<()> {
    ctx.accounts.crank3()
}

impl<'info> Validate<'info> for Crank3<'info> {
    fn validate(&self) -> Result<()> {
        // make sure that only the vault manager can call this crank.
        assert_keys_eq!(self.manager, self.vault.manager);
        // make sure the custody accounts are the same as those in the vault.
        assert_keys_eq!(
            self.pool_token_a_custody_account,
            self.vault.pool_token_a_custody_account
        );
        assert_keys_eq!(
            self.pool_token_b_custody_account,
            self.vault.pool_token_b_custody_account
        );
        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
        );

        // 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
        );

        // make sure the stake info account is in the vault.
        assert_keys_eq!(self.fees_account_a, self.vault.fees_account_a);
        assert_keys_eq!(self.fees_account_b, self.vault.fees_account_b);

        // raydium pool program
        assert_keys_eq!(self.pool_program, self.vault.pool_program_id);

        // make sure the pool token accounts are the same as the ones in the amm.
        assert_keys_eq!(
            self.pool_coin_token_account,
            self.vault.pool_coin_token_account
        );
        assert_keys_eq!(self.pool_pc_token_account, self.vault.pool_pc_token_account);

        // make sure vault is connected to this amm.
        assert_keys_eq!(self.vault.amm, self.amm);

        // amm checks
        assert_keys_eq!(self.amm_authority, self.vault.amm_authority);
        assert_keys_eq!(self.amm_open_orders, self.vault.amm_open_orders);
        assert_keys_eq!(self.amm_target, self.vault.amm_target);

        // serum checks
        assert_keys_eq!(self.serum_market, self.vault.serum_market);

        // 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);
        assert_keys_eq!(self.farm_program, self.vault.farm_program_id);

        Ok(())
    }
}