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

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

#[derive(Accounts)]
pub struct RemoveLiquidity<'info> {
    // Vault Information
    #[account(mut)]
    pub vault: Box<Account<'info, Vault>>,
    // 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>>,

    // User Token Accounts
    #[account(mut)]
    pub token_a_funding_account: Box<Account<'info, TokenAccount>>,
    #[account(mut)]
    pub token_b_funding_account: Box<Account<'info, TokenAccount>>,

    // The following are for pool liquidity mining.
    // The owner should be the vault manager.
    /// also known as token a
    #[account(mut)]
    pub pool_coin_token_account: Box<Account<'info, TokenAccount>>,
    /// also known as token b
    #[account(mut)]
    pub pool_pc_token_account: Box<Account<'info, TokenAccount>>,

    // The Rayutdium AMM Pool Account
    #[account(mut)]
    pub amm: Box<Account<'info, AmmInfoV4>>,

    // Pool Metadata
    #[account(mut)]
    /// CHECK: Raydium verifies
    pub pool_withdraw_queue: AccountInfo<'info>,
    #[account(mut)]
    pub pool_temp_lp_token_account: Box<Account<'info, TokenAccount>>,

    /// AMM Metadata
    #[account(mut)]
    pub lp_mint: Box<Account<'info, Mint>>,
    /// CHECK: Raydium verifies
    pub amm_authority: AccountInfo<'info>,
    #[account(mut)]
    /// CHECK: Raydium verifies
    pub amm_open_orders: AccountInfo<'info>,
    #[account(mut)]
    /// CHECK: Raydium verifies
    pub amm_target: AccountInfo<'info>,

    // Serum Metadata
    #[account(mut)]
    /// CHECK: Raydium verifies
    pub serum_market: AccountInfo<'info>,
    #[account(mut)]
    /// CHECK: Raydium verifies
    pub serum_coin_vault_account: AccountInfo<'info>,
    #[account(mut)]
    /// CHECK: Raydium verifies
    pub serum_pc_vault_account: AccountInfo<'info>,
    /// CHECK: Raydium verifies
    pub serum_vault_signer: AccountInfo<'info>,

    #[account(mut)]
    /// CHECK: Raydium verifies
    pub serum_bids: AccountInfo<'info>,
    #[account(mut)]
    /// CHECK: Raydium verifies
    pub serum_asks: AccountInfo<'info>,
    #[account(mut)]
    /// CHECK: Raydium verifies
    pub serum_event_queue: AccountInfo<'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 pool_token_a_custody_account: Box<Account<'info, TokenAccount>>,
    #[account(mut)]
    pub pool_token_b_custody_account: Box<Account<'info, TokenAccount>>,

    // Raydium Pool Program.
    /// CHECK: Raydium verifies
    pub pool_program: AccountInfo<'info>,

    // Serum Program
    /// CHECK: Raydium verifies
    pub serum_program_id: AccountInfo<'info>,

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

    pub system_program: Program<'info, System>,
}

impl<'info> RemoveLiquidity<'info> {
    fn remove_liquidity(&mut self) -> Result<()> {
        // make sure withdraws are allowed.
        invariant!(self.vault.withdraws_allowed);

        // get the lp balances form the user.
        let lp_tokens_debt = self.user_info_account.lp_tokens_debt;
        msg!("Read balances. lp_tokens_debt: {}", lp_tokens_debt);

        let lp_remove_amount = lp_tokens_debt;

        if lp_remove_amount == 0 {
            msg!("Error: Zero balance. Forgot to unlock funds?");
            return Err(ProgramError::InsufficientFunds.into());
        }

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

        let initial_token_a_custody_balance =
            account::get_token_balance(&self.pool_token_a_custody_account.to_account_info())?;
        let initial_token_b_custody_balance =
            account::get_token_balance(&self.pool_token_b_custody_account.to_account_info())?;
        let initial_lp_tokens_balance =
            account::get_token_balance(&self.pool_lp_custody_account.to_account_info())?;

        raydium::remove_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_withdraw_queue.to_account_info(),
                self.pool_temp_lp_token_account.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(),
                self.serum_program_id.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(),
                self.serum_event_queue.to_account_info(),
                self.serum_bids.to_account_info(),
                self.serum_asks.to_account_info(),
                self.system_program.to_account_info(),
            ],
            signer_seeds,
            lp_remove_amount,
        )?;

        // check tokens received
        let tokens_a_received = account::get_balance_increase(
            &self.pool_token_a_custody_account.to_account_info(),
            initial_token_a_custody_balance,
        )?;
        let tokens_b_received = account::get_balance_increase(
            &self.pool_token_b_custody_account.to_account_info(),
            initial_token_b_custody_balance,
        )?;
        if tokens_a_received == 0 && tokens_b_received == 0 {
            // msg!("Error: Remove liquidity instruction didn't result in any of the tokens received");
            return Err(ProgramError::Custom(190).into());
        }
        let _ = account::check_tokens_spent(
            &self.pool_lp_custody_account.to_account_info(),
            initial_lp_tokens_balance,
            lp_remove_amount,
        )?;
        let transfer_accounts = Transfer {
            from: self.pool_token_a_custody_account.to_account_info(),
            to: self.token_a_funding_account.to_account_info(),
            authority: self.vault.to_account_info(),
        };
        token::transfer(
            CpiContext::new(self.spl_token_program.to_account_info(), transfer_accounts)
                .with_signer(signer_seeds),
            tokens_a_received,
        )?;

        let transfer_accounts = Transfer {
            from: self.pool_token_b_custody_account.to_account_info(),
            to: self.token_b_funding_account.to_account_info(),
            authority: self.vault.to_account_info(),
        };
        token::transfer(
            CpiContext::new(self.spl_token_program.to_account_info(), transfer_accounts)
                .with_signer(signer_seeds),
            tokens_b_received,
        )?;

        self.user_info_account
            .remove_liquidity(tokens_a_received, tokens_b_received)?;
        self.user_info_account
            .remove_lp_tokens_debt(lp_remove_amount)?;
        self.vault
            .remove_liquidity(tokens_a_received, tokens_b_received)?;

        Ok(())
    }
}

pub fn handler(ctx: Context<RemoveLiquidity>) -> Result<()> {
    ctx.accounts.remove_liquidity()?;
    Ok(())
}

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

        // make sure user is same user as user info account
        assert_keys_eq!(self.user_info_account.user_account, self.user_account);
        // make sure the user accounts are the correct mints.
        assert_keys_eq!(
            self.token_a_funding_account.mint,
            self.vault.pool_token_a_mint
        );
        assert_keys_eq!(
            self.token_b_funding_account.mint,
            self.vault.pool_token_b_mint
        );

        // make sure the owner of all user accounts is the user.
        assert_keys_eq!(self.token_a_funding_account.owner, self.user_account,);
        assert_keys_eq!(self.token_b_funding_account.owner, self.user_account,);

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

        // amm checks
        assert_keys_eq!(self.lp_mint, self.vault.pool_lp_token_mint);
        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);

        // Check all unchecked accounts for AMM metadata.
        assert_keys_eq!(self.serum_market, self.vault.serum_market);
        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 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
        );

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

        // serum program id
        assert_keys_eq!(self.serum_program_id, self.vault.serum_program_id);

        Ok(())
    }
}