/-
  WP34 — the SHA-256 step function is a bijection of the 256-bit state for fixed (K_t, W_t).
  `step` and `stepInv` are transcribed from `step` / `step_inverse` in `src/wp34/sha256.py`.
-/
import Wp34.Sha256Words

namespace Wp34

/-- The eight 32-bit working registers `(a, b, c, d, e, f, g, h)`. -/
structure State where
  a : Word
  b : Word
  c : Word
  d : Word
  e : Word
  f : Word
  g : Word
  h : Word
  deriving DecidableEq, Repr

/-- One SHA-256 step (FIPS 180-4 §6.2.2 step 3) with round constant `k` and message word `w`. -/
def step (s : State) (k w : Word) : State :=
  let T1 := s.h + Sigma1 s.e + Ch s.e s.f s.g + k + w
  let T2 := Sigma0 s.a + Maj s.a s.b s.c
  { a := T1 + T2, b := s.a, c := s.b, d := s.c, e := s.d + T1, f := s.e, g := s.f, h := s.g }

/-- Inverse step (`step_inverse` in the Python reference): the new state is
`(a₁, a, b, c, e₁, e, f, g)`; recover `T2` from `(a, b, c)`, then `T1 = a₁ − T2`, `d = e₁ − T1`,
`h = T1 − Σ1(e) − Ch(e,f,g) − k − w`. -/
def stepInv (n : State) (k w : Word) : State :=
  let T2 := Sigma0 n.b + Maj n.b n.c n.d
  let T1 := n.a - T2
  { a := n.b, b := n.c, c := n.d, d := n.e - T1,
    e := n.f, f := n.g, g := n.h, h := T1 - Sigma1 n.f - Ch n.f n.g n.h - k - w }

/-! #### Proof strategy

The direct `bv_decide` call on the 320-bit goal exceeds the SAT timeout (see README): the
instance mixes the non-linear Σ/Ch/Maj sub-terms with the re-association of five chained 32-bit
additions. The statement is *linear* in the sub-terms `T2 = Σ0(a) + Maj(a,b,c)`, `Σ1(e)` and
`Ch(e,f,g)`, so we abstract them with `generalize`; the two surviving components are then closed
by two small SAT lemmas on the pure addition chains (`add4_sub4`, `sub4_add4`, 160-bit
instances, ~1.5 s each) and core's `BitVec.add_sub_cancel` / `sub_add_cancel`. The other six
components are definitional. -/

/-- `h + A + B + k + w − A − B − k − w = h` (SAT). -/
theorem add4_sub4 (h A B k w : Word) : h + A + B + k + w - A - B - k - w = h := by
  bv_decide (config := { timeout := 120 })

/-- `x − A − B − k − w + A + B + k + w = x` (SAT). -/
theorem sub4_add4 (x A B k w : Word) : x - A - B - k - w + A + B + k + w = x := by
  bv_decide (config := { timeout := 120 })

/-- `stepInv` is a left inverse of `step`: every SHA-256 step can be undone given `(k, w)`. -/
theorem stepInv_step (s : State) (k w : Word) : stepInv (step s k w) k w = s := by
  obtain ⟨a, b, c, d, e, f, g, h⟩ := s
  simp only [step, stepInv, State.mk.injEq]
  generalize Sigma0 a + Maj a b c = T2
  generalize Sigma1 e = A
  generalize Ch e f g = B
  simp only [BitVec.add_sub_cancel, add4_sub4, and_self]

/-- `stepInv` is a right inverse of `step`. -/
theorem step_stepInv (n : State) (k w : Word) : step (stepInv n k w) k w = n := by
  obtain ⟨a₁, a, b, c, e₁, e, f, g⟩ := n
  simp only [step, stepInv, State.mk.injEq]
  generalize Sigma0 a + Maj a b c = T2
  generalize Sigma1 e = A
  generalize Ch e f g = B
  simp only [sub4_add4, BitVec.sub_add_cancel, and_self]

/-- For fixed `(k, w)` the step function is injective on the state. -/
theorem step_injective (k w : Word) : Function.Injective (fun s => step s k w) := by
  intro s s' h
  have := congrArg (fun t => stepInv t k w) h
  simpa only [stepInv_step] using this

/-- For fixed `(k, w)` the step function is surjective on the state. -/
theorem step_surjective (k w : Word) : ∀ n, ∃ s, step s k w = n :=
  fun n => ⟨stepInv n k w, step_stepInv n k w⟩

/-- The step function is a bijection of the 256-bit state for every fixed `(k, w)`. -/
theorem step_bijective (k w : Word) : Bijective (fun s => step s k w) :=
  ⟨step_injective k w, step_surjective k w⟩

/-! ### Iterated steps: any number of rounds with known `(K_t, W_t)` is invertible. -/

/-- Apply the steps for the list of `(k, w)` pairs, first element first. -/
def steps (s : State) : List (Word × Word) → State
  | [] => s
  | (k, w) :: rest => steps (step s k w) rest

/-- Undo the steps, last element first. -/
def stepsInv (n : State) : List (Word × Word) → State
  | [] => n
  | (k, w) :: rest => stepInv (stepsInv n rest) k w

theorem stepsInv_steps (kws : List (Word × Word)) : ∀ s, stepsInv (steps s kws) kws = s := by
  induction kws with
  | nil => intro s; rfl
  | cons kw rest ih =>
    intro s
    obtain ⟨k, w⟩ := kw
    simp only [steps, stepsInv, ih, stepInv_step]

theorem steps_stepsInv (kws : List (Word × Word)) : ∀ n, steps (stepsInv n kws) kws = n := by
  induction kws with
  | nil => intro n; rfl
  | cons kw rest ih =>
    intro n
    obtain ⟨k, w⟩ := kw
    simp only [steps, stepsInv, step_stepInv, ih]

/-- The `r`-round compression core (before the Davies–Meyer feed-forward) is a bijection of the
state for every fixed sequence of round constants and message words. -/
theorem steps_bijective (kws : List (Word × Word)) : Bijective (fun s => steps s kws) :=
  ⟨fun s s' h => by
      have := congrArg (fun t => stepsInv t kws) h
      simpa only [stepsInv_steps] using this,
   fun n => ⟨stepsInv n kws, steps_stepsInv kws n⟩⟩

end Wp34