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

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

#[derive(Accounts)]
pub struct AddLiquidity<'info> {
    // Vault Information
    #[account(mut)]
    pub vault: Box<Account<'info, Vault>>,
    // The LP Token Mint of the 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>>,
    #[account(mut)]
    pub lp_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>>,
    /// AMM Metadata
    #[account(mut)]
    pub lp_mint: Account<'info, Mint>,
    #[account(mut)]
    /// CHECK: Raydium will verify
    pub amm_authority: 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)]
    /// CHECK: Raydium will verify
    pub amm_open_orders: UncheckedAccount<'info>,

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

    // Raydium Pool Program.
    /// CHECK: Raydium will verify
    pub pool_program: UncheckedAccount<'info>,

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

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

impl<'info> AddLiquidity<'info> {
    fn add_liquidity(
        &mut self,
        max_token_a_deposit_amount: u64,
        max_token_b_deposit_amount: u64,
    ) -> Result<()> {
        // read user balances
        let initial_token_a_user_balance =
            account::get_token_balance(&self.token_a_funding_account.to_account_info())?;
        let initial_token_b_user_balance =
            account::get_token_balance(&self.token_b_funding_account.to_account_info())?;
        let initial_lp_user_balance =
            account::get_token_balance(&self.lp_funding_account.to_account_info())?;

        // calculate deposit amounts
        let (max_token_a_deposit_amount, max_token_b_deposit_amount) =
            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(),
                max_token_a_deposit_amount,
                max_token_b_deposit_amount,
            )?;

        // Deposit tokens into the pool
        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);
        invariant!(max_token_a_deposit_amount > 0 && max_token_b_deposit_amount > 0,);
        msg!("Fetched user balances.");

        invariant!(self.vault.deposits_allowed);

        raydium::add_liquidity(
            &[
                self.user_account.to_account_info(),
                self.token_a_funding_account.to_account_info(),
                self.token_b_funding_account.to_account_info(),
                self.lp_funding_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(),
            ],
            max_token_a_deposit_amount,
            max_token_b_deposit_amount,
        )?;
        msg!("Finished add liquidity.");

        // // check amounts spent and received
        let tokens_a_spent = account::check_tokens_spent(
            &self.token_a_funding_account.to_account_info(),
            initial_token_a_user_balance,
            max_token_a_deposit_amount,
        )?;
        let tokens_b_spent = account::check_tokens_spent(
            &self.token_b_funding_account.to_account_info(),
            initial_token_b_user_balance,
            max_token_b_deposit_amount,
        )?;
        let lp_tokens_received = account::check_tokens_received(
            &self.lp_funding_account.to_account_info(),
            initial_lp_user_balance,
            1,
        )?;

        // transfer LP tokens to the custody
        msg!(
            "Transfer LP tokens from user. tokens_a_spent: {}, tokens_b_spent: {}, lp_tokens_received: {}",
            tokens_a_spent,
            tokens_b_spent,
            lp_tokens_received
        );

        // update user stats.
        self.user_info_account
            .add_liquidity(tokens_a_spent, tokens_b_spent)?;
        self.vault.add_liquidity(tokens_a_spent, tokens_b_spent)?;
        msg!("Completed add liquidity.");

        Ok(())
    }
}

pub fn handler(ctx: Context<AddLiquidity>, token_a_amount: u64, token_b_amount: u64) -> Result<()> {
    ctx.accounts.add_liquidity(token_a_amount, token_b_amount)?;
    emit!(VaultAddLiquidityEvent {
        vault: ctx.accounts.vault.key(),
        token_a_amount,
        token_b_amount,
    });
    Ok(())
}

impl<'info> Validate<'info> for AddLiquidity<'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 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,);
        assert_keys_eq!(self.lp_funding_account.owner, 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
        );
        assert_keys_eq!(self.lp_funding_account.mint, self.vault.pool_lp_token_mint);

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

        // Check all unchecked accounts for AMM metadata.
        assert_keys_eq!(self.lp_mint, self.vault.pool_lp_token_mint);
        assert_keys_eq!(self.amm_target, self.vault.amm_target);
        assert_keys_eq!(self.serum_market, self.vault.serum_market);
        assert_keys_eq!(self.amm_open_orders, self.vault.amm_open_orders);

        assert_keys_eq!(
            self.pool_lp_custody_account,
            self.vault.pool_lp_custody_account
        );
        assert_keys_eq!(self.pool_program, self.vault.pool_program_id);

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

        Ok(())
    }
}