import Mathlib.Data.Set.Basic
import Mathlib.Data.Set.Insert
import InductiveVerification.Message

-- Define the `Event` type
inductive Event
| Says : Agent → Agent → Msg → Event
| Gets : Agent → Msg → Event
| Notes : Agent → Msg → Event

-- Define the `initState` function
class HasInitState (α : Type) where
  initState : α → Set Msg
  
variable [ hasInitStateAgent : HasInitState Agent ]
  
open HasInitState
attribute [simp] initState
  
-- Define the `bad` set
abbrev DecidableMem ( A : Set Agent ) := (a : Agent) →  Decidable (a ∈ A)
class Bad where
  bad : Set Agent 
  decidableBad : DecidableMem bad
  Spy_in_bad : Agent.Spy ∈ bad
  Server_not_bad : Agent.Server ∉ bad

instance [Bad] : DecidableMem Bad.bad := Bad.decidableBad
open Bad

attribute [simp, grind .] Spy_in_bad
attribute [simp] Server_not_bad

instance decidableAgentEq : DecidableEq Agent :=
  λ a b =>
    match a, b with
    | Agent.Spy, Agent.Spy => isTrue rfl
    | Agent.Spy, Agent.Server => isFalse (λ h => by contradiction)
    | Agent.Spy, Agent.Friend m => isFalse (λ h => by contradiction)
    | Agent.Server, Agent.Spy => isFalse (λ h => by contradiction)
    | Agent.Server, Agent.Server => isTrue rfl
    | Agent.Server, Agent.Friend m => isFalse (λ h => by contradiction)
    | Agent.Friend n, Agent.Spy => isFalse (λ h => by contradiction)
    | Agent.Friend n, Agent.Server => isFalse (λ h => by contradiction)
    | Agent.Friend n, Agent.Friend m =>
      if h : n = m then isTrue (congrArg Agent.Friend h) else isFalse (λ h' => h (Agent.Friend.inj h'))

-- instance decidableBad (a: Agent) : Decidable (a ∈ bad) :=
-- by rw[bad]; apply Set.decidableSingleton

-- Define the `knows` function
def knows [Bad] : Agent → List Event → Set Msg
| A, [] => initState A
| Agent.Spy, Event.Says _ _ X :: evs => insert X (knows Agent.Spy evs)
| Agent.Spy, Event.Gets _ _ :: evs => knows Agent.Spy evs
| Agent.Spy, Event.Notes A X :: evs =>
  if A ∈ bad then insert X (knows Agent.Spy evs) else knows Agent.Spy evs
| A, Event.Says A' _ X :: evs =>
  if A = A' then insert X (knows A evs) else knows A evs
| A, Event.Gets A' X :: evs =>
  if A = A' then insert X (knows A evs) else knows A evs
| A, Event.Notes A' X :: evs =>
  if A = A' then insert X (knows A evs) else knows A evs
  
attribute [simp] knows

-- Define the `spies` abbreviation
abbrev spies (evs : List Event) [Bad] : Set Msg := knows Agent.Spy evs
attribute [simp] spies

-- Define the `used` function
def used : List Event → Set Msg
| [] => ⋃ (B : Agent), parts (initState B)
| Event.Says _ _ X :: evs => parts {X} ∪ used evs
| Event.Gets _ _ :: evs => used evs
| Event.Notes _ X :: evs => parts {X} ∪ used evs

-- Lemmas for `used`

 lemma used_mono : used (evs) ⊆ used (ev :: evs) := by
  cases ev <;> simp[used]
  
 lemma Notes_imp_used {A : Agent} {X : Msg} {evs : List Event} :
  Event.Notes A X ∈ evs → X ∈ used evs := by
  induction evs with
  | nil => intro h; cases h
  | cons ev evs ih =>
    intro h
    cases h with
    | tail => cases ev
              · simp [used]; right; aapply ih
              · simp [used]; aapply ih
              · simp [used]; right; aapply ih
    | head => simp [used]; left; apply parts.inj; simp

 lemma Says_imp_used {A B : Agent} {X : Msg} {evs : List Event} :
  Event.Says A B X ∈ evs → X ∈ used evs := by
  induction evs with
  | nil => intro h; cases h
  | cons ev evs ih =>
    intro h
    cases h with
    | tail => cases ev
              · simp [used]; right; aapply ih
              · simp [used]; aapply ih
              · simp [used]; right; aapply ih
    | head => simp [used]; left; apply parts.inj; simp

