use anchor_lang::prelude::*;
use anchor_lang::system_program;

declare_id!("7UP91JrXSfLSPkGE6SKLREX6Vm7MmuDbk1m4cUuVr3M7");

const BOND_VAULT_SPACE: usize = 8;

fn checked_stake_totals(total_bonds: u64, total_staked: u64, amount: u64) -> Option<(u64, u64)> {
    Some((
        total_bonds.checked_add(1)?,
        total_staked.checked_add(amount)?,
    ))
}

fn checked_withdrawal_totals(
    total_bonds: u64,
    total_staked: u64,
    amount: u64,
) -> Option<(u64, u64)> {
    Some((
        total_bonds.checked_sub(1)?,
        total_staked.checked_sub(amount)?,
    ))
}

#[program]
pub mod ai_accountability_bonds {
    use super::*;

    pub fn initialize(ctx: Context<Initialize>) -> Result<()> {
        let deployment_authorized = false;
        require!(
            deployment_authorized,
            AccountabilityError::BootstrapDisabled
        );
        let state = &mut ctx.accounts.bonds_state;
        state.authority = ctx.accounts.authority.key();
        state.total_bonds = 0;
        state.total_staked = 0;
        state.created_at = Clock::get()?.unix_timestamp;
        Ok(())
    }

    pub fn stake_bond(ctx: Context<StakeBond>, amount: u64) -> Result<()> {
        require!(amount > 0, AccountabilityError::ZeroStake);

        // Transfer SOL from agent to bond vault PDA
        system_program::transfer(
            CpiContext::new(
                ctx.accounts.system_program.to_account_info(),
                system_program::Transfer {
                    from: ctx.accounts.agent.to_account_info(),
                    to: ctx.accounts.bond_vault.to_account_info(),
                },
            ),
            amount,
        )?;

        let bond = &mut ctx.accounts.bond;
        bond.agent = ctx.accounts.agent.key();
        bond.stake_amount = amount;
        bond.staked_at = Clock::get()?.unix_timestamp;
        bond.slashed = false;
        bond.withdrawn = false;
        bond.vault = ctx.accounts.bond_vault.key();
        bond.state = ctx.accounts.bonds_state.key();

        let state = &mut ctx.accounts.bonds_state;
        let (total_bonds, total_staked) =
            checked_stake_totals(state.total_bonds, state.total_staked, amount)
                .ok_or(AccountabilityError::ArithmeticOverflow)?;
        state.total_bonds = total_bonds;
        state.total_staked = total_staked;

        emit!(BondStaked {
            agent: bond.agent,
            amount,
            staked_at: bond.staked_at,
        });
        Ok(())
    }

    pub fn slash_bond(ctx: Context<SlashBond>) -> Result<()> {
        let _ = ctx;
        // A unilateral authority flag is not an adjudication process and the
        // original program named no governed slash recipient. Reject instead of
        // permanently freezing or misdirecting a participant's lamports until a
        // separately reviewed dispute, decision, and settlement design exists.
        err!(AccountabilityError::SlashingUnavailable)
    }

    /// Return the stake and close both per-agent PDAs. This makes a completed
    /// bond reusable, returns its rent, and keeps aggregate active-bond
    /// accounting aligned with live vaults.
    pub fn withdraw_bond(ctx: Context<WithdrawBond>) -> Result<()> {
        let bond = &mut ctx.accounts.bond;
        require!(!bond.slashed, AccountabilityError::BondSlashed);
        require!(!bond.withdrawn, AccountabilityError::AlreadyWithdrawn);

        let amount = bond.stake_amount;
        bond.withdrawn = true;

        let state = &mut ctx.accounts.bonds_state;
        let (total_bonds, total_staked) =
            checked_withdrawal_totals(state.total_bonds, state.total_staked, amount)
                .ok_or(AccountabilityError::AccountingInvariantViolation)?;
        state.total_bonds = total_bonds;
        state.total_staked = total_staked;

        // The program owns the initialized vault, so this direct lamport movement
        // preserves the vault's rent reserve and cannot spend a system-owned PDA.
        let rent_reserve = Rent::get()?.minimum_balance(BOND_VAULT_SPACE);
        let vault_lamports = **ctx
            .accounts
            .bond_vault
            .to_account_info()
            .try_borrow_lamports()?;
        let agent_lamports = **ctx.accounts.agent.to_account_info().try_borrow_lamports()?;
        let post_withdrawal_vault_balance = vault_lamports
            .checked_sub(amount)
            .ok_or(AccountabilityError::InsufficientVaultBalance)?;
        require!(
            post_withdrawal_vault_balance >= rent_reserve,
            AccountabilityError::InsufficientVaultBalance
        );
        let post_withdrawal_agent_balance = agent_lamports
            .checked_add(amount)
            .ok_or(AccountabilityError::ArithmeticOverflow)?;
        **ctx
            .accounts
            .bond_vault
            .to_account_info()
            .try_borrow_mut_lamports()? = post_withdrawal_vault_balance;
        **ctx
            .accounts
            .agent
            .to_account_info()
            .try_borrow_mut_lamports()? = post_withdrawal_agent_balance;

        emit!(BondWithdrawn {
            agent: bond.agent,
            amount,
            withdrawn_at: Clock::get()?.unix_timestamp,
        });
        Ok(())
    }
}

