How many square roots can a complex matrix with distinct non zero eigenvalues can take?

I got the above problem in an exam, the problem stated to show there is at most $2^n$ such matrices.

What I did was to show the matrix is similar to a matrix with a diagonal matrix with distinct complex numbers, and now if $A$ is such a matrix and $B$ is a square root. Then $B$ is diagonalizable so since $A$ and $B$ commute there is a matrix $P$ such that that $P^{-1}BP$ and $P^{-1}AP$ are both are diagonal.

So my argument was that to consider the diagonal case (w.l.g) and show that each diagonal entry of $P^{-1}BP$ should be root of the corresponding diganoal entry of $P^{-1}AP$. So we can have atmost 2 options for each and overall we get $2^n$.

Now the problem I just realized is that I just saw that for any invertible diagonalizable matrix $A^2$, $A$ is also diagonalizable. So then we can apply the same logic as what I just did to any such invertible matrix. But I also know the identity has infinitely many roots. So where did I go wrong?

Thank You

It seems to me that you missed one subtlety when you wrote "w.l.g". If $A=PDP^{-1}$ and $E^2=D$, then $(PEP^{-1})^2=A$. But who says that there are no other invertible $Q$ with $(QEQ^{-1})^2=A$? That's precisely what happens here.
When you know that all diagonal elements in $D$ are distinct, the following happens: if $(QEQ^{-1})^2=A=(PEP^{-1})^2$, we can rewrite this as $$(Q^{-1}P) D=D(Q^{-1}P).$$ Because all diagonal entries of $D$ are distinct, this implies that $Q^{-1}P=I$, that is $Q=P$. That's why you get precisely $2^n$ square roots.
When eigenvalues are repeated, this goes off the window. For $I$, you can take the $2^n$ matrices $D_k$ with $1$ and $-1$ in the diagonal, so $D_k^2=I$. But now, for any invertible $P$, the matrixi $PD_kP^{-1}$ is a root, and distinct $P$ will give us mostly distinct matrices; that's how we get infinitely many roots.