/-
SilverSight.WireFormat — shim + Q16_16/Q0_16 WireFormat instances
-/
import SilverSight.Semantics.WireFormat
import SilverSight.Semantics.Schema

namespace SilverSight.Semantics
open SilverSight.FixedPoint

-- ═══════════════════════════════════════════════════════════════════════════
-- Nat-level 4-byte little-endian byte-reassembly identity (testBit proof)
-- ═══════════════════════════════════════════════════════════════════════════

private theorem nat_byte_reassembly (n : Nat) (hn : n < 4294967296) :
    (n &&& 0xFF) ||| (((n >>> 8) &&& 0xFF) <<< 8) |||
    (((n >>> 16) &&& 0xFF) <<< 16) ||| (((n >>> 24) &&& 0xFF) <<< 24) = n := by
  apply Nat.eq_of_testBit_eq; intro i
  have h255 : ∀ k, Nat.testBit 255 k = decide (k < 8) := fun k => by
    simp only [show (255 : Nat) = 2 ^ 8 - 1 from by decide, Nat.testBit_two_pow_sub_one]
  simp only [Nat.testBit_or, Nat.testBit_and, Nat.testBit_shiftLeft,
             Nat.testBit_shiftRight, h255]
  rcases Nat.lt_or_ge i 8 with hi | hi
  · simp [decide_eq_true_eq.mpr hi,
          decide_eq_false_iff_not.mpr (by omega : ¬ i ≥ 8),
          decide_eq_false_iff_not.mpr (by omega : ¬ i ≥ 16),
          decide_eq_false_iff_not.mpr (by omega : ¬ i ≥ 24)]
  · rcases Nat.lt_or_ge i 16 with hi16 | hi16
    · simp [decide_eq_false_iff_not.mpr (by omega : ¬ i < 8),
            decide_eq_true_eq.mpr (show i ≥ 8 from hi),
            show 8 + (i - 8) = i from by omega,
            decide_eq_true_eq.mpr (show i - 8 < 8 from by omega),
            decide_eq_false_iff_not.mpr (by omega : ¬ i ≥ 16),
            decide_eq_false_iff_not.mpr (by omega : ¬ i ≥ 24)]
    · rcases Nat.lt_or_ge i 24 with hi24 | hi24
      · simp [decide_eq_false_iff_not.mpr (by omega : ¬ i < 8),
              decide_eq_true_eq.mpr (show i ≥ 8 from by omega),
              decide_eq_false_iff_not.mpr (by omega : ¬ (i - 8 < 8)),
              decide_eq_true_eq.mpr (show i ≥ 16 from hi16),
              show 16 + (i - 16) = i from by omega,
              decide_eq_true_eq.mpr (show i - 16 < 8 from by omega),
              decide_eq_false_iff_not.mpr (by omega : ¬ i ≥ 24)]
      · rcases Nat.lt_or_ge i 32 with hi32 | hi32
        · simp [decide_eq_false_iff_not.mpr (by omega : ¬ i < 8),
                decide_eq_true_eq.mpr (show i ≥ 8 from by omega),
                decide_eq_false_iff_not.mpr (by omega : ¬ (i - 8 < 8)),
                decide_eq_true_eq.mpr (show i ≥ 16 from by omega),
                decide_eq_false_iff_not.mpr (by omega : ¬ (i - 16 < 8)),
                decide_eq_true_eq.mpr (show i ≥ 24 from hi24),
                show 24 + (i - 24) = i from by omega,
                decide_eq_true_eq.mpr (show i - 24 < 8 from by omega)]
        · have hti : n.testBit i = false :=
            Nat.testBit_lt_two_pow (calc n < 4294967296 := hn
              _ = 2 ^ 32 := by decide
              _ ≤ 2 ^ i  := Nat.pow_le_pow_right (by decide) hi32)
          simp [hti,
                decide_eq_false_iff_not.mpr (by omega : ¬ i < 8),
                decide_eq_false_iff_not.mpr (by omega : ¬ (i - 8 < 8)),
                decide_eq_false_iff_not.mpr (by omega : ¬ (i - 16 < 8)),
                decide_eq_false_iff_not.mpr (by omega : ¬ (i - 24 < 8))]

-- ═══════════════════════════════════════════════════════════════════════════
-- Nat-level 2-byte little-endian byte-reassembly identity
-- ═══════════════════════════════════════════════════════════════════════════