#[derive(Accounts)]
pub struct Initialize<'info> {
    #[account(
        init,
        payer = authority,
        space = 8 + 32 + 8 + 8 + 8,
        seeds = [b"accountability"],
        bump
    )]
    pub bonds_state: Account<'info, AccountabilityState>,
    #[account(mut)]
    pub authority: Signer<'info>,
    pub system_program: Program<'info, System>,
}

#[derive(Accounts)]
pub struct StakeBond<'info> {
    #[account(
        mut,
        seeds = [b"accountability"],
        bump
    )]
    pub bonds_state: Account<'info, AccountabilityState>,
    #[account(
        init,
        payer = agent,
        space = 8 + 32 + 8 + 8 + 1 + 1 + 32 + 32,
        seeds = [b"acc_bond", agent.key().as_ref()],
        bump
    )]
    pub bond: Account<'info, AccountabilityBond>,
    #[account(
        init,
        payer = agent,
        space = BOND_VAULT_SPACE,
        seeds = [b"vault", agent.key().as_ref()],
        bump
    )]
    pub bond_vault: Account<'info, BondVault>,
    #[account(mut)]
    pub agent: Signer<'info>,
    pub system_program: Program<'info, System>,
}

#[derive(Accounts)]
pub struct SlashBond<'info> {
    #[account(
        mut,
        seeds = [b"accountability"],
        bump,
        has_one = authority
    )]
    pub bonds_state: Account<'info, AccountabilityState>,
    #[account(
        mut,
        seeds = [b"acc_bond", bond.agent.as_ref()],
        bump,
        constraint = bond.state == bonds_state.key() @ AccountabilityError::InvalidBondState
    )]
    pub bond: Account<'info, AccountabilityBond>,
    pub authority: Signer<'info>,
}

#[derive(Accounts)]
pub struct WithdrawBond<'info> {
    #[account(
        mut,
        seeds = [b"accountability"],
        bump
    )]
    pub bonds_state: Account<'info, AccountabilityState>,
    #[account(
        mut,
        close = agent,
        seeds = [b"acc_bond", agent.key().as_ref()],
        bump,
        constraint = bond.agent == agent.key() @ AccountabilityError::Unauthorized,
        constraint = bond.state == bonds_state.key() @ AccountabilityError::InvalidBondState
    )]
    pub bond: Account<'info, AccountabilityBond>,
    #[account(
        mut,
        close = agent,
        seeds = [b"vault", agent.key().as_ref()],
        bump,
        constraint = bond.vault == bond_vault.key() @ AccountabilityError::InvalidBondVault
    )]
    pub bond_vault: Account<'info, BondVault>,
    #[account(mut)]
    pub agent: Signer<'info>,
}

#[account]
pub struct AccountabilityState {
    pub authority: Pubkey,
    pub total_bonds: u64,
    pub total_staked: u64,
    pub created_at: i64,
}

#[account]
pub struct AccountabilityBond {
    pub agent: Pubkey,
    pub stake_amount: u64,
    pub staked_at: i64,
    pub slashed: bool,
    pub withdrawn: bool,
    pub vault: Pubkey,
    pub state: Pubkey,
}

#[account]
pub struct BondVault {}

#[event]
pub struct BondStaked {
    pub agent: Pubkey,
    pub amount: u64,
    pub staked_at: i64,
}

#[event]
pub struct BondSlashed {
    pub agent: Pubkey,
    pub amount: u64,
    pub slashed_at: i64,
}

#[event]
pub struct BondWithdrawn {
    pub agent: Pubkey,
    pub amount: u64,
    pub withdrawn_at: i64,
}

#[error_code]
pub enum AccountabilityError {
    #[msg("Stake amount must be greater than zero")]
    ZeroStake,
    #[msg("Bond has already been slashed")]
    AlreadySlashed,
    #[msg("Bond has already been withdrawn")]
    AlreadyWithdrawn,
    #[msg("Bond has been slashed and cannot be withdrawn")]
    BondSlashed,
    #[msg("Unauthorized")]
    Unauthorized,
    #[msg("Arithmetic overflow")]
    ArithmeticOverflow,
    #[msg("Bond accounting invariant violated")]
    AccountingInvariantViolation,
    #[msg("Bond vault does not contain the required withdrawable balance")]
    InsufficientVaultBalance,
    #[msg("Bond is not linked to the supplied accountability state")]
    InvalidBondState,
    #[msg("Bond is not linked to the supplied vault")]
    InvalidBondVault,
    #[msg("Slashing is disabled until an adjudication and settlement design is approved")]
    SlashingUnavailable,
    #[msg("Bootstrap is disabled until a source-backed partner authority design is approved")]
    BootstrapDisabled,
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn stake_then_withdraw_restores_active_aggregates() {
        let after_stake = checked_stake_totals(0, 0, 42).expect("valid stake");
        assert_eq!(after_stake, (1, 42));
        assert_eq!(
            checked_withdrawal_totals(after_stake.0, after_stake.1, 42),
            Some((0, 0))
        );
    }

    #[test]
    fn lifecycle_accounting_fails_closed_on_overflow_or_underflow() {
        assert_eq!(checked_stake_totals(u64::MAX, 0, 1), None);
        assert_eq!(checked_stake_totals(0, u64::MAX, 1), None);
        assert_eq!(checked_withdrawal_totals(0, 1, 1), None);
        assert_eq!(checked_withdrawal_totals(1, 0, 1), None);
    }
}
