module Platform.Linux.Nvidia.ProgramReadback

import Platform.Linux.Nvidia.PlanHost
import Platform.Linux.Nvidia.PlanHostRequest
import Runtime.NativePhysicalEmbeddedArtifact
import Runtime.NativeFilePlan
import Runtime.NativePhysicalProgram
import Model.Word64
import Model.Parameter
import Data.Bytes
import Std.Natural

-- A bounded, versioned capture of one retired native program session. The
-- payload's interpretation belongs to the system; this owner records its byte
-- extent, contiguous semaphore slots and fresh-channel issue count. No clock
-- subtraction occurs here. Reads happen only after the caller's final fence.
-- v1: magic[8], version:u64, dataBytes:u64, semaphoreCount:u64, payload,
-- semaphoreCount * prSemaphoreStride bytes, then issueCount:u64; little-endian.
def nvidiaProgramReadbackMagic : Bytes =
  b"ALPHAPRD"

def nvidiaProgramReadbackVersion : Nat =
  1

def nvidiaProgramReadbackHeader =
  (lambda unrestricted dataBytes : Nat .
    (lambda unrestricted semaphores : Nat .
      (bytes-append
        nvidiaProgramReadbackMagic
        (bytes-append
          (dataBytesWord64LE (modelWord64FromNaturalTruncated nvidiaProgramReadbackVersion))
          (bytes-append
            (dataBytesWord64LE (modelWord64FromNaturalTruncated dataBytes))
            (dataBytesWord64LE (modelWord64FromNaturalTruncated semaphores)))))))

def nvidiaProgramReadbackHeaderBytes : Nat =
  (bytes-length (nvidiaProgramReadbackHeader 0 0))

-- Envelope identity validation has retired before device execution begins.
-- Reuse that full identity buffer, not prExpected's single-word allocation.
def nvidiaProgramReadbackHeaderFits :
  (equal
    Nat
    (naturalLessOrEqual nvidiaProgramReadbackHeaderBytes nativePhysicalEmbeddedIdentityLength)
    1) =
  (refl Nat 1)

def nvidiaProgramReadback =
  (lambda unrestricted path : (family NativePhysicalOperand) .
    (lambda unrestricted semaphoreHost : Nat .
      (lambda unrestricted dataHost : Nat .
        (lambda unrestricted dataBytes : Nat .
          (lambda unrestricted semaphores : Nat .
            (lambda unrestricted counter : Nat .
              (lambda unrestricted tail : (family NativePhysicalCommands) .
                (let unrestricted semBytes =
                  (naturalMultiply semaphores prSemaphoreStride)
                  in
                  (let unrestricted semOffset =
                    (naturalAdd nvidiaProgramReadbackHeaderBytes dataBytes)
                    in
                    (let unrestricted counterOffset =
                      (naturalAdd semOffset semBytes)
                      in
                      (phCopy
                        prIdentityScratch
                        (nvidiaProgramReadbackHeader dataBytes semaphores)
                        (nativeFileOpen
                          path
                          nativeFileCandidateFlags
                          prSlotResult
                          prSlotStatus
                          b"program-readback-open"
                          (nativeFileWrite
                            (phSlot prSlotResult)
                            (phState prIdentityScratch)
                            nvidiaProgramReadbackHeaderBytes
                            (phImm 0)
                            prSlotStatus
                            b"program-readback-header"
                            (nativeFileWrite
                              (phSlot prSlotResult)
                              (phImm dataHost)
                              dataBytes
                              (phImm nvidiaProgramReadbackHeaderBytes)
                              prSlotStatus
                              b"program-readback-data"
                              (nativeFileWrite
                                (phSlot prSlotResult)
                                (phImm semaphoreHost)
                                semBytes
                                (phImm semOffset)
                                prSlotStatus
                                b"program-readback-fences"
                                (nativeFileWrite
                                  (phSlot prSlotResult)
                                  (phState counter)
                                  (bytes-length
                                    (dataBytesWord64LE (modelWord64FromNaturalTruncated 0)))
                                  (phImm counterOffset)
                                  prSlotStatus
                                  b"program-readback-issues"
                                  (nativeFileClose
                                    (phSlot prSlotResult)
                                    prSlotStatus
                                    b"program-readback-close"
                                    tail)))))))))))))))))