The number of polynomial functions $f:A\to A$ is $|A|^2$ if and only if $x^2=x$ for all $x\in A$.

Your proof of the implication $a)\implies b)$ is lacking; you prove that every polynomial function $f:\ A\ \longrightarrow\ A$ is of the form $f(x)=bx+a_0$ for some $b,a_0\in A$. This proves that there are at most $n^2$ polynomial functions $f:\ A\ \longrightarrow\ A$. You should show that they are all distinct.

For the converse; again show that the linear functions are all distinct. This yields $n^2$ distinct polynomial functions $f:\ A\ \longrightarrow\ A$ of the form $f(x)=bx+a_0$. So all polynomial functions are of this form. In particular ther exist $b,a_0\in A$ such that $x^2=bx+a_0$ for all $x\in A$. Plugging in $x=0$ and $x=1$ shows that $a_0=0$ and $b=1$, so $x^2=x$ for all $x\in A$.

---

Here's a more abstract alternative, that doesn't turn out to be any easier at all, but perhaps it offers another perspective:

Let $B:=\operatorname{Map}(A,A)$ denote the set of functions $f:\ A\ \longrightarrow\ A$. Note that $B$ is a group with respect to pointwise addition. Polynomials with coefficients in $A$ induce functions from $A$ to $A$ by evaluation, yielding a group homomorphism $$\varepsilon:\ A[X]\ \longrightarrow\ B:\ f\ \longmapsto\ (x\ \mapsto\ f(x)).$$ Because $x^2=x$ for all $x\in A$ we see that $\langle X^2-X\rangle\in\ker\varepsilon$, and it is not hard to see that in fact the ideal $(X^2-X)\subset A[X]$ is contained in $\ker\varepsilon$. To see that $\ker\varepsilon=(X^2-X)$ it suffices to note that a linear polynomial $bX+a\in A[X]$ induces the zero map if and only if $a=b=0$; indeed from $$\forall x\in A:\ bx+a=0,$$ it immediately follows that $a=0$ and $b=0$ by plugging in $x=0$ and $x=1$. Hence by the first isomorphism theorem the image of $\varepsilon$ is isomorphic to $A[X]/(X^2-X)$ which is a ring of $n^2$ elements.