-- AVM ISA v1 (Lean-only): Type Safety Proofs

import SilverSight.AVMIsa.TypeCheck
import SilverSight.AVMIsa.Step

namespace SilverSight.AVMIsa

/-- Relation asserting that stack values match the expected stack type signature. -/
inductive StackMatches : List AnyVal → List AvmTy → Prop where
  | nil : StackMatches [] []
  | cons {ty : AvmTy} (val : AvmVal ty) {vs : List AnyVal} {tys : List AvmTy}
      (h : StackMatches vs tys) : StackMatches (⟨ty, val⟩ :: vs) (ty :: tys)

/-- Relation asserting that a local frame matches the expected local types. -/
inductive OptionMatches : Option AnyVal → Option AvmTy → Prop where
  | none : OptionMatches none none
  | some {ty : AvmTy} (val : AvmVal ty) : OptionMatches (some ⟨ty, val⟩) (some ty)

def LocalsMatches (locals : List (Option AnyVal)) (signatures : List (Option AvmTy)) : Prop :=
  locals.length = signatures.length ∧
  ∀ i, OptionMatches (locals.getD i none) (signatures.getD i none)

inductive SafeOutcome : Outcome State → Prop where
  | ok (s' : State) : SafeOutcome (Outcome.ok s')
  | divByZero : SafeOutcome (Outcome.err StepError.divisionByZero)
  | stackOverflow : SafeOutcome (Outcome.err StepError.stackOverflow)

theorem StackMatches.length_eq {vs : List AnyVal} {tys : List AvmTy} (h : StackMatches vs tys) :
    vs.length = tys.length := by
  induction h with
  | nil => rfl
  | cons _ _ ih => simp [ih]

/-- Helper: pop1 matches. -/
theorem pop1_safety {vs : List AnyVal} {ty : AvmTy} {tys : List AvmTy} (h : StackMatches vs (ty :: tys)) :
    match vs with
    | [] => False
    | x :: xs => x.ty = ty ∧ StackMatches xs tys := by
  cases h
  exact ⟨rfl, by assumption⟩

/-- Locals lookup safety helper. -/
theorem getLocal_safety {locals : List (Option AnyVal)} {signatures : List (Option AvmTy)}
    (h : LocalsMatches locals signatures) (i : Nat) {ty : AvmTy}
    (hloc : signatures[i]? = some (some ty)) :
    ∃ val : AvmVal ty, getLocal? { pc := 0, stack := [], locals := locals, halted := false } i = some ⟨ty, val⟩ := by
  obtain ⟨hlen, hforall⟩ := h
  have hsig_eq : signatures.getD i none = some ty := by
    rw [List.getD_eq_getElem?_getD]
    simp [hloc]
  have hmatches := hforall i
  rw [hsig_eq] at hmatches
  generalize hopt : locals.getD i none = opt at hmatches
  cases hmatches with
  | some val =>
      use val
      unfold getLocal?
      exact hopt

/-- Locals update safety helper. -/
theorem setLocal_safety {locals : List (Option AnyVal)} {signatures : List (Option AvmTy)}
    (h : LocalsMatches locals signatures) (i : Nat) {ty : AvmTy} (val : AvmVal ty)
    (hbound : i < signatures.length) :
    LocalsMatches (List.set locals i (some ⟨ty, val⟩)) (List.set signatures i (some ty)) := by
  obtain ⟨hlen, hforall⟩ := h
  constructor
  · simp [hlen]
  · intro j
    by_cases hj : j = i
    · rw [hj]
      have hbound_locals : i < locals.length := by omega
      rw [List.getD_eq_getElem?_getD, List.getElem?_set_self hbound_locals]
      rw [List.getD_eq_getElem?_getD, List.getElem?_set_self hbound]
      exact OptionMatches.some val
    · -- j ≠ i
      have hneq : i ≠ j := by omega
      rw [List.getD_eq_getElem?_getD, List.getElem?_set_ne hneq]
      rw [List.getD_eq_getElem?_getD, List.getElem?_set_ne hneq]
      exact hforall j

