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

#[derive(Accounts)]
pub struct Crank2<'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>>,

    // Raydium Programs.
    /// CHECK: Raydium will verify
    pub raydium_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>,
    // Serum Program
    /// CHECK: Raydium will verify
    pub serum_program_id: UncheckedAccount<'info>,
    #[account(mut)]
    /// CHECK: Raydium will verify
    pub serum_bids: UncheckedAccount<'info>,
    #[account(mut)]
    /// CHECK: Raydium will verify
    pub serum_asks: UncheckedAccount<'info>,
    #[account(mut)]
    /// CHECK: Raydium will verify
    pub serum_event_queue: UncheckedAccount<'info>,
    #[account(mut)]
    /// CHECK: Raydium will verify
    pub serum_coin_vault_account: Box<Account<'info, TokenAccount>>,
    #[account(mut)]
    /// CHECK: Raydium will verify
    pub serum_pc_vault_account: UncheckedAccount<'info>,
    /// CHECK: Raydium will verify
    pub serum_vault_signer: UncheckedAccount<'info>,

    // 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 verifies
    pub farm_authority: UncheckedAccount<'info>,

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

impl<'info> Crank2<'info> {
    fn crank2(&mut self) -> Result<()> {
        check_min_crank_interval(&self.vault)?;
        self.vault.update_crank_time()?;
        self.vault.crank_step = 2;

        // read reward balances
        let dual_rewards = self.farm_reward_token_b_account.key() != zero::id();
        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
        };
        msg!(
            "Read reward balances. token_a_reward_balance: {}, token_b_reward_balance: {}",
            token_a_reward_balance,
            token_b_reward_balance
        );

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

        let token_a_reward_mint = account::get_token_account_mint(
            &self.farm_token_a_reward_custody_account.to_account_info(),
        )?;
        let token_a_custody_mint =
            account::get_token_account_mint(&self.pool_token_a_custody_account.to_account_info())?;
        let token_b_custody_mint =
            account::get_token_account_mint(&self.pool_token_b_custody_account.to_account_info())?;

        if token_a_reward_mint == token_a_custody_mint {
            pda::transfer_tokens_with_seeds(
                &self.farm_token_a_reward_custody_account.to_account_info(),
                &self.pool_token_a_custody_account.to_account_info(),
                &self.vault.to_account_info(),
                signer_seeds,
                token_a_reward_balance,
            )?;
        } else if token_a_reward_mint == token_b_custody_mint {
            pda::transfer_tokens_with_seeds(
                &self.farm_token_a_reward_custody_account.to_account_info(),
                &self.pool_token_b_custody_account.to_account_info(),
                &self.vault.to_account_info(),
                signer_seeds,
                token_a_reward_balance,
            )?;
        }

        if dual_rewards {
            let token_b_reward_mint = account::get_token_account_mint(
                &self.farm_token_b_reward_custody_account.to_account_info(),
            )?;
            if token_b_reward_mint == token_b_custody_mint {
                pda::transfer_tokens_with_seeds(
                    &self.farm_token_b_reward_custody_account.to_account_info(),
                    &self.pool_token_b_custody_account.to_account_info(),
                    &self.vault.to_account_info(),
                    signer_seeds,
                    token_b_reward_balance,
                )?;
            } else if token_b_reward_mint == token_a_custody_mint {
                pda::transfer_tokens_with_seeds(
                    &self.farm_token_b_reward_custody_account.to_account_info(),
                    &self.pool_token_a_custody_account.to_account_info(),
                    &self.vault.to_account_info(),
                    signer_seeds,
                    token_b_reward_balance,
                )?;
            }
        }

