module Checkpoint.RegionManifest

import Checkpoint.Envelope
import Checkpoint.Schema
import Data.Bytes
import Data.SHA256
import Data.SHA256Digest
import Learning.Transition
import Model.Word64
import Representation.Partition
import Std.Foundation
import Std.List
import Std.Natural

-- A logical region is (the complete logical schema digest, its declaration
-- ordinal). Names/formats describe that region; no rank, path, stream offset or
-- device address enters the identity. Byte coordinates preserve opaque state
-- and permit splitting a serialized value provided restore joins it before use.
-- Declaration order is semantic and stable across physical placements.
family CheckpointLogicalRegion : Type 0
constructor CheckpointLogicalRegionOf
field unrestricted logicalCheckpointNameValue : Bytes
field unrestricted logicalCheckpointFormatValue : Bytes
field unrestricted logicalCheckpointExtentValue : Nat

end-family

family CheckpointCapture : Type 0
constructor CheckpointCaptureOf
field unrestricted checkpointCaptureRoundValue : (family LearningRound)
field unrestricted checkpointCaptureCursorValue : Nat

end-family

family CheckpointRegionPartition : Type 0
constructor CheckpointRegionPartitionOf
field unrestricted checkpointPartitionOrdinalValue : Nat
field unrestricted checkpointPartitionIntervalsValue : (family StdList (family LogicalInterval))

end-family

family CheckpointPreparedShard : Type 0
constructor CheckpointPreparedShardOf
field unrestricted checkpointShardIntervalValue : (family LogicalInterval)
field unrestricted checkpointShardFormatValue : Bytes
field unrestricted checkpointShardCaptureValue : (family CheckpointCapture)
field unrestricted checkpointShardBytesValue : Nat
field unrestricted checkpointShardDigestValue : Bytes

end-family

family CheckpointCapturedRegion : Type 0
constructor CheckpointCapturedRegionOf
field unrestricted checkpointCapturedOrdinalValue : Nat
field unrestricted checkpointCapturedShardsValue : (family StdList (family CheckpointPreparedShard))

end-family

def logicalCheckpointName =
  (lambda unrestricted value : (family CheckpointLogicalRegion) .
    (eliminate
      CheckpointLogicalRegion
      (lambda unrestricted current : (family CheckpointLogicalRegion) . Bytes)
      value
      (branch CheckpointLogicalRegionOf name format extent . name)))

def logicalCheckpointFormat =
  (lambda unrestricted value : (family CheckpointLogicalRegion) .
    (eliminate
      CheckpointLogicalRegion
      (lambda unrestricted current : (family CheckpointLogicalRegion) . Bytes)
      value
      (branch CheckpointLogicalRegionOf name format extent . format)))

def logicalCheckpointExtent =
  (lambda unrestricted value : (family CheckpointLogicalRegion) .
    (eliminate
      CheckpointLogicalRegion
      (lambda unrestricted current : (family CheckpointLogicalRegion) . Nat)
      value
      (branch CheckpointLogicalRegionOf name format extent . extent)))

def checkpointCaptureRound =
  (lambda unrestricted value : (family CheckpointCapture) .
    (eliminate
      CheckpointCapture
      (lambda unrestricted current : (family CheckpointCapture) . (family LearningRound))
      value
      (branch CheckpointCaptureOf round cursor . round)))

def checkpointCaptureCursor =
  (lambda unrestricted value : (family CheckpointCapture) .
    (eliminate
      CheckpointCapture
      (lambda unrestricted current : (family CheckpointCapture) . Nat)
      value
      (branch CheckpointCaptureOf round cursor . cursor)))

def checkpointPartitionOrdinal =
  (lambda unrestricted value : (family CheckpointRegionPartition) .
    (eliminate
      CheckpointRegionPartition
      (lambda unrestricted current : (family CheckpointRegionPartition) . Nat)
      value
      (branch CheckpointRegionPartitionOf ordinal intervals . ordinal)))

def checkpointPartitionIntervals =
  (lambda unrestricted value : (family CheckpointRegionPartition) .
    (eliminate
      CheckpointRegionPartition
      (lambda unrestricted current : (family CheckpointRegionPartition) .
        (family StdList (family LogicalInterval)))
      value
      (branch CheckpointRegionPartitionOf ordinal intervals . intervals)))

def checkpointShardInterval =
  (lambda unrestricted value : (family CheckpointPreparedShard) .
    (eliminate
      CheckpointPreparedShard
      (lambda unrestricted current : (family CheckpointPreparedShard) . (family LogicalInterval))
      value
      (branch CheckpointPreparedShardOf interval format capture extent digest . interval)))

