use anchor_lang::prelude::*;

pub mod errors;
pub mod solver;
pub mod state;

use errors::RuneError;
use state::*;

rune_wallet!();

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

    /// Allocate the rig and write the skeleton it will spend the rest of its
    /// life inside. Bone lengths and joint limits are fixed at this point and
    /// there is no instruction anywhere in this program that changes them.
    pub fn initialize(ctx: Context<Initialize>, bone: [i32; JOINT_COUNT]) -> Result<()> {
        let mut rig = ctx.accounts.rig.load_init()?;
        rig.authority = ctx.accounts.authority.key();
        rig.bone = bone;
        rig.angle = REST;
        rig.lower = LOWER;
        rig.upper = UPPER;
        rig.root = ROOT;
        rig.target = ROOT;
        rig.residual = i64::MAX;
        rig.prior = i64::MAX;
        rig.iterations = 0;
        rig.steps = 0;
        rig.slot = Clock::get()?.slot;
        rig.stall = 0;
        rig.sealed = 0;
        rig.bump = ctx.bumps.rig;

        for j in 0..JOINT_COUNT {
            require!(rig.bone[j] > 0, RuneError::DegenerateBone);
            require!(rig.lower[j] < rig.upper[j], RuneError::InvertedLimit);
            require!(rig.angle[j] >= rig.lower[j], RuneError::RestOutsideLimit);
            require!(rig.angle[j] <= rig.upper[j], RuneError::RestOutsideLimit);
        }

        msg!("rig live, joints {}, chains {}", JOINT_COUNT, CHAIN_COUNT);
        Ok(())
    }

    /// Move the target for one chain. The target is planar and is expressed in
    /// the same fixed point units as every other length in the account.
    pub fn aim(ctx: Context<Step>, chain: u8, x: i32, y: i32) -> Result<()> {
        let mut rig = ctx.accounts.rig.load_mut()?;
        require!(rig.sealed == 0, RuneError::Sealed);
        require!((chain as usize) < CHAIN_COUNT, RuneError::ChainRange);
        let c = chain as usize;
        let dx = (x - rig.root[c * 2]) as i64;
        let dy = (y - rig.root[c * 2 + 1]) as i64;
        let reach = chain_reach(&rig, c);
        require!(dx * dx + dy * dy <= reach * reach, RuneError::OutOfReach);
        rig.target[c * 2] = x;
        rig.target[c * 2 + 1] = y;
        rig.prior = i64::MAX;
        rig.stall = 0;
        msg!("aim chain {} to {} {}", chain, x, y);
        Ok(())
    }

    /// Run the solver. Every pass walks all four chains from the effector back
    /// to the root, so the cost of this instruction is passes * chains * bones
    /// and the ceiling on passes is set by the compute budget, not by taste.
    pub fn step(ctx: Context<Step>, passes: u8) -> Result<()> {
        let mut rig = ctx.accounts.rig.load_mut()?;
        require!(rig.sealed == 0, RuneError::Sealed);
        require!(passes > 0 && passes <= MAX_PASSES, RuneError::PassRange);

        let slot = Clock::get()?.slot;
        require!(slot > rig.slot, RuneError::SameSlot);

        let mut total: i64 = 0;
        for _ in 0..passes {
            total = 0;
            for c in 0..CHAIN_COUNT {
                total = total.saturating_add(solver::pass(&mut rig, c));
            }
            rig.iterations = rig.iterations.saturating_add(1);
        }

        solver::carry(&mut rig);

        rig.prior = rig.residual;
        rig.residual = total;
        rig.steps = rig.steps.saturating_add(1);
        rig.slot = slot;

        if rig.prior != i64::MAX && rig.residual.saturating_add(RESIDUAL_EPS) >= rig.prior {
            rig.stall = rig.stall.saturating_add(1);
        } else {
            rig.stall = 0;
        }

        if rig.stall >= STALL_LIMIT {
            rig.sealed = 1;
            msg!("sealed, residual {}, iterations {}", rig.residual, rig.iterations);
        }

        msg!(
            "step {} passes {} residual {} stall {} slot {}",
            rig.steps,
            passes,
            rig.residual,
            rig.stall,
            slot
        );
        Ok(())
    }
}

fn chain_reach(rig: &Rig, chain: usize) -> i64 {
    let mut sum: i64 = 0;
    for k in 0..CHAIN_LEN {
        sum += rig.bone[CHAIN[chain][k]] as i64;
    }
    sum
}

#[derive(Accounts)]
pub struct Initialize<'info> {
    #[account(
        init,
        payer = authority,
        space = Rig::LEN,
        seeds = [b"rig"],
        bump
    )]
    pub rig: AccountLoader<'info, Rig>,
    #[account(mut)]
    pub authority: Signer<'info>,
    pub system_program: Program<'info, System>,
}

#[derive(Accounts)]
pub struct Step<'info> {
    #[account(mut, seeds = [b"rig"], bump = rig.load()?.bump)]
    pub rig: AccountLoader<'info, Rig>,
    pub caller: Signer<'info>,
}