refactor!: update staking code #2

[sainoe]

Description

Closes: #XXXX

---

Author Checklist

All items are required. Please add a note to the item if the item is not applicable and
please add links to any relevant follow up issues.

I have...

• included the correct type prefix in the PR title
• confirmed ! in the type prefix if API or client breaking change
• targeted the correct branch (see PR Targeting)
• provided a link to the relevant issue or specification
• reviewed "Files changed" and left comments if necessary
• included the necessary unit and integration tests
• added a changelog entry to CHANGELOG.md
• updated the relevant documentation or specification, including comments for documenting Go code
• confirmed all CI checks have passed

Reviewers Checklist

All items are required. Please add a note if the item is not applicable and please add
your handle next to the items reviewed if you only reviewed selected items.

I have...

• confirmed the correct type prefix in the PR title
• confirmed all author checklist items have been addressed
• reviewed state machine logic, API design and naming, documentation is accurate, tests and test coverage