Pinocchio Programs

use light_account_pinocchio::{
    prepare_compressed_account_on_init, CompressedCpiContext, CpiAccounts, CpiAccountsConfig,
    CpiContextWriteAccounts, CreateMints, CreateMintsStaticAccounts, CreateTokenAccountCpi,
    InstructionDataInvokeCpiWithAccountInfo, InvokeLightSystemProgram, LightAccount, LightConfig,
    SingleMintParams,
};
use pinocchio::sysvars::{clock::Clock, Sysvar};

pub fn process(
    ctx: &InitializeAccounts<'_>,
    params: &InitializeParams,
    remaining_accounts: &[AccountInfo],
) -> Result<(), LightSdkTypesError> {
    // 1. Build CPI accounts
    let config = CpiAccountsConfig::new_with_cpi_context(crate::LIGHT_CPI_SIGNER);
    let cpi_accounts = CpiAccounts::new_with_config(
        ctx.payer,
        &remaining_accounts[params.create_accounts_proof.system_accounts_offset as usize..],
        config,
    );

    // 2. Get address tree info + config
    let address_tree_info = &params.create_accounts_proof.address_tree_info;
    let address_tree_pubkey = address_tree_info.get_tree_pubkey(&cpi_accounts)?;
    let light_config = LightConfig::load_checked(ctx.compressible_config, &crate::ID)?;
    let current_slot = Clock::get()?.slot;

    // 3. Create pool PDA (write to CPI context)
    {
        let cpi_context = CompressedCpiContext::first();
        let mut new_address_params = Vec::with_capacity(1);
        let mut account_infos = Vec::with_capacity(1);
        let pool_key = *ctx.pool.key();

        prepare_compressed_account_on_init(
            &pool_key, &address_tree_pubkey, address_tree_info,
            params.create_accounts_proof.output_state_tree_index,
            0, &crate::ID,
            &mut new_address_params, &mut account_infos,
        )?;

        // Initialize pool state (zero-copy)
        {
            let mut data = ctx.pool.try_borrow_mut_data()?;
            let pool_state: &mut PoolState = bytemuck::from_bytes_mut(
                &mut data[8..8 + core::mem::size_of::<PoolState>()]
            );
            pool_state.set_decompressed(&light_config, current_slot);
            pool_state.token_a_mint = *ctx.mint_a().key();
            pool_state.token_b_mint = *ctx.mint_b().key();
            // ... remaining fields
        }

        // Write to CPI context
        let instruction_data = InstructionDataInvokeCpiWithAccountInfo {
            mode: 1,
            bump: crate::LIGHT_CPI_SIGNER.bump,
            invoking_program_id: crate::LIGHT_CPI_SIGNER.program_id.into(),
            proof: params.create_accounts_proof.proof.0,
            new_address_params,
            account_infos,
            // ...
        };
        instruction_data.invoke_write_to_cpi_context_first(
            CpiContextWriteAccounts {
                fee_payer: cpi_accounts.fee_payer(),
                authority: cpi_accounts.authority()?,
                cpi_context: cpi_accounts.cpi_context()?,
                cpi_signer: crate::LIGHT_CPI_SIGNER,
            }
        )?;
    }

    // 4. Create mints
    CreateMints { /* ... */ }.invoke(&cpi_accounts)?;

    // 5. Create vaults (rent-free)
    CreateTokenAccountCpi { /* ... */ }.rent_free(/* ... */).invoke_signed(/* ... */)?;

    Ok(())
}

use light_client::interface::{
    AccountInterface, AccountSpec, ColdContext, LightProgramInterface, PdaSpec,
};
use light_account::token::Token;
use pinocchio_swap::{LightAccountVariant, PoolState, PoolStateSeeds, VaultSeeds};

/// Flat SDK struct. All fields populated at construction from pool state data.
pub struct SwapSdk {
    pub pool_state_pubkey: Pubkey,
    pub token_a_mint: Pubkey,
    pub token_b_mint: Pubkey,
    pub token_a_vault: Pubkey,
    pub token_b_vault: Pubkey,
    pub pool_authority: Pubkey,
}

impl SwapSdk {
    pub fn new(pool_state_pubkey: Pubkey, pool_data: &[u8]) -> Result<Self, SwapSdkError> {
        let pool = PoolState::deserialize(&mut &pool_data[8..])?;
        // ... derive addresses from pool state
        Ok(Self { pool_state_pubkey, /* ... */ })
    }
}