/-- Safety of checkPrim: if checkPrim succeeds, evalPrim is safe and type-matches. -/
theorem evalPrim_safety {p : Prim} {s : State} {tys : List AvmTy}
    {tys' : List AvmTy} (hstack : StackMatches s.stack tys)
    (hprim : checkPrim p tys = some tys') :
    ∃ o, evalPrim p s = o ∧ SafeOutcome o ∧ (∀ s', o = Outcome.ok s' → StackMatches s'.stack tys' ∧ s'.locals = s.locals) := by
  rcases s with ⟨pc, stack, locals, halted⟩
  cases p <;> cases tys <;> try (unfold checkPrim at hprim; dsimp at hprim; contradiction)
  · -- addSatQ0
    rename_i ty tys_tail
    cases ty <;> try (unfold checkPrim at hprim; dsimp at hprim; contradiction)
    cases tys_tail
    · unfold checkPrim at hprim; dsimp at hprim; contradiction
    · rename_i ty2 tys_tail'
      cases ty2 <;> try (unfold checkPrim at hprim; dsimp at hprim; contradiction)
      cases hstack with
      | cons val1 h1 =>
          cases h1 with
          | cons val2 h2 =>
              cases val1; cases val2
              unfold checkPrim at hprim; dsimp at hprim
              injection hprim with heq_tys'
              unfold evalPrim pop1 push1
              dsimp
              split
              · refine ⟨_, rfl, ⟨SafeOutcome.stackOverflow, by intro s' h; contradiction⟩⟩
              · refine ⟨_, rfl, ⟨SafeOutcome.ok _, ?_⟩⟩
                intro s' heq_s'
                injection heq_s' with heq_state
                rw [← heq_state]
                rw [← heq_tys']
                exact ⟨StackMatches.cons (AvmVal.q0 _) h2, rfl⟩
  · -- subSatQ0
    rename_i ty tys_tail
    cases ty <;> try (unfold checkPrim at hprim; dsimp at hprim; contradiction)
    cases tys_tail
    · unfold checkPrim at hprim; dsimp at hprim; contradiction
    · rename_i ty2 tys_tail'
      cases ty2 <;> try (unfold checkPrim at hprim; dsimp at hprim; contradiction)
      cases hstack with
      | cons val1 h1 =>
          cases h1 with
          | cons val2 h2 =>
              cases val1; cases val2
              unfold checkPrim at hprim; dsimp at hprim
              injection hprim with heq_tys'
              unfold evalPrim pop1 push1
              dsimp
              split
              · refine ⟨_, rfl, ⟨SafeOutcome.stackOverflow, by intro s' h; contradiction⟩⟩
              · refine ⟨_, rfl, ⟨SafeOutcome.ok _, ?_⟩⟩
                intro s' heq_s'
                injection heq_s' with heq_state
                rw [← heq_state]
                rw [← heq_tys']
                exact ⟨StackMatches.cons (AvmVal.q0 _) h2, rfl⟩
  · -- addSatQ16
    rename_i ty tys_tail
    cases ty <;> try (unfold checkPrim at hprim; dsimp at hprim; contradiction)
    cases tys_tail
    · unfold checkPrim at hprim; dsimp at hprim; contradiction
    · rename_i ty2 tys_tail'
      cases ty2 <;> try (unfold checkPrim at hprim; dsimp at hprim; contradiction)
      cases hstack with
      | cons val1 h1 =>
          cases h1 with
          | cons val2 h2 =>
              cases val1; cases val2
              unfold checkPrim at hprim; dsimp at hprim
              injection hprim with heq_tys'
              unfold evalPrim pop1 push1
              dsimp
              split
              · refine ⟨_, rfl, ⟨SafeOutcome.stackOverflow, by intro s' h; contradiction⟩⟩
              · refine ⟨_, rfl, ⟨SafeOutcome.ok _, ?_⟩⟩
                intro s' heq_s'
                injection heq_s' with heq_state
                rw [← heq_state]
                rw [← heq_tys']
                exact ⟨StackMatches.cons (AvmVal.q16 _) h2, rfl⟩
  · -- subSatQ16
    rename_i ty tys_tail
    cases ty <;> try (unfold checkPrim at hprim; dsimp at hprim; contradiction)
    cases tys_tail
    · unfold checkPrim at hprim; dsimp at hprim; contradiction
    · rename_i ty2 tys_tail'
      cases ty2 <;> try (unfold checkPrim at hprim; dsimp at hprim; contradiction)
      cases hstack with
      | cons val1 h1 =>
          cases h1 with
          | cons val2 h2 =>
              cases val1; cases val2
              unfold checkPrim at hprim; dsimp at hprim
              injection hprim with heq_tys'
              unfold evalPrim pop1 push1
              dsimp
              split
              · refine ⟨_, rfl, ⟨SafeOutcome.stackOverflow, by intro s' h; contradiction⟩⟩
              · refine ⟨_, rfl, ⟨SafeOutcome.ok _, ?_⟩⟩
                intro s' heq_s'
                injection heq_s' with heq_state
                rw [← heq_state]
                rw [← heq_tys']
                exact ⟨StackMatches.cons (AvmVal.q16 _) h2, rfl⟩
  · -- mulSatQ16
    rename_i ty tys_tail
    cases ty <;> try (unfold checkPrim at hprim; dsimp at hprim; contradiction)
    cases tys_tail
    · unfold checkPrim at hprim; dsimp at hprim; contradiction
    · rename_i ty2 tys_tail'
      cases ty2 <;> try (unfold checkPrim at hprim; dsimp at hprim; contradiction)
      cases hstack with
      | cons val1 h1 =>
          cases h1 with
          | cons val2 h2 =>
              cases val1; cases val2
              unfold checkPrim at hprim; dsimp at hprim
              injection hprim with heq_tys'
              unfold evalPrim pop1 push1
              dsimp
              split
              · refine ⟨_, rfl, ⟨SafeOutcome.stackOverflow, by intro s' h; contradiction⟩⟩
              · refine ⟨_, rfl, ⟨SafeOutcome.ok _, ?_⟩⟩
                intro s' heq_s'
                injection heq_s' with heq_state
                rw [← heq_state]
                rw [← heq_tys']
                exact ⟨StackMatches.cons (AvmVal.q16 _) h2, rfl⟩
  · -- divSatQ16
    rename_i ty tys_tail
    cases ty <;> try (unfold checkPrim at hprim; dsimp at hprim; contradiction)
    cases tys_tail
    · unfold checkPrim at hprim; dsimp at hprim; contradiction
    · rename_i ty2 tys_tail'
      cases ty2 <;> try (unfold checkPrim at hprim; dsimp at hprim; contradiction)
      cases hstack with
      | cons val1 h1 =>
          cases h1 with
          | cons val2 h2 =>
              cases val1; cases val2
              unfold checkPrim at hprim; dsimp at hprim
              injection hprim with heq_tys'
              unfold evalPrim pop1 push1
              dsimp
              split
              · refine ⟨_, rfl, ⟨SafeOutcome.divByZero, by intro s' h; contradiction⟩⟩
              · split
                · refine ⟨_, rfl, ⟨SafeOutcome.stackOverflow, by intro s' h; contradiction⟩⟩
                · refine ⟨_, rfl, ⟨SafeOutcome.ok _, ?_⟩⟩
                  intro s' heq_s'
                  injection heq_s' with heq_state
                  rw [← heq_state]
                  rw [← heq_tys']
                  exact ⟨StackMatches.cons (AvmVal.q16 _) h2, rfl⟩
  · -- ltQ16
    rename_i ty tys_tail
    cases ty <;> try (unfold checkPrim at hprim; dsimp at hprim; contradiction)
    cases tys_tail
    · unfold checkPrim at hprim; dsimp at hprim; contradiction
    · rename_i ty2 tys_tail'
      cases ty2 <;> try (unfold checkPrim at hprim; dsimp at hprim; contradiction)
      cases hstack with
      | cons val1 h1 =>
          cases h1 with
          | cons val2 h2 =>
              cases val1; cases val2
              unfold checkPrim at hprim; dsimp at hprim
              injection hprim with heq_tys'
              unfold evalPrim pop1 push1
              dsimp
              split
              · refine ⟨_, rfl, ⟨SafeOutcome.stackOverflow, by intro s' h; contradiction⟩⟩
              · refine ⟨_, rfl, ⟨SafeOutcome.ok _, ?_⟩⟩
                intro s' heq_s'
                injection heq_s' with heq_state
                rw [← heq_state]
                rw [← heq_tys']
                exact ⟨StackMatches.cons (AvmVal.b _) h2, rfl⟩
  · -- eqQ16
    rename_i ty tys_tail
    cases ty <;> try (unfold checkPrim at hprim; dsimp at hprim; contradiction)
    cases tys_tail
    · unfold checkPrim at hprim; dsimp at hprim; contradiction
    · rename_i ty2 tys_tail'
      cases ty2 <;> try (unfold checkPrim at hprim; dsimp at hprim; contradiction)
      cases hstack with
      | cons val1 h1 =>
          cases h1 with
          | cons val2 h2 =>
              cases val1; cases val2
              unfold checkPrim at hprim; dsimp at hprim
              injection hprim with heq_tys'
              unfold evalPrim pop1 push1
              dsimp
              split
              · refine ⟨_, rfl, ⟨SafeOutcome.stackOverflow, by intro s' h; contradiction⟩⟩
              · refine ⟨_, rfl, ⟨SafeOutcome.ok _, ?_⟩⟩
                intro s' heq_s'
                injection heq_s' with heq_state
                rw [← heq_state]
                rw [← heq_tys']
                exact ⟨StackMatches.cons (AvmVal.b _) h2, rfl⟩
  · -- and
    rename_i ty tys_tail
    cases ty <;> try (unfold checkPrim at hprim; dsimp at hprim; contradiction)
    cases tys_tail
    · unfold checkPrim at hprim; dsimp at hprim; contradiction
    · rename_i ty2 tys_tail'
      cases ty2 <;> try (unfold checkPrim at hprim; dsimp at hprim; contradiction)
      cases hstack with
      | cons val1 h1 =>
          cases h1 with
          | cons val2 h2 =>
              cases val1; cases val2
              unfold checkPrim at hprim; dsimp at hprim
              injection hprim with heq_tys'
              unfold evalPrim pop1 push1
              dsimp
              split
              · refine ⟨_, rfl, ⟨SafeOutcome.stackOverflow, by intro s' h; contradiction⟩⟩
              · refine ⟨_, rfl, ⟨SafeOutcome.ok _, ?_⟩⟩
                intro s' heq_s'
                injection heq_s' with heq_state
                rw [← heq_state]
                rw [← heq_tys']
                exact ⟨StackMatches.cons (AvmVal.b _) h2, rfl⟩
  · -- or
    rename_i ty tys_tail
    cases ty <;> try (unfold checkPrim at hprim; dsimp at hprim; contradiction)
    cases tys_tail
    · unfold checkPrim at hprim; dsimp at hprim; contradiction
    · rename_i ty2 tys_tail'
      cases ty2 <;> try (unfold checkPrim at hprim; dsimp at hprim; contradiction)
      cases hstack with
      | cons val1 h1 =>
          cases h1 with
          | cons val2 h2 =>
              cases val1; cases val2
              unfold checkPrim at hprim; dsimp at hprim
              injection hprim with heq_tys'
              unfold evalPrim pop1 push1
              dsimp
              split
              · refine ⟨_, rfl, ⟨SafeOutcome.stackOverflow, by intro s' h; contradiction⟩⟩
              · refine ⟨_, rfl, ⟨SafeOutcome.ok _, ?_⟩⟩
                intro s' heq_s'
                injection heq_s' with heq_state
                rw [← heq_state]
                rw [← heq_tys']
                exact ⟨StackMatches.cons (AvmVal.b _) h2, rfl⟩
  · -- not
    rename_i ty tys_tail
    cases ty <;> try (unfold checkPrim at hprim; dsimp at hprim; contradiction)
    cases hstack with
    | cons val h' =>
        cases val
        unfold checkPrim at hprim; dsimp at hprim
        injection hprim with heq_tys'
        unfold evalPrim pop1 push1
        dsimp
        split
        · refine ⟨_, rfl, ⟨SafeOutcome.stackOverflow, by intro s' h; contradiction⟩⟩
        · refine ⟨_, rfl, ⟨SafeOutcome.ok _, ?_⟩⟩
          intro s' heq_s'
          injection heq_s' with heq_state
          rw [← heq_state]
          rw [← heq_tys']
          exact ⟨StackMatches.cons (AvmVal.b _) h', rfl⟩

/-- Type preservation theorem for single instruction step. -/
theorem step_preservation {instr : Instr} {s : State} {tys : List AvmTy} {tys_locals : List (Option AvmTy)}
    (hstack : StackMatches s.stack tys) (hlocals : LocalsMatches s.locals tys_locals)
    (hpc : s.pc = 0) (hhalt : s.halted = false)
    {tys' : List AvmTy} {tys_locals' : List (Option AvmTy)}
    (hcheck : checkInstr instr tys tys_locals = some (tys', tys_locals'))
    (hnon_jump : match instr with | .jump _ | .jumpIf _ | .halt => False | _ => True) :
    ∃ o, step [instr] s = o ∧ SafeOutcome o ∧ (∀ s', o = Outcome.ok s' → StackMatches s'.stack tys' ∧ LocalsMatches s'.locals tys_locals') := by
  rcases s with ⟨pc, stack, locals, halted⟩
  subst hpc hhalt
  cases instr
  case push val =>
      unfold checkInstr at hcheck
      injection hcheck with h
      injection h with hst hloc
      subst hst hloc
      unfold step pop1
      simp
      split
      · refine ⟨SafeOutcome.stackOverflow, by intro s' h; contradiction⟩
      · refine ⟨SafeOutcome.ok _, ?_⟩
        intro s' heq_s'
        injection heq_s' with heq_state
        rw [← heq_state]
        exact ⟨StackMatches.cons val.val hstack, hlocals⟩
  case pop =>
      unfold checkInstr at hcheck
      dsimp at hcheck
      split at hcheck
      case h_1 x xs heq =>
          injection hcheck with h
          injection h with hst hloc
          subst hst hloc
          cases hstack
          unfold step pop1
          simp
          refine ⟨SafeOutcome.ok _, ⟨by assumption, hlocals⟩⟩
      case h_2 heq => contradiction
  case dup =>
      unfold checkInstr at hcheck
      dsimp at hcheck
      split at hcheck
      case h_1 x xs heq =>
          injection hcheck with h
          injection h with hst hloc
          subst hst hloc
          cases hstack with
          | cons val h' =>
              unfold step pop1
              simp
              split
              · refine ⟨SafeOutcome.stackOverflow, by intro s' h; contradiction⟩
              · refine ⟨SafeOutcome.ok _, ?_⟩
                intro s' heq_s'
                injection heq_s' with heq_state
                rw [← heq_state]
                exact ⟨StackMatches.cons val (StackMatches.cons val h'), hlocals⟩
      case h_2 heq => contradiction
  case swap =>
      unfold checkInstr at hcheck
      dsimp at hcheck
      split at hcheck
      case h_1 x y xs heq =>
          injection hcheck with h
          injection h with hst hloc
          subst hst hloc
          cases hstack with
          | cons val1 h1 =>
              cases h1 with
              | cons val2 h2 =>
                  unfold step pop1
                  simp
                  refine ⟨SafeOutcome.ok _, ⟨StackMatches.cons val2 (StackMatches.cons val1 h2), hlocals⟩⟩
      case h_2 heq => contradiction
  case load i =>
      unfold checkInstr at hcheck
      dsimp at hcheck
      split at hcheck <;> try contradiction
      rename_i ty heq
      injection hcheck with h
      injection h with hst hloc
      subst hst hloc
      obtain ⟨val, hloc_val⟩ := getLocal_safety hlocals i heq
      simp [getLocal?] at hloc_val
      unfold step pop1 getLocal?
      simp
      rw [hloc_val]
      simp
      split
      · refine ⟨SafeOutcome.stackOverflow, by intro s' h; contradiction⟩
      · refine ⟨SafeOutcome.ok _, ?_⟩
        intro s' heq_s'
        injection heq_s' with heq_state
        rw [← heq_state]
        exact ⟨StackMatches.cons val hstack, hlocals⟩
  case store i =>
      unfold checkInstr at hcheck
      dsimp at hcheck
      split at hcheck
      · -- isTrue
        rename_i hbound
        split at hcheck
        case h_1 ty tys_tail heq =>
            injection hcheck with h
            injection h with hst hloc
            subst hst hloc
            cases hstack with
            | cons val h' =>
                unfold step pop1 setLocal?
                simp
                have hbound_locals : i < locals.length := by
                  rw [hlocals.1]
                  exact hbound
                rw [if_pos hbound_locals]
                simp
                refine ⟨SafeOutcome.ok _, ⟨h', setLocal_safety hlocals i val hbound⟩⟩
        case h_2 heq => contradiction
      · -- isFalse
        contradiction
  case jump target => contradiction
  case jumpIf target => contradiction
  case prim p =>
      unfold checkInstr at hcheck
      dsimp at hcheck
      split at hcheck
      case h_1 stack' heq =>
          injection hcheck with h
          injection h with hst hloc
          subst hst hloc
          obtain ⟨o, heval, hsafe, hst'⟩ := evalPrim_safety hstack heq
          unfold step
          have hlookup : [Instr.prim p][0]? = some (Instr.prim p) := rfl
          rw [hlookup]
          simp
          rw [heval]
          cases o
          case ok s1 =>
              refine ⟨SafeOutcome.ok _, ?_⟩
              intro s' heq_s'
              injection heq_s' with heq_state
              rw [← heq_state]
              obtain ⟨hst_matches, hlocals_eq⟩ := hst' s1 rfl
              dsimp at hlocals hlocals_eq ⊢
              rw [hlocals_eq]
              exact ⟨hst_matches, hlocals⟩
          case err e =>
              refine ⟨hsafe, by intro s' h; contradiction⟩
      case h_2 heq => contradiction
  case halt => contradiction

end SilverSight.AVMIsa
