tentacles 0.1.1

use anchor_lang::prelude::*;
use anchor_spl::token::{self, Transfer};
use crate::state::{MemberInfo, InitializeWallet, AddMember, Distribute};
use crate::utils::{assert_ata, assert_ata_for_pubkey, assert_ata_for_keys};
use crate::error::TentaClesError;

pub fn initialize_wallet(ctx: Context<InitializeWallet>, name: String, total_shares: u64, bump: u8) -> Result<()> {
    let wallet = &mut ctx.accounts.wallet;
    let wallet_key = wallet.key();  // Save wallet key to a variable
    // Assert that the provided token account is the associated token account
    assert_ata(
        &ctx.accounts.wallet_token_account.to_account_info(),
        &wallet_key,
        &ctx.accounts.mint.key(),
        Some(TentaClesError::InvalidAccountData),
    )?;

    wallet.name = name;
    wallet.total_shares = total_shares;
    wallet.total_inflow = 0;
    wallet.last_inflow = 0;
    wallet.disburse_cycles = 0;
    wallet.total_members = 0;
    wallet.tentacles = wallet.to_account_info().key();
    wallet.authority = *ctx.accounts.authority.key;
    wallet.remaining_flow = ctx.accounts.wallet_token_account.amount;
    wallet.mint = ctx.accounts.mint.to_account_info().key();
    wallet.bump_seed = bump;
    wallet.token_account = ctx.accounts.wallet_token_account.to_account_info().key();
    wallet.total_available_shares = total_shares;
    wallet.members = Vec::with_capacity(5);
    Ok(())
}

pub fn add_member(ctx: Context<AddMember>, shares: u64) -> Result<()> {
        let wallet = &mut ctx.accounts.wallet;

        // Check that the number of members does not exceed 5
        if wallet.total_members >= 5 {
            return Err(TentaClesError::MaxMembersReached.into());
        }

        // Check that there are enough shares available
        if wallet.total_available_shares < shares {
            return Err(TentaClesError::NotEnoughAvailableShares.into());
        }

        // Check if the member is already in the wallet
        let member_pubkey = *ctx.accounts.member.key;
        if wallet.members.iter().any(|m| m.member == member_pubkey) {
            return Err(TentaClesError::DoubleMember.into());
        }

        // Borrow mint before mutable borrow of wallet
        let mint = wallet.mint;
        let member_ata = ctx.accounts.member_token_account.key();

        // Assert that the mint of the wallet matches the mint of the member token account
        //To-do => handles this assertion properly
        assert_ata_for_pubkey(&member_ata, &*ctx.accounts.member.key, &mint, Some(TentaClesError::MintMismatch))?;

        let member_info = MemberInfo {
            member: *ctx.accounts.member.key,
            shares,
            disburse_cycles: 0,
            member_token_account: ctx.accounts.member_token_account.key(),
        };

        wallet.total_members += 1;
        wallet.total_available_shares -= shares;
        wallet.members.push(member_info);

        Ok(())
    }

    /*
     during distribution by any member
     - snapshot after transfer - remaining to distribute

     - increment user disburse cycle by 1
     - have wallet current disburse cycle field

 

     - distribution
     - distribution state if needed processed, processing, null
     */
pub fn distribute(ctx: Context<Distribute>, member_pubkey: Pubkey) -> Result<()> {
    let wallet = &mut ctx.accounts.wallet;
    let wallet_token_account = &ctx.accounts.wallet_token_account;
    let member_token_account = &ctx.accounts.member_token_account;
    let token_program = &ctx.accounts.token_program;
    let pda_as_authority = &ctx.accounts.pda_as_authority;

    // Fetch the current balance of the wallet's token account
    let current_balance = wallet_token_account.amount;
    msg!("Current balance before: {}", current_balance);

    // Check if the wallet token address passed in is the correct one
    assert_ata_for_keys(&wallet_token_account.key(), &wallet.key(), &wallet.mint, Some(TentaClesError::MintMismatch))?;

    // Calculate the total shares and validate members
    let total_shares = wallet.total_shares;
    if total_shares == 0 {
        return Err(TentaClesError::NotEnoughAvailableShares.into());
    }

    // Find the member index in the wallet
    let member_index = wallet.members.iter().position(|m| m.member == member_pubkey)
        .ok_or_else(|| TentaClesError::InvalidAccountData)?;

    // Ensure no member claims twice for a cycle
    if wallet.members[member_index].disburse_cycles == wallet.disburse_cycles && wallet.remaining_flow != 0 {
        return Err(TentaClesError::UnclaimedAmountExists.into());
    }

    // Calculate the amount to distribute to the member
    let amount_to_distribute;
    if wallet.remaining_flow == 0 {
        // First member claiming in this cycle
        amount_to_distribute = (current_balance * wallet.members[member_index].shares) / total_shares;

        // Update last inflow and remaining flow
        wallet.last_inflow = current_balance;
        wallet.remaining_flow = current_balance - amount_to_distribute;

        wallet.total_inflow += current_balance;

        // Update wallet's disburse cycles
        wallet.disburse_cycles += 1;
    } else {
        // Subsequent member claiming in this cycle
        amount_to_distribute = (wallet.last_inflow * wallet.members[member_index].shares) / total_shares;

        // Update remaining flow
        wallet.remaining_flow -= amount_to_distribute;
    }
    msg!("Amount to distribute: {}", amount_to_distribute);

    // Perform the token transfer
    {
        let wallet_name_as_bytes = wallet.name.as_bytes();
        let wallet_seeds = &[b"split_wallet".as_ref(), wallet_name_as_bytes, &[wallet.bump_seed]];
        let wallet_signer = &[&wallet_seeds[..]];

        let cpi_accounts = Transfer {
            from: wallet_token_account.to_account_info(),
            to: member_token_account.to_account_info(),
            authority: pda_as_authority.to_account_info(),
        };
        let cpi_context = CpiContext::new_with_signer(token_program.to_account_info(), cpi_accounts, wallet_signer);
        token::transfer(cpi_context, amount_to_distribute)?;
    }

    let balance_after = wallet_token_account.amount;
    msg!("Balance After: {}", balance_after);

    // Update the member's disburse cycles and unclaimed amount
    {
        let member_info = &mut wallet.members[member_index];

        member_info.disburse_cycles += 1;
    }

    // If this is the last member in the cycle, reset the remaining flow to zero
    let all_members_claimed = wallet.members.iter().all(|m| m.disburse_cycles == wallet.disburse_cycles);
    if all_members_claimed {
        wallet.remaining_flow = 0;
    } else {
        // Update the wallet's remaining flow
        wallet.remaining_flow = balance_after;
    }

    Ok(())
}