private theorem nat_byte_reassembly16 (n : Nat) (hn : n < 65536) :
    (n &&& 0xFF) ||| (((n >>> 8) &&& 0xFF) <<< 8) = n := by
  apply Nat.eq_of_testBit_eq; intro i
  have h255 : ∀ k, Nat.testBit 255 k = decide (k < 8) := fun k => by
    simp only [show (255 : Nat) = 2 ^ 8 - 1 from by decide, Nat.testBit_two_pow_sub_one]
  simp only [Nat.testBit_or, Nat.testBit_and, Nat.testBit_shiftLeft,
             Nat.testBit_shiftRight, h255]
  rcases Nat.lt_or_ge i 8 with hi | hi
  · simp [decide_eq_true_eq.mpr hi,
          decide_eq_false_iff_not.mpr (by omega : ¬ i ≥ 8)]
  · rcases Nat.lt_or_ge i 16 with hi16 | hi16
    · simp [decide_eq_false_iff_not.mpr (by omega : ¬ i < 8),
            decide_eq_true_eq.mpr (show i ≥ 8 from hi),
            show 8 + (i - 8) = i from by omega,
            decide_eq_true_eq.mpr (show i - 8 < 8 from by omega)]
    · have hti : n.testBit i = false :=
        Nat.testBit_lt_two_pow (calc n < 65536 := hn
          _ = 2 ^ 16 := by decide
          _ ≤ 2 ^ i  := Nat.pow_le_pow_right (by decide) hi16)
      simp [hti,
            decide_eq_false_iff_not.mpr (by omega : ¬ i < 8),
            decide_eq_false_iff_not.mpr (by omega : ¬ (i - 8 < 8))]

-- ═══════════════════════════════════════════════════════════════════════════
-- UInt32 byte helpers (toNat level)
-- ═══════════════════════════════════════════════════════════════════════════

private theorem u32_byte0_nat (u : UInt32) :
    ((u &&& 0xFF).toUInt8).toUInt32.toNat = u.toNat &&& 255 := by
  simp only [UInt32.toNat_and, UInt32.toNat_toUInt8, UInt8.toNat_toUInt32,
             show (0xFF : UInt32).toNat = 255 from rfl]
  exact Nat.mod_eq_of_lt (Nat.lt_of_le_of_lt Nat.and_le_right (by decide))

private theorem u32_byte8_nat (u : UInt32) :
    (((u >>> 8) &&& 0xFF).toUInt8).toUInt32.toNat = u.toNat >>> 8 &&& 255 := by
  simp only [UInt32.toNat_and, UInt32.toNat_shiftRight, UInt32.toNat_toUInt8,
             UInt8.toNat_toUInt32, show (0xFF : UInt32).toNat = 255 from rfl,
             show (8 : UInt32).toNat % 32 = 8 from rfl]
  exact Nat.mod_eq_of_lt (Nat.lt_of_le_of_lt Nat.and_le_right (by decide))

private theorem u32_byte16_nat (u : UInt32) :
    (((u >>> 16) &&& 0xFF).toUInt8).toUInt32.toNat = u.toNat >>> 16 &&& 255 := by
  simp only [UInt32.toNat_and, UInt32.toNat_shiftRight, UInt32.toNat_toUInt8,
             UInt8.toNat_toUInt32, show (0xFF : UInt32).toNat = 255 from rfl,
             show (16 : UInt32).toNat % 32 = 16 from rfl]
  exact Nat.mod_eq_of_lt (Nat.lt_of_le_of_lt Nat.and_le_right (by decide))

private theorem u32_byte24_nat (u : UInt32) :
    (((u >>> 24) &&& 0xFF).toUInt8).toUInt32.toNat = u.toNat >>> 24 &&& 255 := by
  simp only [UInt32.toNat_and, UInt32.toNat_shiftRight, UInt32.toNat_toUInt8,
             UInt8.toNat_toUInt32, show (0xFF : UInt32).toNat = 255 from rfl,
             show (24 : UInt32).toNat % 32 = 24 from rfl]
  exact Nat.mod_eq_of_lt (Nat.lt_of_le_of_lt Nat.and_le_right (by decide))

-- ═══════════════════════════════════════════════════════════════════════════
-- UInt32 shifted-byte helpers (include the <<< k in toNat)
-- ═══════════════════════════════════════════════════════════════════════════

