module arch.system.state.resource_machine.prelude;

import prelude open;
import arch.node.types.crypto open using {Digest};

syntax alias ValueHash := Digest;
syntax alias DeltaHash := Digest;
syntax alias LabelHash := Digest;
syntax alias LogicHash := Digest;
syntax alias Quantity := Nat;
syntax alias Balance := Nat;

type Nonce :=
  mkNonce@{
    nonce : Nat;
  };

deriving instance
eqNonce : Eq Nonce;

deriving instance
ordNonce : Ord Nonce;

axiom RandSeed : Type;

trait
type FixedSize A :=
  mkFixedSize@{
    bitSize : Nat;
  };

instance
fixedSizeNat256 : FixedSize Nat :=
  mkFixedSize@{
    bitSize := 256;
  };

instance
fixedSizeByteString256 : FixedSize ByteString :=
  mkFixedSize@{
    bitSize := 256;
  };

instance
fixedSizeNonce256 : FixedSize Nonce :=
  mkFixedSize@{
    bitSize := 256;
  };

instance
fixedSizeRandSeed256 : FixedSize RandSeed :=
  mkFixedSize@{
    bitSize := 256;
  };

trait
type Arithmetic A :=
  mkArithmetic@{
    {{fixedSize}} : FixedSize A;
    add : A -> A -> A;
    sub : A -> A -> A;
  };

instance
arithmeticNat : Arithmetic Nat :=
  mkArithmetic@{
    add := \{x y := x + y};
    sub := sub;
  };

trait
type Hash H :=
  mkHash@{
    {{fixedSize}} : FixedSize H;
    hash : H -> Digest;
  };

axiom computeNatHash : Nat -> Digest;

instance
hashNat : Hash Nat :=
  mkHash@{
    fixedSize := fixedSizeNat256;
    hash := computeNatHash;
  };

axiom computeByteStringHash : ByteString -> Digest;

instance
hashByteString : Hash ByteString :=
  mkHash@{
    fixedSize := fixedSizeByteString256;
    hash := computeByteStringHash;
  };

trait
type ISet S A :=
  mkISet@{
    newSet : S;
    size : S -> Nat;
    insert : A -> S -> S;
    union : S -> S -> S;
    intersection : S -> S -> S;
    difference : S -> S -> S;
    contains : A -> S -> Bool;
  };

instance
iSetForStdSet {A} {{Ord A}} : ISet (Set A) A :=
  mkISet@{
    newSet := Set.empty;
    size := Set.size;
    insert := Set.insert;
    union := Set.union;
    intersection := Set.intersection;
    difference := Set.difference;
    contains := Set.isMember;
  };

trait
type IOrderedSet S A :=
  mkIOrderedSet@{
    {{iset}} : ISet S A;
    -- Returns elements in order as a list
    toList : S -> List A;
    -- Creates from a list, preserving order of first occurrence
    fromList : List A -> S;
  };

type OrderedSet A :=
  mkOrderedSet@{
    elements : Set A;
    permutation : List Nat;
  -- List of indices
  };

setToListWithPerm
  {A} {{Ord A}} (indices : List Nat) (elements : Set.Set A) : List A :=
  let
    -- First convert set to sorted list
    sorted : List A := Set.toList elements;
    -- Then map each index to the element at that position
    access (index : Nat) : Option A :=
      case splitAt index sorted of
        | mkPair _ (x :: _) := some x
        | mkPair _ _ := none;
  in catOptions (map access indices);

orderedSetToList {A} {{Ord A}} (s : OrderedSet A) : List A :=
  setToListWithPerm (OrderedSet.permutation s) (OrderedSet.elements s);

findPosition {A} {{Ord A}} (x : A) (elements : Set.Set A) : Nat :=
  let
    sorted : List A := Set.toList elements;
    go : List A -> Nat
      | [] := 0
      | (y :: ys) :=
        if 
          | x <= y := 0
          | else := suc (go ys);
  in go sorted;

orderedSetFromList {A} {{Ord A}} : List A -> OrderedSet A :=
  foldl
    \{acc x :=
      if 
        | Set.isMember x (OrderedSet.elements acc) := acc
        | else :=
          let
            pos := findPosition x (OrderedSet.elements acc);
          in mkOrderedSet
            (Set.insert x (OrderedSet.elements acc))
            (OrderedSet.permutation acc ++ [pos])}
    (mkOrderedSet Set.empty []);

instance
orderedSetInstance {A} {{Ord A}} : IOrderedSet (OrderedSet A) A :=
  mkIOrderedSet@{
    -- ISet instance
    iset :=
      mkISet@{
        newSet := mkOrderedSet Set.empty [];
        size := \{s := Set.size (OrderedSet.elements s)};
        insert :=
          \{x s :=
            if 
              | Set.isMember x (OrderedSet.elements s) := s
              | else :=
                mkOrderedSet
                  (Set.insert x (OrderedSet.elements s))
                  (OrderedSet.permutation s
                    ++ [Set.size (OrderedSet.elements s)])};
        -- Union
        union :=
          \{s1 s2 :=
            orderedSetFromList (orderedSetToList s1 ++ orderedSetToList s2)};
        
        -- Intersection
        intersection :=
          \{s1 s2 :=
            orderedSetFromList
              (filter
                \{x := Set.isMember x (OrderedSet.elements s2)}
                (orderedSetToList s1))};
        
        -- Difference
        difference :=
          \{s1 s2 :=
            orderedSetFromList
              (filter
                \{x := not (Set.isMember x (OrderedSet.elements s2))}
                (orderedSetToList s1))};
        
        -- Contains
        contains := \{x s := Set.isMember x (OrderedSet.elements s)};
      };
    
    -- toList
    toList := orderedSetToList;
    
    -- fromList
    fromList := orderedSetFromList;
  };

trait
type IMap M K V :=
  mkIMap@{
    emptyMap : M;
    insert : K -> V -> M -> M;
    lookup : K -> M -> Option V;
    delete : K -> M -> M;
    keys : M -> List K;
    values : M -> List V;
  };

instance
iMapForStdMap {K} {{Ord K}} {V : Type} : IMap (Map K V) K V :=
  mkIMap@{
    emptyMap := Map.empty;
    insert := \{k v m := Map.insert k v m};
    lookup := \{k m := Map.lookup k m};
    delete := \{k m := Map.delete k m};
    keys := \{m := Map.keys m};
    values := \{m := Map.values m};
  };