def checkpointShardFormat =
  (lambda unrestricted value : (family CheckpointPreparedShard) .
    (eliminate
      CheckpointPreparedShard
      (lambda unrestricted current : (family CheckpointPreparedShard) . Bytes)
      value
      (branch CheckpointPreparedShardOf interval format capture extent digest . format)))

def checkpointShardCapture =
  (lambda unrestricted value : (family CheckpointPreparedShard) .
    (eliminate
      CheckpointPreparedShard
      (lambda unrestricted current : (family CheckpointPreparedShard) . (family CheckpointCapture))
      value
      (branch CheckpointPreparedShardOf interval format capture extent digest . capture)))

def checkpointShardBytes =
  (lambda unrestricted value : (family CheckpointPreparedShard) .
    (eliminate
      CheckpointPreparedShard
      (lambda unrestricted current : (family CheckpointPreparedShard) . Nat)
      value
      (branch CheckpointPreparedShardOf interval format capture extent digest . extent)))

def checkpointShardDigest =
  (lambda unrestricted value : (family CheckpointPreparedShard) .
    (eliminate
      CheckpointPreparedShard
      (lambda unrestricted current : (family CheckpointPreparedShard) . Bytes)
      value
      (branch CheckpointPreparedShardOf interval format capture extent digest . digest)))

def checkpointCapturedOrdinal =
  (lambda unrestricted value : (family CheckpointCapturedRegion) .
    (eliminate
      CheckpointCapturedRegion
      (lambda unrestricted current : (family CheckpointCapturedRegion) . Nat)
      value
      (branch CheckpointCapturedRegionOf ordinal shards . ordinal)))

def checkpointCapturedShards =
  (lambda unrestricted value : (family CheckpointCapturedRegion) .
    (eliminate
      CheckpointCapturedRegion
      (lambda unrestricted current : (family CheckpointCapturedRegion) .
        (family StdList (family CheckpointPreparedShard)))
      value
      (branch CheckpointCapturedRegionOf ordinal shards . shards)))

def checkpointManifestVersion : Nat =
  1

def checkpointManifestWordFits =
  (lambda unrestricted value : Nat . (naturalLessOrEqual value 18446744073709551615))

def checkpointCaptureValid =
  (lambda unrestricted capture : (family CheckpointCapture) .
    (eliminate
      LearningRound
      (lambda unrestricted current : (family LearningRound) . Nat)
      (checkpointCaptureRound capture)
      (branch
        LearningRoundOf
        epoch
        version
        .
        (naturalAnd
          (naturalNonzero epoch)
          (naturalAnd
            (checkpointManifestWordFits epoch)
            (naturalAnd
              (checkpointManifestWordFits version)
              (checkpointManifestWordFits (checkpointCaptureCursor capture))))))))

def checkpointCaptureEqual =
  (lambda unrestricted expected : (family CheckpointCapture) .
    (lambda unrestricted actual : (family CheckpointCapture) .
      (naturalAnd
        (learningRoundEqual (checkpointCaptureRound expected) (checkpointCaptureRound actual))
        (naturalEqual (checkpointCaptureCursor expected) (checkpointCaptureCursor actual)))))

def checkpointLogicalRegionValid =
  (lambda unrestricted region : (family CheckpointLogicalRegion) .
    (naturalAnd
      (naturalNonzero (bytes-length (logicalCheckpointName region)))
      (naturalAnd
        (naturalNonzero (bytes-length (logicalCheckpointFormat region)))
        (naturalAnd
          (checkpointManifestWordFits (bytes-length (logicalCheckpointName region)))
          (naturalAnd
            (checkpointManifestWordFits (bytes-length (logicalCheckpointFormat region)))
            (checkpointManifestWordFits (logicalCheckpointExtent region)))))))

def checkpointLogicalRegionsValid =
  (stdListFold
    (family CheckpointLogicalRegion)
    Nat
    (lambda unrestricted region : (family CheckpointLogicalRegion) .
      (lambda unrestricted rest : Nat . (naturalAnd (checkpointLogicalRegionValid region) rest)))
    1)

-- Empty logical regions have an empty shard list. Empty shard records are
-- refused: they add no ownership, and otherwise permit duplicate empty claims.
def checkpointNonemptyIntervals =
  (stdListFold
    (family LogicalInterval)
    Nat
    (lambda unrestricted interval : (family LogicalInterval) .
      (lambda unrestricted rest : Nat .
        (naturalAnd (naturalLess (intervalBegin interval) (intervalEnd interval)) rest)))
    1)

