module Runtime.TextArtifactRequest

import Data.Bytes
import Model.Word64
import Std.Natural

-- One request record drives both operations of every R6 text artifact.  The
-- native host consumes exactly one path argument naming this record; no
-- working-directory convention or positional train-only ABI is involved.
family TextArtifactOperation : Type 0
constructor TextArtifactTrain
constructor TextArtifactPredict

end-family

family TextArtifactRequest : Type 0
constructor TextArtifactRequestValue
field unrestricted textArtifactOperation : (family TextArtifactOperation)
field unrestricted textArtifactInputPath : Bytes
field unrestricted textArtifactCheckpointInputPath : Bytes
field unrestricted textArtifactCheckpointOutputPath : Bytes
field unrestricted textArtifactPredictionsOutputPath : Bytes
field unrestricted textArtifactResultOutputPath : Bytes
field unrestricted textArtifactCheckpointTemporaryPath : Bytes
field unrestricted textArtifactCheckpointDirectoryPath : Bytes

end-family

def textArtifactRequestMagic : Bytes =
  b"ALPHATX1"

-- Version 2 adds the checkpoint's temporary path and its directory: the
-- host writes the checkpoint to the temporary, then renames it over the
-- output path and syncs the directory (Checkpoint.Envelope's publication).
def textArtifactRequestVersion : Nat =
  2

-- The operation is deliberately represented by a harmless Linux syscall
-- number.  After the shared prediction prefix, getpid returns and training
-- continues; exit_group(0) terminates a predict request before the optimizer.
-- This keeps one static ELF and one native command stream without introducing
-- a second host interpreter or CPU-learning branch.
def textArtifactTrainOperationWord : Nat =
  39

def textArtifactPredictOperationWord : Nat =
  231

def textArtifactPathExtent : Nat =
  256

def textArtifactPathMaximumLength : Nat =
  255

def textArtifactRequestPaths : Nat =
  7

-- magic[8] + version:u64 + operation:u64 + seven NUL-padded paths[256].
def textArtifactRequestEncodedLength : Nat =
  (naturalAdd 24 (naturalMultiply textArtifactRequestPaths textArtifactPathExtent))

def textArtifactZeroBytes =
  (lambda unrestricted count : Nat .
    (dataBytesRepeatByte (byte 0) count))

def textArtifactOperationWord =
  (lambda unrestricted operation : (family TextArtifactOperation) .
    (eliminate
      TextArtifactOperation
      (lambda unrestricted current : (family TextArtifactOperation) . Nat)
      operation
      (branch TextArtifactTrain . textArtifactTrainOperationWord)
      (branch TextArtifactPredict . textArtifactPredictOperationWord)))

def textArtifactPathValid =
  (lambda unrestricted path : Bytes .
    (naturalAnd
      (naturalNonzero (bytes-length path))
      (naturalLessOrEqual (bytes-length path) textArtifactPathMaximumLength)))

def textArtifactEncodePath =
  (lambda unrestricted path : Bytes .
    (bytes-append
      path
      (textArtifactZeroBytes
        (naturalSaturatingSubtract textArtifactPathExtent (bytes-length path)))))

def textArtifactRequestFieldsValid =
  (lambda unrestricted request : (family TextArtifactRequest) .
    (eliminate
      TextArtifactRequest
      (lambda unrestricted current : (family TextArtifactRequest) . Nat)
      request
      (branch
        TextArtifactRequestValue
        operation
        input
        checkpointInput
        checkpointOutput
        predictionsOutput
        resultOutput
        checkpointTemporary
        checkpointDirectory
        .
        (naturalAnd
          (textArtifactPathValid input)
          (naturalAnd
            (textArtifactPathValid checkpointInput)
            (naturalAnd
              (textArtifactPathValid checkpointOutput)
              (naturalAnd
                (textArtifactPathValid predictionsOutput)
                (naturalAnd
                  (textArtifactPathValid resultOutput)
                  (naturalAnd
                    (textArtifactPathValid checkpointTemporary)
                    (textArtifactPathValid checkpointDirectory))))))))))

def textArtifactRequestEncodeUnchecked =
  (lambda unrestricted request : (family TextArtifactRequest) .
    (eliminate
      TextArtifactRequest
      (lambda unrestricted current : (family TextArtifactRequest) . Bytes)
      request
      (branch
        TextArtifactRequestValue
        operation
        input
        checkpointInput
        checkpointOutput
        predictionsOutput
        resultOutput
        checkpointTemporary
        checkpointDirectory
        .
        (bytes-append
          textArtifactRequestMagic
          (bytes-append
            (dataBytesWord64LE (modelWord64FromNaturalTruncated textArtifactRequestVersion))
            (bytes-append
              (dataBytesWord64LE
                (modelWord64FromNaturalTruncated (textArtifactOperationWord operation)))
              (bytes-append
                (textArtifactEncodePath input)
                (bytes-append
                  (textArtifactEncodePath checkpointInput)
                  (bytes-append
                    (textArtifactEncodePath checkpointOutput)
                    (bytes-append
                      (textArtifactEncodePath predictionsOutput)
                      (bytes-append
                        (textArtifactEncodePath resultOutput)
                        (bytes-append
                          (textArtifactEncodePath checkpointTemporary)
                          (textArtifactEncodePath checkpointDirectory)))))))))))))

-- Invalid requests have no wire representation.  This makes the encoder
-- fail-closed while keeping construction and static evaluation lightweight.
def textArtifactRequestEncode =
  (lambda unrestricted request : (family TextArtifactRequest) .
    (nat-eliminate
      (lambda unrestricted current : Nat . Bytes)
      b""
      (lambda unrestricted predecessor : Nat .
        (lambda unrestricted induction : Bytes .
          (textArtifactRequestEncodeUnchecked request)))
      (textArtifactRequestFieldsValid request)))

def textArtifactEncodedFieldTerminated =
  (lambda unrestricted offset : Nat .
    (lambda unrestricted raw : Bytes .
      (naturalAnd
        (naturalNonzero
          (byte-to-nat (dataBytesByteAtValidated raw offset)))
        (naturalIsZero
          (byte-to-nat
            (dataBytesByteAtValidated
              raw
              (naturalAdd offset textArtifactPathMaximumLength)))))))

-- Native hosts perform these same checks before opening any supplied path.
-- Path fields must begin non-NUL and terminate inside their fixed extent.
def textArtifactRequestBytesValid =
  (lambda unrestricted raw : Bytes .
    (naturalAnd
      (naturalEqual (bytes-length raw) textArtifactRequestEncodedLength)
      (naturalAnd
        (bytes-equal (dataBytesTakeValidated 8 raw) textArtifactRequestMagic)
        (naturalAnd
          (bytes-equal
            (dataBytesTakeValidated 8 (dataBytesDropValidated 8 raw))
            (dataBytesWord64LE
              (modelWord64FromNaturalTruncated textArtifactRequestVersion)))
          (naturalAnd
            (naturalOr
              (bytes-equal
                (dataBytesTakeValidated 8 (dataBytesDropValidated 16 raw))
                (dataBytesWord64LE
                  (modelWord64FromNaturalTruncated textArtifactTrainOperationWord)))
              (bytes-equal
                (dataBytesTakeValidated 8 (dataBytesDropValidated 16 raw))
                (dataBytesWord64LE
                  (modelWord64FromNaturalTruncated textArtifactPredictOperationWord))))
            (naturalAnd
              (textArtifactEncodedFieldTerminated 24 raw)
              (naturalAnd
                (textArtifactEncodedFieldTerminated 280 raw)
                (naturalAnd
                  (textArtifactEncodedFieldTerminated 536 raw)
                  (naturalAnd
                    (textArtifactEncodedFieldTerminated 792 raw)
                    (naturalAnd
                      (textArtifactEncodedFieldTerminated 1048 raw)
                      (naturalAnd
                        (textArtifactEncodedFieldTerminated 1304 raw)
                        (textArtifactEncodedFieldTerminated 1560 raw))))))))))))