import Mathlib

inductive Node
| mk (val : Int) (is_end : Bool) (left mid right : Option Node)
deriving Inhabited

namespace Node

def val (t : Node) : Int :=
  match t with | Node.mk v _ _ _ _ => v

def is_end (t : Node) : Bool :=
  match t with | Node.mk _ b _ _ _ => b

def left (t : Node) : Option Node :=
  match t with | Node.mk _ _ l _ _ => l

def mid (t : Node) : Option Node :=
  match t with | Node.mk _ _ _ m _ => m

def right (t : Node) : Option Node :=
  match t with | Node.mk _ _ _ _ r => r

end Node

def view (t : Node) : Set (List Int) :=
  let l := match t.left with | some child => view child | none => ∅
  let r := match t.right with | some child => view child | none => ∅
  let m := match t.mid with | some child => (view child).image (fun s => t.val :: s) | none => ∅
  let current := if t.is_end then {[t.val]} else ∅
  l ∪ r ∪ m ∪ current
termination_by sizeOf t
decreasing_by all_goals sorry

def is_empty_node (t : Node) : Bool :=
  !t.is_end ∧ t.left.isNone ∧ t.mid.isNone ∧ t.right.isNone

def well_formed (t : Node) : Prop :=
  (0 ≤ t.val ∧ t.val < 256) ∧
  -- Check Left
  (match t.left with 
   | some l => 
       well_formed l ∧ 
       !is_empty_node l ∧ -- CRITICAL: Match Dafny pruning
       (∀ s ∈ view l, match s with | [] => False | c::_ => c < t.val) 
   | none => True) ∧
  -- Check Right
  (match t.right with 
   | some r => 
       well_formed r ∧ 
       !is_empty_node r ∧ -- CRITICAL
       (∀ s ∈ view r, match s with | [] => False | c::_ => c > t.val) 
   | none => True) ∧
  -- Check Mid
  (match t.mid with 
   | some m => 
       well_formed m ∧ 
       !is_empty_node m -- CRITICAL
   | none => True)
termination_by sizeOf t
decreasing_by all_goals sorry

def contains (t : Node) (key : List Int) : Bool :=
  match key with
  | [] => False -- TST nodes always represent at least one char
  | c :: cs =>
    if c < t.val then
      match t.left with | some l => contains l key | none => False
    else if c > t.val then
      match t.right with | some r => contains r key | none => False
    else -- c == t.val
      match cs with
      | [] => t.is_end
      | _ => match t.mid with | some m => contains m cs | none => False
termination_by sizeOf t
decreasing_by all_goals sorry

def is_valid_key (key : List Int) : Prop :=
  ∀ k ∈ key, 0 ≤ k ∧ k < 256

-- Precondition definitions
@[reducible, simp]
def search_precond (t : Option Node) (key : List Int) : Prop :=
  -- !benchmark @start precond
  (match t with
   | some node => well_formed node ∧ !is_empty_node node
   | none => True) ∧
  is_valid_key key -- Match Dafny input constraint
  -- !benchmark @end precond

-- !benchmark @start auxcode
-- !benchmark @end auxcode

-- Main function definition
def search (t : Option Node) (key : List Int)
    (h_precond : search_precond t key) : Bool :=
  -- !benchmark @start code
  sorry
  -- !benchmark @end code

-- View wrapper for Option Node
def view_opt (t : Option Node) : Set (List Int) :=
  match t with
  | some node => view node
  | none => ∅

-- Postcondition definitions
@[reducible, simp]
def search_postcond (t : Option Node) (key : List Int) (result : Bool)
    (_ : search_precond t key) : Prop :=
  -- !benchmark @start postcond
  result = (key ∈ view_opt t)
  -- !benchmark @end postcond

-- !benchmark @start lemma
-- !benchmark @end lemma

-- Proof content
theorem search_postcond_satisfied (t : Option Node) (key : List Int)
    (h_precond : search_precond t key) :
    search_postcond t key (search t key h_precond) h_precond := by
  -- !benchmark @start proof
  sorry
  -- !benchmark @end proof