private theorem u32_byte8_shifted (u : UInt32) :
    ((((u >>> 8) &&& 0xFF).toUInt8).toUInt32 <<< 8).toNat =
    (u.toNat >>> 8 &&& 255) <<< 8 := by
  simp only [UInt32.toNat_shiftLeft, show (8 : UInt32).toNat % 32 = 8 from rfl,
             u32_byte8_nat]
  exact Nat.mod_eq_of_lt (by
    have := @Nat.and_le_right (u.toNat >>> 8) 255
    rw [Nat.shiftLeft_eq]
    calc (u.toNat >>> 8 &&& 255) * 2 ^ 8 ≤ 255 * 2 ^ 8 := by nlinarith
      _ < 2 ^ 32 := by decide)

private theorem u32_byte16_shifted (u : UInt32) :
    ((((u >>> 16) &&& 0xFF).toUInt8).toUInt32 <<< 16).toNat =
    (u.toNat >>> 16 &&& 255) <<< 16 := by
  simp only [UInt32.toNat_shiftLeft, show (16 : UInt32).toNat % 32 = 16 from rfl,
             u32_byte16_nat]
  exact Nat.mod_eq_of_lt (by
    have := @Nat.and_le_right (u.toNat >>> 16) 255
    rw [Nat.shiftLeft_eq]
    calc (u.toNat >>> 16 &&& 255) * 2 ^ 16 ≤ 255 * 2 ^ 16 := by nlinarith
      _ < 2 ^ 32 := by decide)

private theorem u32_byte24_shifted (u : UInt32) :
    ((((u >>> 24) &&& 0xFF).toUInt8).toUInt32 <<< 24).toNat =
    (u.toNat >>> 24 &&& 255) <<< 24 := by
  simp only [UInt32.toNat_shiftLeft, show (24 : UInt32).toNat % 32 = 24 from rfl,
             u32_byte24_nat]
  exact Nat.mod_eq_of_lt (by
    have := @Nat.and_le_right (u.toNat >>> 24) 255
    rw [Nat.shiftLeft_eq]
    calc (u.toNat >>> 24 &&& 255) * 2 ^ 24 ≤ 255 * 2 ^ 24 := by nlinarith
      _ < 2 ^ 32 := by decide)

-- ═══════════════════════════════════════════════════════════════════════════
-- UInt32 four-byte round-trip
-- ═══════════════════════════════════════════════════════════════════════════

private theorem u32_byte_reassembly (u : UInt32) :
    ((u &&& 0xFF).toUInt8).toUInt32 |||
    ((((u >>> 8) &&& 0xFF).toUInt8).toUInt32 <<< 8) |||
    ((((u >>> 16) &&& 0xFF).toUInt8).toUInt32 <<< 16) |||
    ((((u >>> 24) &&& 0xFF).toUInt8).toUInt32 <<< 24) = u := by
  apply UInt32.toNat.inj
  simp only [UInt32.toNat_or]
  rw [u32_byte0_nat u, u32_byte8_shifted u, u32_byte16_shifted u, u32_byte24_shifted u]
  exact nat_byte_reassembly u.toNat u.toNat_lt

-- ═══════════════════════════════════════════════════════════════════════════
-- UInt16 byte helpers
-- ═══════════════════════════════════════════════════════════════════════════

private theorem u16_byte0_nat (u : UInt16) :
    ((u &&& 0xFF).toUInt8).toUInt16.toNat = u.toNat &&& 255 := by
  simp only [UInt16.toNat_and, UInt16.toNat_toUInt8, UInt8.toNat_toUInt16,
             show (0xFF : UInt16).toNat = 255 from rfl]
  exact Nat.mod_eq_of_lt (Nat.lt_of_le_of_lt Nat.and_le_right (by decide))

private theorem u16_byte8_nat (u : UInt16) :
    (((u >>> 8) &&& 0xFF).toUInt8).toUInt16.toNat = u.toNat >>> 8 &&& 255 := by
  simp only [UInt16.toNat_and, UInt16.toNat_shiftRight, UInt16.toNat_toUInt8,
             UInt8.toNat_toUInt16, show (0xFF : UInt16).toNat = 255 from rfl,
             show (8 : UInt16).toNat % 16 = 8 from rfl]
  exact Nat.mod_eq_of_lt (Nat.lt_of_le_of_lt Nat.and_le_right (by decide))

private theorem u16_byte8_shifted (u : UInt16) :
    ((((u >>> 8) &&& 0xFF).toUInt8).toUInt16 <<< 8).toNat =
    (u.toNat >>> 8 &&& 255) <<< 8 := by
  simp only [UInt16.toNat_shiftLeft, show (8 : UInt16).toNat % 16 = 8 from rfl,
             u16_byte8_nat]
  exact Nat.mod_eq_of_lt (by
    have := @Nat.and_le_right (u.toNat >>> 8) 255
    rw [Nat.shiftLeft_eq]
    calc (u.toNat >>> 8 &&& 255) * 2 ^ 8 ≤ 255 * 2 ^ 8 := by nlinarith
      _ < 2 ^ 16 := by decide)