-- Knowledge subset rules
lemma knows_subset_knows_Cons [Bad] {A : Agent} {ev : Event} {evs : List Event} :
  knows A (evs) ⊆ knows A (ev :: evs) := by
  by_cases h :A = Agent.Spy
  · rw[h]
    intro _ _
    cases ev
    · simp [knows]; right; assumption
    · simp [knows]; assumption
    · simp [knows]; split_ifs
      · right; assumption
      · assumption
  · intro _ _
    cases ev <;> (simp [knows]; split_ifs; right; assumption; assumption)

lemma initState_subset_knows [Bad] :
  ∀ {A : Agent} {evs : List Event},
    initState A ⊆ knows A evs := by
  intro A evs
  induction evs with
  | nil => simp [knows]
  | cons e evs ih =>
    apply subset_trans
    · exact ih
    · apply knows_subset_knows_Cons

-- Lemmas for `knows`
@[simp]
lemma knows_Spy_Says [Bad] {A B : Agent} {X : Msg} {evs : List Event} :
  knows Agent.Spy (Event.Says A B X :: evs) = insert X (knows Agent.Spy evs) := by
  simp [knows]

@[simp]
lemma knows_Spy_Notes [Bad] {A : Agent} {X : Msg} {evs : List Event} :
  knows Agent.Spy (Event.Notes A X :: evs) =
    if A ∈ bad then insert X (knows Agent.Spy evs) else knows Agent.Spy evs := by
  simp [knows]

@[simp]
lemma knows_Spy_Gets [Bad] {A : Agent} {X : Msg} {evs : List Event} :
  knows Agent.Spy (Event.Gets A X :: evs) = knows Agent.Spy evs := by
  simp [knows]

lemma Says_imp_knows_Spy [Bad] {A B : Agent} {X : Msg} {evs : List Event} :
  Event.Says A B X ∈ evs → X ∈ knows Agent.Spy evs := by
  induction evs with
  | nil => intro h; cases h
  | cons ev evs ih =>
    intro h
    cases h with
    | tail h => cases ev
                · simp [knows]; right; aapply ih
                · simp [knows]; aapply ih
                · simp [knows]; split_ifs
                  · right; aapply ih
                  · aapply ih
    | head h => simp [knows];

lemma Notes_imp_knows_Spy [Bad] {A : Agent} {X : Msg} {evs : List Event} :
  Event.Notes A X ∈ evs → A ∈ bad → X ∈ knows Agent.Spy evs := by
  induction evs with
  | nil => intro h; cases h
  | cons ev evs ih =>
    intro h h_bad
    cases h with
    | head => simp [knows]; split_ifs; left; trivial
    | tail _ h => apply knows_subset_knows_Cons; apply ih; apply h; apply h_bad

-- Elimination rules: derive contradictions from old Says events containing
-- items known to be fresh
@[grind ., grind! .]
lemma Says_imp_parts_knows_Spy [Bad] :
  ∀ {A B : Agent} {X : Msg} {evs : List Event},
    Event.Says A B X ∈ evs → X ∈ parts (knows Agent.Spy evs) := by
  intro A B X evs h
  apply parts.inj
  apply Says_imp_knows_Spy
  exact h

lemma knows_Spy_partsEs [Bad] :
  ∀ {A B : Agent} {X : Msg} {evs : List Event},
    Event.Says A B X ∈ evs → X ∈ parts (knows Agent.Spy evs) := by
  exact Says_imp_parts_knows_Spy

lemma Says_imp_analz_Spy [InvKey] [Bad]
  {A B : Agent} {X : Msg} {evs : List Event} :
    Event.Says A B X ∈ evs → X ∈ analz (knows Agent.Spy evs) := by
  intro h
  apply analz.inj
  apply Says_imp_knows_Spy
  exact h