-- One linear pass over regions and fragments, never over logical byte counts.
-- A group ordinal identifies a field in this schema, not a participant.
def checkpointPartitionsFrom =
  (lambda unrestricted regions : (family StdList (family CheckpointLogicalRegion)) .
    (eliminate
      StdList
      (lambda unrestricted current : (family StdList (family CheckpointLogicalRegion)) .
        (pi unrestricted ordinal : Nat .
          (pi unrestricted parts : (family StdList (family CheckpointRegionPartition)) . Nat)))
      regions
      (branch
        StdListEmpty
        .
        (lambda unrestricted ordinal : Nat .
          (lambda unrestricted parts : (family StdList (family CheckpointRegionPartition)) .
            (stdBoolToNatural (stdListIsEmpty (family CheckpointRegionPartition) parts)))))
      (branch
        StdListCons
        region
        tail
        induction
        .
        (lambda unrestricted ordinal : Nat .
          (lambda unrestricted parts : (family StdList (family CheckpointRegionPartition)) .
            (eliminate
              StdList
              (lambda unrestricted current : (family StdList (family CheckpointRegionPartition)) .
                Nat)
              parts
              (branch StdListEmpty . 0)
              (branch
                StdListCons
                part
                after
                ignored
                .
                (naturalAnd
                  (naturalEqual ordinal (checkpointPartitionOrdinal part))
                  (naturalAnd
                    (intervalPartitionValid
                      (logicalCheckpointExtent region)
                      (checkpointPartitionIntervals part))
                    (naturalAnd
                      (checkpointNonemptyIntervals (checkpointPartitionIntervals part))
                      (induction (succ ordinal) after)))))))))))

def checkpointPartitionsValid =
  (lambda unrestricted regions : (family StdList (family CheckpointLogicalRegion)) .
    (lambda unrestricted parts : (family StdList (family CheckpointRegionPartition)) .
      (naturalAnd
        (checkpointLogicalRegionsValid regions)
        (checkpointPartitionsFrom regions 0 parts))))

def checkpointPreparedShardValid =
  (lambda unrestricted region : (family CheckpointLogicalRegion) .
    (lambda unrestricted capture : (family CheckpointCapture) .
      (lambda unrestricted shard : (family CheckpointPreparedShard) .
        (naturalAnd
          (checkpointCaptureEqual capture (checkpointShardCapture shard))
          (naturalAnd
            (bytes-equal (logicalCheckpointFormat region) (checkpointShardFormat shard))
            (naturalAnd
              (naturalEqual
                (checkpointShardBytes shard)
                (intervalExtent (checkpointShardInterval shard)))
              (naturalAnd
                (naturalEqual
                  (bytes-length (checkpointShardDigest shard))
                  checkpointEnvelopeDigestBytes)
                (naturalAnd
                  (intervalWithin (logicalCheckpointExtent region) (checkpointShardInterval shard))
                  (naturalLess
                    (intervalBegin (checkpointShardInterval shard))
                    (intervalEnd (checkpointShardInterval shard)))))))))))

def checkpointCapturedRegionValid =
  (lambda unrestricted region : (family CheckpointLogicalRegion) .
    (lambda unrestricted ordinal : Nat .
      (lambda unrestricted capture : (family CheckpointCapture) .
        (lambda unrestricted group : (family CheckpointCapturedRegion) .
          (let unrestricted shards =
            (checkpointCapturedShards group)
            in
            (naturalAnd
              (naturalEqual ordinal (checkpointCapturedOrdinal group))
              (naturalAnd
                (intervalPartitionValid
                  (logicalCheckpointExtent region)
                  (stdListMap
                    (family CheckpointPreparedShard)
                    (family LogicalInterval)
                    checkpointShardInterval
                    shards))
                (stdListFold
                  (family CheckpointPreparedShard)
                  Nat
                  (lambda unrestricted shard : (family CheckpointPreparedShard) .
                    (lambda unrestricted rest : Nat .
                      (naturalAnd (checkpointPreparedShardValid region capture shard) rest)))
                  1
                  shards))))))))

def checkpointManifestFrom =
  (lambda unrestricted capture : (family CheckpointCapture) .
    (lambda unrestricted regions : (family StdList (family CheckpointLogicalRegion)) .
      (eliminate
        StdList
        (lambda unrestricted current : (family StdList (family CheckpointLogicalRegion)) .
          (pi unrestricted ordinal : Nat .
            (pi unrestricted groups : (family StdList (family CheckpointCapturedRegion)) . Nat)))
        regions
        (branch
          StdListEmpty
          .
          (lambda unrestricted ordinal : Nat .
            (lambda unrestricted groups : (family StdList (family CheckpointCapturedRegion)) .
              (stdBoolToNatural (stdListIsEmpty (family CheckpointCapturedRegion) groups)))))
        (branch
          StdListCons
          region
          tail
          induction
          .
          (lambda unrestricted ordinal : Nat .
            (lambda unrestricted groups : (family StdList (family CheckpointCapturedRegion)) .
              (eliminate
                StdList
                (lambda unrestricted current : (family StdList (family CheckpointCapturedRegion)) .
                  Nat)
                groups
                (branch StdListEmpty . 0)
                (branch
                  StdListCons
                  group
                  after
                  ignored
                  .
                  (naturalAnd
                    (checkpointCapturedRegionValid region ordinal capture group)
                    (induction (succ ordinal) after))))))))))

