import Mathlib.Data.List.Sort

abbrev NatSet := { xs : List ℕ // List.Sorted (·<·) xs }

namespace NatSet

def empty : NatSet := ⟨[], List.sorted_nil⟩
def singleton (x : ℕ) : NatSet := ⟨[x], List.sorted_singleton x⟩

def remove (x : ℕ) : NatSet → NatSet
| ⟨[], _⟩ => empty
| ⟨n::ns, h⟩ =>
  let tail := ⟨ns, List.Sorted.of_cons h⟩
  if n = x
  then tail
  else ⟨n::tail, h⟩

def list_insert (x : ℕ) : List ℕ → List ℕ
| [] => [x]
| a::as =>
  if x = a
  then as
  else if x < a
  then x::a::as
  else a::(list_insert x as)

theorem list_insert_mem : b ∈ list_insert x xs → b = x ∨ b ∈ xs := by
  induction xs with simp [list_insert]
  | cons a as ih =>
    split
    next heq =>
      intro hmem
      right; right
      assumption
    next hne =>
      split
      · simp
      · next hnlt =>
          simp [*]
          intro h
          cases h
          · right; left; assumption
          cases ih ‹_›
          · left; assumption
          · right; right; assumption

theorem list_insert_sorted (hs : List.Sorted (·<·) xs) : List.Sorted (·<·) (list_insert x xs) := by
  induction xs with simp [list_insert]
  | cons a as ih =>
    split
    · exact List.Sorted.of_cons hs
    · split
      · exact List.Sorted.cons ‹_› hs
      · have : a < x := by
          rename_i hne hnlt
          apply Nat.lt_of_le_of_ne
          exact Nat.le_of_not_lt hnlt
          exact fun a_1 => hne (Eq.symm a_1)

        simp only [List.sorted_cons]

        constructor
        · intro b hmem
          cases list_insert_mem hmem
          <;> simp_all only [imp_self, List.sorted_cons, and_true, not_lt]
        · simp_all only [imp_self, List.sorted_cons, and_true, not_lt]

def insert (x : ℕ) : NatSet → NatSet
| ⟨[], _⟩ => singleton x
| ⟨n::ns, h⟩ =>
  if heq : x = n
  then ⟨n::ns, h⟩
  else if hlt : x < n
  then ⟨x::n::ns, List.Sorted.cons hlt h⟩
  else
    let tail : NatSet := ⟨ns, List.Sorted.of_cons h⟩
    let ⟨tail, h₂⟩ := tail.insert x
    ⟨
      n::tail,
      by
        have : ∀b ∈ tail, b = x ∨ b ∈ ns := by
          intro b h
          induction tail
          · contradiction
          · simp [*] at *

            sorry
        have : ∀b ∈ ns, n < b := by
          simp only [List.sorted_cons] at h
          exact h.left
        simp only [List.sorted_cons]
        constructor
        · assumption

        -- have : n < x := by
        --   rw [Nat.lt_iff_le_and_ne]
        --   constructor
        --   · simp at hlt
        --     assumption
        --   · intro h
        --     exact heq h.symm
        sorry

    ⟩

instance : EmptyCollection NatSet where
  emptyCollection := empty

instance : Singleton ℕ NatSet where
  singleton := singleton

instance : Insert ℕ NatSet where
  insert := insert

instance : LawfulSingleton ℕ NatSet where
  insert_empty_eq := by simp [Insert.insert, EmptyCollection.emptyCollection, empty, insert,
    Singleton.singleton, implies_true]