-- Compatibility for the old "spies" function
lemma spies_partsEs [Bad] :
  ∀ {A B : Agent} {X : Msg} {evs : List Event},
    Event.Says A B X ∈ evs → X ∈ parts (knows Agent.Spy evs) := by
  exact knows_Spy_partsEs

lemma Says_imp_spies [Bad] :
  ∀ {A B : Agent} {X : Msg} {evs : List Event},
    Event.Says A B X ∈ evs → X ∈ knows Agent.Spy evs := by
  exact Says_imp_knows_Spy

lemma parts_insert_spies [Bad] :
    parts (insert X (knows Agent.Spy evs)) = parts {X} ∪ parts (knows Agent.Spy evs) :=
by
  apply parts_insert

lemma analz_spies_mono [InvKey] [Bad]
  { h : M ∈ analz (knows Agent.Spy evs) } :
  M ∈ analz (knows Agent.Spy (ev :: evs))
:= by
  aapply analz_mono; exact knows_subset_knows_Cons

lemma analz_spies_mono_neg [InvKey] [Bad]
  { h : M ∉ analz (knows Agent.Spy (ev :: evs)) } :
  M ∉ analz (knows Agent.Spy evs)
:= by
  intro h₁; apply h; aapply analz_spies_mono

-- Knowledge of Agents
lemma knows_subset_knows_Says [Bad] :
  ∀ {A A' B : Agent} {X : Msg} {evs : List Event},
    knows A evs ⊆ knows A (Event.Says A' B X :: evs) := by
  intro A A' B X evs
  apply knows_subset_knows_Cons

lemma knows_subset_knows_Notes [Bad] :
  ∀ {A A' : Agent} {X : Msg} {evs : List Event},
    knows A evs ⊆ knows A (Event.Notes A' X :: evs) := by
  intro A A' X evs
  apply knows_subset_knows_Cons

lemma knows_subset_knows_Gets [Bad] :
  ∀ {A A' : Agent} {X : Msg} {evs : List Event},
    knows A evs ⊆ knows A (Event.Gets A' X :: evs) := by
  intro A A' X evs
  apply knows_subset_knows_Cons

-- Agents know what they say
lemma Says_imp_knows [Bad] :
  ∀ {A B : Agent} {X : Msg} {evs : List Event},
    Event.Says A B X ∈ evs → X ∈ knows A evs := by
  intro A B X evs h
  induction evs with
  | nil => cases h
  | cons ev evs ih =>
    cases h with
    | head => by_cases h: A = Agent.Spy
              · rw [h]; simp [knows]
              · simp [knows]
    | tail => by_cases h: A = Agent.Spy
              · rw[h]; cases ev
                · simp [knows]; right; rw[←h]; aapply ih
                · simp [knows]; rw[←h]; aapply ih
                · simp [knows]; split_ifs
                  · right; rw[←h]; aapply ih
                  · rw[←h]; aapply ih
              · cases ev <;> (
                simp [knows]; split_ifs; right; aapply ih; aapply ih)
    
-- Agents know what they note
lemma Notes_imp_knows [Bad] :
  ∀ {A : Agent} {X : Msg} {evs : List Event},
    Event.Notes A X ∈ evs → X ∈ knows A evs := by
  intro A X evs h
  induction evs with
  | nil => cases h
  | cons ev evs ih =>
    cases h with
    | head => by_cases h: A = Agent.Spy
              · rw [h]; simp [knows, Spy_in_bad]
              · simp [knows]
    | tail => by_cases h: A = Agent.Spy
              · rw[h]; cases ev
                · simp [knows]; right; rw[←h]; aapply ih
                · simp [knows]; rw[←h]; aapply ih
                · simp [knows]; split_ifs
                  · right; rw[←h]; aapply ih
                  · rw[←h]; aapply ih
              · cases ev <;> (
                simp [knows]; split_ifs; right; aapply ih; aapply ih)

-- Agents know what they receive
lemma Gets_imp_knows_agents [Bad] :
  ∀ {A : Agent} {X : Msg} {evs : List Event},
    A ≠ Agent.Spy → Event.Gets A X ∈ evs → X ∈ knows A evs := by
  intro A X evs h_ne h
  induction evs with
  | nil => cases h
  | cons ev evs ih =>
    cases h with
    | head => simp [knows]
    | tail => cases ev <;> (simp [knows]; split_ifs; right; aapply ih; aapply ih)

-- What agents DIFFERENT FROM Spy know
lemma knows_imp_Says_Gets_Notes_initState [Bad] :
  ∀ {A : Agent} {X : Msg} {evs : List Event},
    X ∈ knows A evs → A ≠ Agent.Spy →
    ∃ B, Event.Says A B X ∈ evs ∨ Event.Gets A X ∈ evs ∨ Event.Notes A X ∈ evs ∨ X ∈ initState A := by
  intro _ _ evs h _
  induction evs with
  | nil => simp [knows] at h; exists Agent.Server; simp_all
  | cons ev evs ih =>
    cases ev
    · simp [knows] at h; split_ifs at h with eq
      · cases h with
        | inl h₁ => grind
        | inr h' => apply ih at h'; cases h' with | intro A'; exists A'; grind
      · apply ih at h; cases h with | intro A'; exists A'; grind
    · simp [knows] at h; split_ifs at h with eq
      · cases h with
        | inl h₁ => simp[h₁, eq]; exact ⟨Agent.Server⟩
        | inr h' => apply ih at h'; cases h' with | intro A'; exists A'; grind
      · apply ih at h; cases h with | intro A'; exists A'; grind
    · simp [knows] at h; split_ifs at h with eq
      · cases h with
        | inl h₁ => simp[h₁, eq]; exact ⟨Agent.Server⟩
        | inr h' => apply ih at h'; cases h' with | intro A'; exists A'; grind
      · apply ih at h; cases h with | intro A'; exists A'; grind

-- auxiliary lemma

lemma knows_Spy_imp_Says_Notes_initState_aux
  [Bad] {X : Msg} {ev : Event} {evs : List Event} :
  (∃ A B, Event.Says A B X ∈ evs ∨ Event.Notes A X ∈ evs ∨ X ∈ initState Agent.Spy)
  → (∃ A B,
      Event.Says A B X ∈ ev :: evs ∨ Event.Notes A X ∈ ev :: evs ∨ X ∈ initState Agent.Spy
  ) := by
    intro h
    cases h with | intro A h =>
    cases h with | intro B h => 
    exists A; exists B
    cases h with
    | inl => left; right; assumption
    | inr h => right; cases h
               · left; right; assumption
               · right; assumption

-- What the Spy knows
lemma knows_Spy_imp_Says_Notes_initState [Bad] {X : Msg} {evs : List Event} :
    X ∈ knows Agent.Spy evs →
    ∃ A B, Event.Says A B X ∈ evs ∨ Event.Notes A X ∈ evs ∨ X ∈ initState Agent.Spy := by
  intro h
  induction evs with
  | nil => simp [knows] at h; exists Agent.Server; exists Agent.Server; simp_all
  | cons ev evs ih =>
    cases ev with
    | Says A' B' =>
      simp [knows] at h; cases h with
      | inl => exists A'; exists B'; simp_all;
      | inr h => 
        apply ih at h; aapply knows_Spy_imp_Says_Notes_initState_aux
    | Gets => 
      simp [knows] at h; apply ih at h; aapply knows_Spy_imp_Says_Notes_initState_aux
    | Notes A' X' => 
      simp [knows] at h; split_ifs at h with A'_bad
      · cases h with
        | inl h => rw[h]; exists A'; exists A'; simp
        | inr h => apply ih at h; aapply knows_Spy_imp_Says_Notes_initState_aux
      · apply ih at h; aapply knows_Spy_imp_Says_Notes_initState_aux

-- Parts of what the Spy knows are a subset of what is used
@[grind! .]
lemma parts_knows_Spy_subset_used [Bad] :
    parts (knows Agent.Spy evs) ⊆ used evs := by
  induction evs with
  | nil => simp [used, knows]; intro _ _; simp; exists Agent.Spy;
  | cons ev evs ih =>
    cases ev
    · simp[used, knows]; apply subset_trans; apply ih; simp
    · simp[used, knows]; assumption
    · simp[used, knows]; split_ifs with ABad
      · simp; apply subset_trans; apply ih; simp
      · apply subset_trans; apply ih; simp
      
lemma parts_knows_Spy_subset_used_neg [Bad] :
  M ∉ used evs → M ∉ parts (knows Agent.Spy evs) := by
  intro h₁ h₂; apply h₁; aapply parts_knows_Spy_subset_used
      
lemma analz_knows_Spy_subset_used [Bad] [InvKey] :
  analz (knows Agent.Spy evs) ⊆ used evs
:= by
  apply subset_trans ( b := parts (knows Agent.Spy evs) )
  · exact analz_subset_parts
  · exact parts_knows_Spy_subset_used

lemma analz_knows_Spy_subset_used_neg [Bad] [InvKey] :
  M ∉ used evs → M ∉ analz (knows Agent.Spy evs) := by
  intro h₁ h₂; apply h₁; aapply analz_knows_Spy_subset_used

-- Parts of what the Spy knows are a subset of what is used
lemma usedI [Bad] :
    X ∈ parts (knows Agent.Spy evs) → X ∈ used evs := by
  intro h
  apply parts_knows_Spy_subset_used
  exact h

-- Initial state messages are part of the used set

lemma initState_into_used {B : Agent} {evs : List Event} :
    parts (initState B) ⊆ used evs := by
  induction evs with
  | nil => simp [used]; apply Set.subset_iUnion (parts ∘ initState) B
  | cons ev =>
    cases ev <;> (simp[used]; try apply Set.subset_union_of_subset_right)
    all_goals assumption

-- Simplification rules for `used`

@[simp]
lemma used_Says {A B : Agent} {X : Msg} {evs : List Event} :
    used (Event.Says A B X :: evs) = parts {X} ∪ used evs := by
  simp [used]


@[simp]
lemma used_Notes {A : Agent} {X : Msg} {evs : List Event} :
    used (Event.Notes A X :: evs) = parts {X} ∪ used evs := by
  simp [used]


@[simp]
lemma used_Gets {A : Agent} {X : Msg} {evs : List Event} :
    used (Event.Gets A X :: evs) = used evs := by
  simp [used]


lemma used_nil_subset {evs : List Event} :
    used [] ⊆ used evs := by
  have B := Agent.Server
  induction evs with
  | nil => simp
  | cons head tail ih =>
  apply subset_trans (b := used tail)
  · assumption
  · exact used_mono
  
-- Keys for parts insert
open InvKey

omit hasInitStateAgent in

lemma keysFor_parts_insert {K : Key} {X : Msg} {G H : Set Msg} [InvKey]:
    K ∈ keysFor (parts (insert X G)) →
    X ∈ synth (analz H) →
    K ∈ keysFor (parts (G ∪ H)) ∪ invKey '' {K | Msg.Key K ∈ parts H} := by
  intro h h₂
  rw[parts_union, keysFor_union]
  rcases h with ⟨K', ⟨⟨X', h⟩, r⟩⟩
  by_cases KkfpG: K∈ keysFor (parts G)
  · left; left; assumption
  · rw[Set.union_assoc]; right; rw[←keysFor_synth]; rw[parts_insert] at h
    cases h with
    | inl h => 
      apply Crypt_imp_invKey_keysFor at h; rw[r] at h
      apply keysFor_mono (a := parts {X})
      · apply synth_subset_iff.mpr; rw[←synth_subset_iff]
        apply subset_trans (b := parts (synth (analz H)))
        · apply parts_mono; simp; assumption
        · rw[parts_synth, parts_analz]; apply sup_le
          · exact synth_increasing
          · apply synth_mono; exact analz_subset_parts
      · assumption
    | inr h => apply Crypt_imp_invKey_keysFor at h; rw[r] at h; contradiction