/-
  WP34 — Abstract signature correctness in an axiomatised cyclic group (no Mathlib).

  `CyclicGroup G` bundles an additively written commutative group with scalar multiplication by
  naturals, a distinguished generator `gen` of order `order`, and an `xcoord` map invariant under
  negation (the model of a prime-order elliptic-curve group such as secp256k1 with the
  x-coordinate projection). EVERY FIELD OF THE STRUCTURE IS AN AXIOM OF THE MODEL; everything
  stated with `theorem` below is derived from those fields and Lean-core arithmetic on `Nat`.
-/

namespace Wp34

structure CyclicGroup (G : Type) where
  zero : G
  add : G → G → G
  neg : G → G
  nsmul : Nat → G → G
  xcoord : G → Nat
  gen : G
  order : Nat
  -- axioms: commutative group
  add_assoc : ∀ a b c, add (add a b) c = add a (add b c)
  add_comm : ∀ a b, add a b = add b a
  zero_add : ∀ a, add zero a = a
  add_neg : ∀ a, add a (neg a) = zero
  -- axioms: scalar multiplication by naturals (recursive definition, as fields)
  nsmul_zero : ∀ a, nsmul 0 a = zero
  nsmul_succ : ∀ k a, nsmul (k + 1) a = add (nsmul k a) a
  -- axioms: the generator has order `order` (> 1)
  one_lt_order : 1 < order
  nsmul_order : nsmul order gen = zero
  -- axiom: the x-coordinate does not see the sign of a point
  xcoord_neg : ∀ P, xcoord (neg P) = xcoord P

namespace CyclicGroup

variable {G : Type} (C : CyclicGroup G)

/-! ### Derived group facts (theorems, not axioms) -/

theorem add_zero (a : G) : C.add a C.zero = a := by
  rw [C.add_comm, C.zero_add]

theorem neg_add (a : G) : C.add (C.neg a) a = C.zero := by
  rw [C.add_comm, C.add_neg]

/-- Uniqueness of inverses: `a + b = 0 → a = −b`. -/
theorem eq_neg_of_add_eq_zero {a b : G} (h : C.add a b = C.zero) : a = C.neg b := by
  calc a = C.add a C.zero := (C.add_zero a).symm
    _ = C.add a (C.add b (C.neg b)) := by rw [C.add_neg]
    _ = C.add (C.add a b) (C.neg b) := by rw [C.add_assoc]
    _ = C.add C.zero (C.neg b) := by rw [h]
    _ = C.neg b := C.zero_add _

theorem neg_add_rev (a b : G) : C.neg (C.add a b) = C.add (C.neg a) (C.neg b) := by
  symm
  apply C.eq_neg_of_add_eq_zero
  calc C.add (C.add (C.neg a) (C.neg b)) (C.add a b)
      = C.add (C.neg a) (C.add (C.neg b) (C.add b a)) := by
        rw [C.add_assoc, C.add_comm a b]
    _ = C.add (C.neg a) (C.add (C.add (C.neg b) b) a) := by rw [C.add_assoc]
    _ = C.add (C.neg a) a := by rw [C.neg_add, C.zero_add]
    _ = C.zero := C.neg_add a

theorem nsmul_add (m k : Nat) (a : G) : C.nsmul (m + k) a = C.add (C.nsmul m a) (C.nsmul k a) := by
  induction k with
  | zero => rw [Nat.add_zero, C.nsmul_zero, C.add_zero]
  | succ k ih => rw [← Nat.add_assoc, C.nsmul_succ, ih, C.nsmul_succ, C.add_assoc]

theorem nsmul_zero_point (k : Nat) : C.nsmul k C.zero = C.zero := by
  induction k with
  | zero => exact C.nsmul_zero _
  | succ k ih => rw [C.nsmul_succ, ih, C.zero_add]

theorem nsmul_mul (m k : Nat) (a : G) : C.nsmul (m * k) a = C.nsmul m (C.nsmul k a) := by
  induction m with
  | zero => rw [Nat.zero_mul, C.nsmul_zero, C.nsmul_zero]
  | succ m ih => rw [Nat.succ_mul, C.nsmul_add, ih, C.nsmul_succ]