        // 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())?;
        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 balance");
            return Ok(());
        }

        // rebalance
        // 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 pool_ratio == 0.0 {
            msg!("Can't balance: Pool ratio is zero");
            return Ok(());
        }
        if custody_ratio > 0.0 && (custody_ratio - pool_ratio).abs() * 100.0 / pool_ratio < 3.0 {
            msg!("Nothing to do: Already balanced");
            return Ok(());
        }

        // compute ui amount to exchange
        let extra_a_tokens =
            (token_a_balance as f64 * pool_ratio - token_b_balance as f64) / (2.0 * pool_ratio);
        let extra_b_tokens = extra_a_tokens * pool_ratio;
        let reverse = extra_a_tokens < 0.0;
        msg!(
            "Rebalance tokens. reverse: {}, extra_a_tokens: {}, extra_b_tokens: {}",
            reverse,
            extra_a_tokens,
            extra_b_tokens
        );

        let token_a_swap_custody = if reverse {
            self.pool_token_b_custody_account.to_account_info()
        } else {
            self.pool_token_a_custody_account.to_account_info()
        };
        let token_b_swap_custody = if reverse {
            self.pool_token_a_custody_account.to_account_info()
        } else {
            self.pool_token_b_custody_account.to_account_info()
        };
        let coint_extra_amount_in = if !reverse {
            account::to_token_amount(extra_a_tokens.abs(), 0)?
        } else {
            0
        };
        let pc_extra_amount_in = if !reverse {
            0
        } else {
            account::to_token_amount(extra_b_tokens.abs(), 0)?
        };
        if coint_extra_amount_in < 2 && pc_extra_amount_in < 2 {
            msg!("Nothing to do: Not enough tokens to balance");
            return Ok(());
        }

        // get exact swap amounts
        let (amount_in, min_amount_out) = raydium::get_pool_swap_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(),
            coint_extra_amount_in,
            pc_extra_amount_in,
        )?;
        msg!(
            "Swap. amount_in: {}, min_amount_out {}",
            amount_in,
            min_amount_out
        );
        if amount_in == 0 || min_amount_out == 0 {
            msg!("Nothing to do: Not enough tokens to balance");
            return Ok(());
        }

        let initial_tokens_spent_balance = account::get_token_balance(&token_a_swap_custody)?;
        let initial_tokens_received_balance = account::get_token_balance(&token_b_swap_custody)?;

        raydium::swap_with_seeds_fixed_in(
            &[
                self.vault.to_account_info(),
                token_a_swap_custody.clone(),
                token_b_swap_custody.clone(),
                self.raydium_pool_program.to_account_info(),
                self.pool_coin_token_account.to_account_info(),
                self.pool_pc_token_account.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(),
                self.serum_program_id.to_account_info(),
                self.serum_bids.to_account_info(),
                self.serum_asks.to_account_info(),
                self.serum_event_queue.to_account_info(),
                self.serum_coin_vault_account.to_account_info(),
                self.serum_pc_vault_account.to_account_info(),
                self.serum_vault_signer.to_account_info(),
            ],
            signer_seeds,
            amount_in,
            min_amount_out,
        )?;
        let _ = account::check_tokens_spent(
            &token_a_swap_custody,
            initial_tokens_spent_balance,
            amount_in,
        )?;
        let tokens_received = account::check_tokens_received(
            &token_b_swap_custody,
            initial_tokens_received_balance,
            min_amount_out,
        )?;

        msg!(
            "Done. tokens_received: {}, token_a_balance: {}, token_b_balance: {}",
            tokens_received,
            account::get_token_balance(&self.pool_token_a_custody_account.to_account_info())?,
            account::get_token_balance(&self.pool_token_b_custody_account.to_account_info())?
        );
        Ok(())
    }
}
pub fn handler(ctx: Context<Crank2>) -> Result<()> {
    ctx.accounts.crank2()
}

impl<'info> Validate<'info> for Crank2<'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
        );

        // raydium pool program
        assert_keys_eq!(self.raydium_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);
        assert_keys_eq!(self.serum_program_id, self.vault.serum_program_id);
        assert_keys_eq!(self.serum_bids, self.vault.serum_bids);
        assert_keys_eq!(self.serum_asks, self.vault.serum_asks);
        assert_keys_eq!(self.serum_event_queue, self.vault.serum_event_queue);
        assert_keys_eq!(
            self.serum_coin_vault_account,
            self.vault.serum_coin_vault_account
        );
        assert_keys_eq!(
            self.serum_pc_vault_account,
            self.vault.serum_pc_vault_account
        );
        assert_keys_eq!(self.serum_vault_signer, self.vault.serum_vault_signer);

        // 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(())
    }
}