-- ═══════════════════════════════════════════════════════════════════════════
-- UInt16 two-byte round-trip
-- ═══════════════════════════════════════════════════════════════════════════

private theorem u16_byte_reassembly (u : UInt16) :
    ((u &&& 0xFF).toUInt8).toUInt16 |||
    ((((u >>> 8) &&& 0xFF).toUInt8).toUInt16 <<< 8) = u := by
  apply UInt16.toNat.inj
  simp only [UInt16.toNat_or]
  rw [u16_byte0_nat u, u16_byte8_shifted u]
  exact nat_byte_reassembly16 u.toNat u.toNat_lt

-- ═══════════════════════════════════════════════════════════════════════════
-- Q16_16 value round-trip: ofBits (toBits q) = q
-- ═══════════════════════════════════════════════════════════════════════════

private theorem q16_ofBits_toBits (q : Q16_16) : Q16_16.ofBits (Q16_16.toBits q) = q := by
  have hlo : -2147483648 ≤ q.val := by have := q.prop.1; simp [q16MinRaw] at this; exact this
  have hhi : q.val ≤ 2147483647  := by have := q.prop.2; simp [q16MaxRaw] at this; exact this
  simp only [Q16_16.ofBits, Q16_16.toBits, Q16_16.toInt]
  have htbn : (UInt32.ofInt q.val).toNat =
      (q.val % (2 : Int) ^ 32).toNat % 2 ^ 32 := by
    simp only [UInt32.ofInt, UInt32.toNat_ofNat']
  rcases Int.lt_or_le q.val 0 with hneg | hnn
  · have htbn' : (UInt32.ofInt q.val).toNat = (q.val + 4294967296).toNat := by
      rw [htbn]; omega
    rw [htbn']
    have hge : (q.val + 4294967296).toNat ≥ 2147483648 := by omega
    simp only [if_pos hge]
    have hval : ((q.val + 4294967296).toNat : Int) - 4294967296 = q.val := by omega
    rw [hval]
    simp only [Q16_16.ofRawInt, q16MaxRaw, q16MinRaw]
    split_ifs with h1 h2
    · exact absurd h1 (by omega)
    · exact absurd h2 (by omega)
    · exact Subtype.ext rfl
  · have htbn' : (UInt32.ofInt q.val).toNat = q.val.toNat := by
      rw [htbn]; omega
    rw [htbn']
    have hlt : ¬ q.val.toNat ≥ 2147483648 := by omega
    simp only [if_neg hlt]
    have hval : (q.val.toNat : Int) = q.val := by omega
    rw [hval]
    simp only [Q16_16.ofRawInt, q16MaxRaw, q16MinRaw]
    split_ifs with h1 h2
    · exact absurd h1 (by omega)
    · exact absurd h2 (by omega)
    · exact Subtype.ext rfl

-- ═══════════════════════════════════════════════════════════════════════════
-- Q0_16 value round-trip
-- ═══════════════════════════════════════════════════════════════════════════

private theorem q0_ofRawInt_ofInt (q : Q0_16) :
    let u := UInt16.ofInt q.val
    Q0_16.ofRawInt (if u.toNat ≥ 32768 then (u.toNat : Int) - 65536 else u.toNat) = q := by
  have hlo : -32768 ≤ q.val := by have := q.prop.1; simp [q0_16MinRaw] at this; exact this
  have hhi : q.val ≤ 32767   := by have := q.prop.2; simp [q0_16MaxRaw] at this; exact this
  simp only [UInt16.ofInt]
  have htbn : (UInt16.ofNat (q.val % (2 : Int) ^ 16).toNat).toNat =
      (q.val % (2 : Int) ^ 16).toNat % 2 ^ 16 := by
    simp only [UInt16.toNat_ofNat']
  rcases Int.lt_or_le q.val 0 with hneg | hnn
  · have htbn' : (UInt16.ofNat (q.val % (2 : Int) ^ 16).toNat).toNat =
        (q.val + 65536).toNat := by rw [htbn]; omega
    rw [htbn']
    have hge : (q.val + 65536).toNat ≥ 32768 := by omega
    simp only [if_pos hge]
    have hval : ((q.val + 65536).toNat : Int) - 65536 = q.val := by omega
    rw [hval]
    simp only [Q0_16.ofRawInt, q0_16MaxRaw, q0_16MinRaw]
    split_ifs with h1 h2
    · exact absurd h1 (by omega)
    · exact absurd h2 (by omega)
    · exact Subtype.ext rfl
  · have htbn' : (UInt16.ofNat (q.val % (2 : Int) ^ 16).toNat).toNat = q.val.toNat := by
      rw [htbn]; omega
    rw [htbn']
    have hlt : ¬ q.val.toNat ≥ 32768 := by omega
    simp only [if_neg hlt]
    have hval : (q.val.toNat : Int) = q.val := by omega
    rw [hval]
    simp only [Q0_16.ofRawInt, q0_16MaxRaw, q0_16MinRaw]
    split_ifs with h1 h2
    · exact absurd h1 (by omega)
    · exact absurd h2 (by omega)
    · exact Subtype.ext rfl

-- ═══════════════════════════════════════════════════════════════════════════
-- Q16_16 wire format: 4-byte little-endian two's complement
-- ═══════════════════════════════════════════════════════════════════════════

def q16_16RowMajor : WireFormat Q16_16 rowMajor where
  encode := fun q =>
    let u := Q16_16.toBits q
    ByteArray.mk #[
      (u &&& 0xFF).toUInt8,
      ((u >>> 8) &&& 0xFF).toUInt8,
      ((u >>> 16) &&& 0xFF).toUInt8,
      ((u >>> 24) &&& 0xFF).toUInt8
    ]
  decode := fun bs =>
    if h : bs.size = 4 then
      let b0 := (bs[0]'(by omega)).toUInt32
      let b1 := (bs[1]'(by omega)).toUInt32
      let b2 := (bs[2]'(by omega)).toUInt32
      let b3 := (bs[3]'(by omega)).toUInt32
      some (Q16_16.ofBits (b0 ||| (b1 <<< 8) ||| (b2 <<< 16) ||| (b3 <<< 24)))
    else none
  encode_size := by intro q; rfl
  roundTrip := by
    intro q
    simp only [dif_pos (show (ByteArray.mk #[_, _, _, _]).size = 4 from rfl)]
    let u := Q16_16.toBits q
    change some (Q16_16.ofBits (
        ((u &&& 0xFF).toUInt8).toUInt32 |||
        ((((u >>> 8) &&& 0xFF).toUInt8).toUInt32 <<< 8) |||
        ((((u >>> 16) &&& 0xFF).toUInt8).toUInt32 <<< 16) |||
        ((((u >>> 24) &&& 0xFF).toUInt8).toUInt32 <<< 24))) = some q
    congr 1
    rw [u32_byte_reassembly u]
    exact q16_ofBits_toBits q

-- ═══════════════════════════════════════════════════════════════════════════
-- Q0_16 wire format: 2-byte little-endian
-- ═══════════════════════════════════════════════════════════════════════════

def q0_16RowMajor : WireFormat Q0_16 rowMajor where
  encode := fun q =>
    let u : UInt16 := UInt16.ofInt q.val
    ByteArray.mk #[
      (u &&& 0xFF).toUInt8,
      ((u >>> 8) &&& 0xFF).toUInt8
    ]
  decode := fun bs =>
    if h : bs.size = 2 then
      let b0 := (bs[0]'(by omega)).toUInt16
      let b1 := (bs[1]'(by omega)).toUInt16
      let u := b0 ||| (b1 <<< 8)
      let raw : Int :=
        if u.toNat ≥ 32768 then (u.toNat : Int) - 65536 else u.toNat
      some (Q0_16.ofRawInt raw)
    else none
  encode_size := by intro q; rfl
  roundTrip := by
    intro q
    simp only [dif_pos (show (ByteArray.mk #[_, _]).size = 2 from rfl)]
    let u := UInt16.ofInt q.val
    change some (Q0_16.ofRawInt (
        let b0 := ((u &&& 0xFF).toUInt8).toUInt16
        let b1 := (((u >>> 8) &&& 0xFF).toUInt8).toUInt16
        let w := b0 ||| (b1 <<< 8)
        if w.toNat ≥ 32768 then (w.toNat : Int) - 65536 else w.toNat)) = some q
    simp only []
    have hw : ((u &&& 0xFF).toUInt8).toUInt16 |||
              ((((u >>> 8) &&& 0xFF).toUInt8).toUInt16 <<< 8) = u :=
      u16_byte_reassembly u
    rw [hw]
    congr 1
    exact q0_ofRawInt_ofInt q

end SilverSight.Semantics
