import Mathlib.Tactic

theorem quadratic_residue_mod_3 (n : Nat) : n^2 % 3 = 0 ∨ n^2 % 3 = 1 := by
  have h : n % 3 = 0 ∨ n % 3 = 1 ∨ n % 3 = 2 := by omega
  rcases h with (h | h | h)
  · -- Case n % 3 = 0
    rw [← Nat.mod_add_div n 3, h]
    ring_nf
    simp
    left
    omega
  · -- Case n % 3 = 1
    rw [← Nat.mod_add_div n 3, h]
    ring_nf
    simp
    right
    omega
  · -- Case n % 3 = 2
    rw [← Nat.mod_add_div n 3, h]
    ring_nf
    simp
    right
    omega