theorem add_add_add_comm (a b c d : G) :
    C.add (C.add a b) (C.add c d) = C.add (C.add a c) (C.add b d) := by
  calc C.add (C.add a b) (C.add c d) = C.add a (C.add b (C.add c d)) := C.add_assoc _ _ _
    _ = C.add a (C.add (C.add b c) d) := by rw [C.add_assoc b c d]
    _ = C.add a (C.add (C.add c b) d) := by rw [C.add_comm b c]
    _ = C.add a (C.add c (C.add b d)) := by rw [C.add_assoc c b d]
    _ = C.add (C.add a c) (C.add b d) := (C.add_assoc _ _ _).symm

theorem nsmul_add_point (k : Nat) (a b : G) :
    C.nsmul k (C.add a b) = C.add (C.nsmul k a) (C.nsmul k b) := by
  induction k with
  | zero => rw [C.nsmul_zero, C.nsmul_zero, C.nsmul_zero, C.zero_add]
  | succ k ih => rw [C.nsmul_succ, C.nsmul_succ, C.nsmul_succ, ih, C.add_add_add_comm]

theorem nsmul_neg (k : Nat) (a : G) : C.nsmul k (C.neg a) = C.neg (C.nsmul k a) := by
  apply C.eq_neg_of_add_eq_zero
  rw [← C.nsmul_add_point, C.neg_add, C.nsmul_zero_point]

/-- Membership in the cyclic subgroup generated by `gen` (`P = d•G` for some `d`). -/
def InSubgroup (P : G) : Prop := ∃ d, P = C.nsmul d C.gen

theorem nsmul_one (a : G) : C.nsmul 1 a = a := by
  rw [show (1 : Nat) = 0 + 1 from rfl, C.nsmul_succ, C.nsmul_zero, C.zero_add]

theorem gen_inSubgroup : C.InSubgroup C.gen := ⟨1, (C.nsmul_one _).symm⟩

theorem nsmul_inSubgroup (k : Nat) (P : G) (hP : C.InSubgroup P) : C.InSubgroup (C.nsmul k P) := by
  obtain ⟨d, rfl⟩ := hP
  exact ⟨k * d, (C.nsmul_mul k d C.gen).symm⟩

/-- Every point of the subgroup is killed by `order`. -/
theorem nsmul_order_of_mem (P : G) (hP : C.InSubgroup P) : C.nsmul C.order P = C.zero := by
  obtain ⟨d, rfl⟩ := hP
  rw [← C.nsmul_mul, Nat.mul_comm, C.nsmul_mul, C.nsmul_order, C.nsmul_zero_point]

/-- Scalars act modulo `order` on the subgroup. -/
theorem nsmul_mod (k : Nat) (P : G) (hP : C.InSubgroup P) :
    C.nsmul (k % C.order) P = C.nsmul k P := by
  conv => rhs; rw [← Nat.div_add_mod k C.order]
  rw [C.nsmul_add, Nat.mul_comm C.order (k / C.order), C.nsmul_mul, C.nsmul_order_of_mem P hP,
    C.nsmul_zero_point, C.zero_add]