def checkpointManifestValid =
  (lambda unrestricted version : Nat .
    (lambda unrestricted regions : (family StdList (family CheckpointLogicalRegion)) .
      (lambda unrestricted capture : (family CheckpointCapture) .
        (lambda unrestricted groups : (family StdList (family CheckpointCapturedRegion)) .
          (naturalAnd
            (naturalEqual version checkpointManifestVersion)
            (naturalAnd
              (checkpointLogicalRegionsValid regions)
              (naturalAnd
                (checkpointCaptureValid capture)
                (checkpointManifestFrom capture regions 0 groups))))))))

-- Length-prefix every semantic byte string. Builders materialize once rather
-- than repeatedly copying an ever-growing schema suffix. Physical offsets and
-- participant identities are deliberately absent from these canonical bytes.
def checkpointManifestSized =
  (lambda unrestricted value : Bytes .
    (bytes-builder-append
      (bytes-builder-chunk (checkpointEnvelopeWord (bytes-length value)))
      (bytes-builder-chunk value)))

def checkpointRegionEncoding =
  (lambda unrestricted region : (family CheckpointLogicalRegion) .
    (bytes-builder-append
      (checkpointManifestSized (logicalCheckpointName region))
      (bytes-builder-append
        (checkpointManifestSized (logicalCheckpointFormat region))
        (bytes-builder-chunk (checkpointEnvelopeWord (logicalCheckpointExtent region))))))

def checkpointLogicalSchemaBytes =
  (lambda unrestricted regions : (family StdList (family CheckpointLogicalRegion)) .
    (bytes-builder-build
      (bytes-builder-append
        (bytes-builder-chunk b"ALPHA-LOGICAL-CHECKPOINT-1\x00")
        (bytes-builder-append
          (bytes-builder-chunk
            (checkpointEnvelopeWord (stdListLength (family CheckpointLogicalRegion) regions)))
          (stdListFold
            (family CheckpointLogicalRegion)
            BytesBuilder
            (lambda unrestricted region : (family CheckpointLogicalRegion) .
              (lambda unrestricted rest : BytesBuilder .
                (bytes-builder-append (checkpointRegionEncoding region) rest)))
            (bytes-builder-empty)
            regions)))))

def checkpointLogicalSchemaDigest =
  (lambda unrestricted regions : (family StdList (family CheckpointLogicalRegion)) .
    (sha256RawDigestOrEmpty (checkpointLogicalSchemaBytes regions)))

-- Adapt the maintained field schema without its storage or stream offsets.
-- Field kind/name, element count, domain and initialization remain semantic.
-- These regions use byte coordinates; a native consumer must reconstruct and
-- validate its declared representation before exposing a partial typed value.
def checkpointLogicalRegionFromField =
  (lambda unrestricted field : (family CheckpointField) .
    (eliminate
      CheckpointField
      (lambda unrestricted current : (family CheckpointField) . (family CheckpointLogicalRegion))
      field
      (branch
        CheckpointFieldValue
        kind
        name
        storage
        stream
        extent
        elements
        domain
        initialization
        .
        (constructor
          CheckpointLogicalRegion
          CheckpointLogicalRegionOf
          (bytes-builder-build
            (bytes-builder-append
              (checkpointManifestSized (checkpointFieldKindName kind))
              (checkpointManifestSized name)))
          (bytes-builder-build
            (bytes-builder-append
              (checkpointManifestSized domain)
              (bytes-builder-append
                (bytes-builder-chunk (checkpointEnvelopeWord (modelWord64Natural elements)))
                (checkpointManifestSized initialization))))
          (modelWord64Natural extent)))))

def checkpointLogicalRegionsFromFields =
  (lambda unrestricted fields : (family CheckpointFields) .
    (eliminate
      CheckpointFields
      (lambda unrestricted current : (family CheckpointFields) .
        (family StdList (family CheckpointLogicalRegion)))
      fields
      (branch
        CheckpointFieldsEnd
        .
        (constructor StdList StdListEmpty (family CheckpointLogicalRegion)))
      (branch
        CheckpointFieldsNext
        head
        tail
        induction
        .
        (constructor
          StdList
          StdListCons
          (family CheckpointLogicalRegion)
          (checkpointLogicalRegionFromField head)
          induction))))