import Mathlib.Tactic

theorem quadratic_residue_mod_3 (n : Nat) : n^2 % 3 = 0 ∨ n^2 % 3 = 1 := by
  have h : ∀ k : Fin 3, k.val^2 % 3 = 0 ∨ k.val^2 % 3 = 1 := by decide
  have periodic : ∀ m : Nat, m^2 % 3 = (m % 3)^2 % 3 := by
    intro m
    rw [← Nat.mod_add_div m 3]
    ring_nf
    simp [Nat.pow_succ, Nat.mul_mod, Nat.add_mod]
  have key := periodic n
  rw [key]
  have h' : (n % 3) < 3 := Nat.mod_lt n (by decide)
  have : ∃ k : Fin 3, k.val = n % 3 := by
    use ⟨n % 3, h'⟩
  obtain ⟨k, hk⟩ := this
  rw [hk]
  exact h k