In fact you do not need to assume that $A$ is diagonalisable, that is:
Claim: If $B$ is a diagonalisable $n\times n$ matrix with $n$ distinct eigenvalues and $AB=BA$, then $A$ there is a polynomial $f(t)$ such that $A = f(B)$. In particular, $A$ is also diagonalisable.
Proof: If $B$ is diagonalisable, then there is an invertible matrix $P$ such that $PBP^{-1} =D$, where $D$ is a diagonalisable matrix. Assuming that $D = \text{diag}(d_1,\ldots,d_n)$ has distinct diagonal entries, that is, $d_i \neq d_j$ if $i \neq j$, it follows that the characteristic polynomial of $B$ is also its minimal polynomial $f(t)=\prod_{i=1}^n(t-d_i)$ so that $\mathbb C[B]$, the $\mathbb C$-algebra generated by $B$ inside $\mathrm{Mat}_n(\mathbb C)$ is isomorphic to $\mathbb C[t]/\langle f \rangle$, and hence in particular $n$-dimensional.
On the other hand, since $PBP^{-1} = D$ and so $P\mathbb C[B]P^{-1} = \mathbb C[D]$. But since $D$ is diagonal
$$
\mathbb C[D] \subseteq T= \{X \in \mathrm{Mat}_n(\mathbb C)\mid x_{ij} = 0, \forall i \neq j\}
$$
which is also $n$-dimensional, it follows that $\mathbb C[D] = T$. For any matrix $C$ we have $DC=CD$ if and only if $(d_i-d_j)c_{ij} = 0$ and since $(d_i-d_j)\neq 0$ when $i\neq j$ it follows that $C$ is diagonal. Since any two diagonal matrices clearly commute, it follows that
$$
\{C \in \mathrm{Mat}_n(\mathbb C): CD=DC\} = T = \mathbb C[D]
$$
Now if $AB=BA$, then if $C=PAP^{-1}$ we see that $CD=DC$ and hence $C \in \mathbb C[D]$. Since $P^{-1}\mathbb C[D]P = \mathbb C[B]$, it follows that $A \in \mathbb C[B]$.