theorem nsmul_congr_mod (k k' : Nat) (P : G) (hP : C.InSubgroup P) (h : k % C.order = k' % C.order) :
    C.nsmul k P = C.nsmul k' P := by
  rw [← C.nsmul_mod k P hP, h, C.nsmul_mod k' P hP]

/-- `(order − u)•P = −(u•P)` for `u ≤ order` and `P` in the subgroup. -/
theorem nsmul_order_sub (u : Nat) (hu : u ≤ C.order) (P : G) (hP : C.InSubgroup P) :
    C.nsmul (C.order - u) P = C.neg (C.nsmul u P) := by
  apply C.eq_neg_of_add_eq_zero
  rw [← C.nsmul_add, Nat.sub_add_cancel hu, C.nsmul_order_of_mem P hP]

/-! ### Schnorr signatures (BIP340-style, abstract)

Key pair `(d, P = d•G)`, nonce `k` with commitment `R = k•G`, challenge `e` (any natural — in
BIP340 it is `H(R ‖ P ‖ m)`), response `s = (k + e·d) mod n`.
Verification: `s•G = R + e•P`. -/

/-- Schnorr verification equation. -/
def schnorrVerify (P R : G) (e s : Nat) : Prop :=
  C.nsmul s C.gen = C.add R (C.nsmul e P)

/-- **Schnorr correctness**: an honestly generated signature verifies. -/
theorem schnorr_correct (d k e : Nat) :
    C.schnorrVerify (C.nsmul d C.gen) (C.nsmul k C.gen) e ((k + e * d) % C.order) := by
  unfold schnorrVerify
  rw [C.nsmul_mod _ _ C.gen_inSubgroup, C.nsmul_add, C.nsmul_mul]

/-- Variant without the final reduction (`s = k + e·d` as a natural number). -/
theorem schnorr_correct_nat (d k e : Nat) :
    C.schnorrVerify (C.nsmul d C.gen) (C.nsmul k C.gen) e (k + e * d) := by
  unfold schnorrVerify
  rw [C.nsmul_add, C.nsmul_mul]

/-! ### ECDSA (abstract)

Signature `(r, s)` on hash `e` under public key `P`. Verification computes `w = s⁻¹ mod n`,
`u₁ = e·w`, `u₂ = r·w`, `R' = u₁•G + u₂•P` and checks `x(R') mod n = r`. The modular inverse is
given as a witness `w` with `s·w ≡ 1 (mod n)`; scalars `e·w`, `r·w` need not be reduced because
`nsmul_mod` shows scalars act modulo `n` on the subgroup. -/

/-- ECDSA verification predicate (`0 < r, s < n` as in the standard). -/
def ecdsaVerify (P : G) (e r s : Nat) : Prop :=
  0 < r ∧ r < C.order ∧ 0 < s ∧ s < C.order ∧
    ∃ w, (s * w) % C.order = 1 ∧
      (C.xcoord (C.add (C.nsmul (e * w) C.gen) (C.nsmul (r * w) P))) % C.order = r

/-- Helper: `(n − s)(n − w) ≡ s·w (mod n)` for `s, w ≤ n` (pure `Nat` arithmetic). -/
theorem sub_mul_sub_mod (n s w : Nat) (hs : s ≤ n) (hw : w ≤ n) :
    ((n - s) * (n - w)) % n = (s * w) % n := by
  obtain ⟨p, hp⟩ := Nat.exists_eq_add_of_le hs
  obtain ⟨q, hq⟩ := Nat.exists_eq_add_of_le hw
  have h1 : n - s = p := by omega
  have h2 : n - w = q := by omega
  rw [h1, h2]
  -- p*q + n*(w + s) = s*w + n*n, hence p*q ≡ s*w (mod n)
  have key : p * q + n * (w + s) = s * w + n * n := by
    have e1 : n * n = n * w + s * q + p * q := by
      calc n * n = n * (w + q) := by rw [← hq]
        _ = n * w + n * q := Nat.mul_add _ _ _
        _ = n * w + (s + p) * q := by rw [← hp]
        _ = n * w + (s * q + p * q) := by rw [Nat.add_mul]
        _ = n * w + s * q + p * q := by rw [Nat.add_assoc]
    have e2 : n * (w + s) = n * w + s * q + s * w := by
      calc n * (w + s) = n * w + n * s := Nat.mul_add _ _ _
        _ = n * w + s * n := by rw [Nat.mul_comm n s]
        _ = n * w + s * (w + q) := by rw [← hq]
        _ = n * w + (s * w + s * q) := by rw [Nat.mul_add]
        _ = n * w + s * q + s * w := by omega
    rw [e1, e2]; omega
  have : (p * q + n * (w + s)) % n = (s * w + n * n) % n := by rw [key]
  rwa [Nat.add_mul_mod_self_left, Nat.add_mul_mod_self_left] at this

/-- **ECDSA malleability**: if `(r, s)` verifies then so does `(r, n − s)`, using only the
group axioms and `xcoord (−Q) = xcoord Q`. -/
theorem ecdsa_malleable (P : G) (hP : C.InSubgroup P) (e r s : Nat)
    (h : C.ecdsaVerify P e r s) : C.ecdsaVerify P e r (C.order - s) := by
  obtain ⟨hr0, hrn, hs0, hsn, w, hw, hx⟩ := h
  have hn : 0 < C.order := Nat.lt_trans Nat.zero_lt_one C.one_lt_order
  have hG : C.InSubgroup C.gen := C.gen_inSubgroup
  -- reduce the inverse witness modulo n: w0 = w mod n, 0 < w0 < n, s·w0 ≡ 1
  obtain ⟨w0, hw0def⟩ : ∃ w0, w % C.order = w0 := ⟨_, rfl⟩
  have hw0n : w0 < C.order := hw0def ▸ Nat.mod_lt _ hn
  have hw0 : (s * w0) % C.order = 1 := by
    rw [Nat.mul_mod, ← hw0def, Nat.mod_mod, ← Nat.mul_mod, hw]
  have hw0pos : 0 < w0 := by
    rcases Nat.eq_zero_or_pos w0 with h0 | h0
    · rw [h0, Nat.mul_zero, Nat.zero_mod] at hw0; omega
    · exact h0
  -- (n − w0) + w is a multiple of n
  have hsum : C.order - w0 + w = C.order * (w / C.order + 1) := by
    have := Nat.div_add_mod w C.order
    rw [hw0def] at this
    rw [Nat.mul_add, Nat.mul_one]; omega
  -- u•(n − w0) + u•w = 0 for every u and every Q in the subgroup
  have hneg : ∀ (u : Nat) (Q : G), C.InSubgroup Q →
      C.nsmul (u * (C.order - w0)) Q = C.neg (C.nsmul (u * w) Q) := by
    intro u Q hQ
    apply C.eq_neg_of_add_eq_zero
    have h1 : u * (C.order - w0) + u * w = C.order * ((w / C.order + 1) * u) := by
      rw [← Nat.mul_add, hsum]; ac_rfl
    rw [← C.nsmul_add, h1, C.nsmul_mul, C.nsmul_order_of_mem _ (C.nsmul_inSubgroup _ _ hQ)]
  refine ⟨hr0, hrn, by omega, by omega, C.order - w0, ?_, ?_⟩
  · -- (n − s)(n − w0) ≡ s·w0 ≡ 1
    rw [sub_mul_sub_mod _ _ _ (Nat.le_of_lt hsn) (Nat.le_of_lt hw0n), hw0]
  · -- the verification point is negated, so its x-coordinate is unchanged
    rw [hneg e C.gen hG, hneg r P hP, ← C.neg_add_rev, C.xcoord_neg]
    exact hx

/-- **ECDSA correctness**: an honestly generated signature verifies. Hypotheses encode
`r = x(k•G) mod n` and `s ≡ k⁻¹(e + r·d)`, i.e. `k·s ≡ e + r·d (mod n)`, with a witness `w` for
`s⁻¹`. -/
theorem ecdsa_correct (d k e r s w : Nat)
    (hr0 : 0 < r) (hrn : r < C.order) (hs0 : 0 < s) (hsn : s < C.order)
    (hr : (C.xcoord (C.nsmul k C.gen)) % C.order = r)
    (hs : (k * s) % C.order = (e + r * d) % C.order)
    (hw : (s * w) % C.order = 1) :
    C.ecdsaVerify (C.nsmul d C.gen) e r s := by
  refine ⟨hr0, hrn, hs0, hsn, w, hw, ?_⟩
  have hG : C.InSubgroup C.gen := C.gen_inSubgroup
  have hpt : C.add (C.nsmul (e * w) C.gen) (C.nsmul (r * w) (C.nsmul d C.gen)) =
      C.nsmul k C.gen := by
    rw [← C.nsmul_mul, ← C.nsmul_add]
    apply C.nsmul_congr_mod _ _ _ hG
    -- e·w + r·w·d = w·(e + r·d) ≡ w·k·s = k·(s·w) ≡ k
    have e1 : e * w + r * w * d = w * (e + r * d) := by rw [Nat.mul_add]; ac_rfl
    have e2 : w * (k * s) = k * (s * w) := by ac_rfl
    rw [e1, Nat.mul_mod, ← hs, ← Nat.mul_mod, e2, Nat.mul_mod, hw, Nat.mul_one, Nat.mod_mod]
  rw [hpt, hr]

end CyclicGroup

end Wp34