impl LightProgramInterface for SwapSdk {
    type Variant = LightAccountVariant;
    type Instruction = SwapInstruction;

    fn program_id() -> Pubkey { PROGRAM_ID }

    fn instruction_accounts(&self, ix: &Self::Instruction) -> Vec<Pubkey> {
        match ix {
            SwapInstruction::Swap => vec![
                self.pool_state_pubkey,
                self.token_a_vault,
                self.token_b_vault,
                self.token_a_mint,
                self.token_b_mint,
            ],
            // ...
        }
    }

    fn load_specs(
        &self,
        cold_accounts: &[AccountInterface],
    ) -> Result<Vec<AccountSpec<Self::Variant>>, Box<dyn std::error::Error>> {
        let mut specs = Vec::new();
        for account in cold_accounts {
            if account.key == self.pool_state_pubkey {
                let pool = PoolState::deserialize(&mut &account.data()[8..])?;
                let variant = LightAccountVariant::PoolState {
                    seeds: PoolStateSeeds { /* ... */ },
                    data: pool,
                };
                specs.push(AccountSpec::Pda(PdaSpec::new(account.clone(), variant, PROGRAM_ID)));
            } else if account.key == self.token_a_vault {
                let token: Token = Token::deserialize(&mut &account.data()[..])?;
                let variant = LightAccountVariant::Vault(TokenDataWithSeeds {
                    seeds: VaultSeeds { pool: /* ... */, mint: /* ... */ },
                    token_data: token,
                });
                specs.push(AccountSpec::Pda(PdaSpec::new(account.clone(), variant, PROGRAM_ID)));
            }
            // ... token_b_vault, mints
        }
        Ok(specs)
    }
}

use light_program_test::{LightProgramTest, Rpc};
use light_client::interface::{
    create_load_instructions, get_create_accounts_proof,
    AccountSpec, CreateAccountsProofInput, LightProgramInterface,
};


#[tokio::test]
async fn test_pool_lifecycle() {
    let mut rpc = LightProgramTest::new(config).await.unwrap();

    // 1. Initialize pool (rent-free: pool PDA, 2 mints, 2 vaults)
    let proof = get_create_accounts_proof(&rpc, &program_id, vec![
        CreateAccountsProofInput::pda(pool_state),
        CreateAccountsProofInput::mint(mint_a_signer),
        CreateAccountsProofInput::mint(mint_b_signer),
    ]).await.unwrap();

    rpc.create_and_send_transaction(&[init_ix], &payer.pubkey(), &[&payer, &authority])
        .await.unwrap();

    // 2. Swap (hot path)
    rpc.create_and_send_transaction(&[swap_ix], &user.pubkey(), &[&user])
        .await.unwrap();

    // 3. Trigger compression for the purpose of the test.
    const SLOTS_PER_EPOCH: u64 = 13500;
    rpc.warp_slot_forward(SLOTS_PER_EPOCH * 30).await.unwrap();

    // 4. Build SDK from pool state, fetch cold accounts
    let pool_iface = rpc.get_account_interface(&pool_state, None).await.unwrap().value.unwrap();
    assert!(pool_iface.is_cold());

    let sdk = SwapSdk::new(pool_state, pool_iface.data()).unwrap();
    let pubkeys = sdk.instruction_accounts(&SwapInstruction::Swap);
    let accounts = rpc.get_multiple_account_interfaces(pubkeys.iter().collect(), None)
        .await.unwrap().value;
    let cold: Vec<_> = accounts.into_iter().flatten().filter(|a| a.is_cold()).collect();

    // 5. Load cold accounts
    let mut specs = sdk.load_specs(&cold).unwrap();
    // Add user ATAs
    let ata_a = rpc.get_associated_token_account_interface(&user.pubkey(), &mint_a, None)
        .await.unwrap().value.unwrap();
    let ata_b = rpc.get_associated_token_account_interface(&user.pubkey(), &mint_b, None)
        .await.unwrap().value.unwrap();
    specs.push(AccountSpec::Ata(ata_a));
    specs.push(AccountSpec::Ata(ata_b));

    let load_ixs = create_load_instructions(&specs, payer.pubkey(), config_pda, &rpc)
        .await.unwrap();


    // 6. Load and Swap
    let mut all_ixs = load_ixs;
    all_ixs.push(swap_ix);
    rpc.create_and_send_transaction(&all_ixs, &user.pubkey(), &[&user])
        .await.unwrap();
}