#r "/opt/fsi/packages/AWSSDK.Core.dll"
#r "/opt/fsi/packages/AWSSDK.S3.dll"
#nowarn "0760"
#load "plan.fsx"
#load "cache-gc.fsx"

open System
open System.Collections.Generic
open System.IO
open System.Net
open System.Net.Http
open System.Security.Cryptography
open System.Security.Cryptography.X509Certificates
open System.Text
open System.Text.Json
open System.Text.Json.Serialization
open System.Threading.Tasks
open Amazon.Runtime
open Amazon.S3
open Amazon.S3.Model
open Contract
open Plan

[<CLIMutable>]
type CacheManifest =
    { [<JsonPropertyName("schemaVersion")>]
      SchemaVersion: int
      [<JsonPropertyName("cacheVersion")>]
      CacheVersion: string
      [<JsonPropertyName("inputHashes")>]
      InputHashes: Dictionary<string, string>
      [<JsonPropertyName("payloadKey")>]
      PayloadKey: string }

let bucket = "project-delivery"

let credential path =
    require (File.Exists path) "cache credential unavailable"
    let value = File.ReadAllText(path).Trim()
    require (not <| String.IsNullOrWhiteSpace value) "cache credential unavailable"
    value

let root = X509CertificateLoader.LoadCertificateFromFile "/cache-ca/ca.crt"

type CacheClientFactory(root: X509Certificate2) =
    inherit HttpClientFactory()

    override _.CreateHttpClient(_) =
        let handler = new HttpClientHandler()

        handler.ServerCertificateCustomValidationCallback <-
            fun _ certificate chain errors -> validateCertificate root certificate chain errors

        let client = new HttpClient(handler, true)
        client.Timeout <- TimeSpan.FromMinutes 15.0
        client

let configuration = AmazonS3Config()
configuration.ServiceURL <- "https://s3.quorate.internal"
configuration.AuthenticationRegion <- "us-east-1"
configuration.ForcePathStyle <- true
configuration.HttpClientFactory <- CacheClientFactory(root)

let client =
    AmazonS3Client(
        BasicAWSCredentials(credential "/cache-auth/accesskey", credential "/cache-auth/secretkey"),
        configuration
    )

let result (operation: Task<'T>) = operation.GetAwaiter().GetResult()

let parallel work =
    Async.StartAsTask(Async.Parallel work, cancellationToken = cancellationToken)
    |> result
    |> ignore

let sameHashes (left: Dictionary<string, string>) (right: Dictionary<string, string>) =
    not (isNull (box left))
    && not (isNull (box right))
    && left.Count = right.Count
    && left |> Seq.forall (fun item -> right.TryGetValue item.Key = (true, item.Value))

let select (name: string) (assignment: Assignment) =
    try
        let request = GetObjectRequest(BucketName = bucket, Key = assignment.ExportKey)
        use reply = result (client.GetObjectAsync(request, cancellationToken))
        require (reply.HttpStatusCode = HttpStatusCode.OK) $"{name} cache manifest GET failed"
        use reader = new StreamReader(reply.ResponseStream, Encoding.UTF8)
        let manifest = reader.ReadToEnd() |> decode<CacheManifest>

        if
            manifest.SchemaVersion <> 1
            || manifest.CacheVersion <> assignment.CacheVersion
            || not (sameHashes manifest.InputHashes assignment.InputHashes)
        then
            printfn "cache=miss component=%s reason=incompatible" name
            assignment
        else
            let payload = manifest.PayloadKey

            require
                (not <| isNull payload
                 && payload.StartsWith("objects/", StringComparison.Ordinal)
                 && payload.EndsWith(".tar.zst", StringComparison.Ordinal))
                $"invalid {name} cache payload key"

            let digest =
                Path.GetFileNameWithoutExtension(Path.GetFileNameWithoutExtension payload)

            require (validSha256 digest) $"invalid {name} cache digest"

            let prefix =
                assignment.ExportKey[.. assignment.ExportKey.Length - "_latest.json".Length - 1]

            { assignment with
                Origin = "cache"
                Cache =
                    { Key = prefix + payload
                      Sha256 = digest } }
    with :? AmazonS3Exception as error when
        error.StatusCode = HttpStatusCode.Forbidden
        || error.StatusCode = HttpStatusCode.NotFound ->
        printfn "cache=miss component=%s reason=absent" name
        assignment

let download (item: CacheObject) (destination: string) =
    let temporary = destination + ".incoming"

    try
        let request = GetObjectRequest(BucketName = bucket, Key = item.Key)
        use reply = result (client.GetObjectAsync(request, cancellationToken))
        require (reply.HttpStatusCode = HttpStatusCode.OK) "cache payload GET failed"
        use input = reply.ResponseStream
        use output = File.Create temporary
        use algorithm = SHA256.Create()
        use hashing = new CryptoStream(output, algorithm, CryptoStreamMode.Write, true)
        input.CopyTo hashing
        hashing.FlushFinalBlock()
        let digest = Convert.ToHexString(algorithm.Hash).ToLowerInvariant()
        require (digest = item.Sha256) "cache payload digest mismatch"
        File.Move(temporary, destination, true)
    with error ->
        if File.Exists temporary then
            File.Delete temporary

        raise error

let restore (name: string) (assignment: Assignment) (destination: string) =
    let archive = Path.Combine("/tmp", Path.GetRandomFileName() + ".tar.zst")

    try
        download assignment.Cache archive
        extractZstd archive destination
        printfn "cache=restored component=%s" name
    finally
        if File.Exists archive then
            File.Delete archive

let objectDigest (key: string) =
    try
        let request = GetObjectMetadataRequest(BucketName = bucket, Key = key)
        let reply = result (client.GetObjectMetadataAsync(request, cancellationToken))

        require (reply.HttpStatusCode = HttpStatusCode.OK) "cache metadata HEAD failed"
        Some reply.Metadata["x-amz-meta-sha256"]
    with :? AmazonS3Exception as error when error.StatusCode = HttpStatusCode.NotFound ->
        None

let upload (key: string) (path: string) (digest: string) =
    match objectDigest key with
    | Some actual ->
        require (actual = digest) "existing cache payload digest mismatch"
        printfn "cache=present key=%s" key
    | None ->
        use input = File.OpenRead path

        let request =
            PutObjectRequest(BucketName = bucket, Key = key, InputStream = input, IfNoneMatch = "*")

        request.AutoCloseStream <- false
        request.AutoResetStreamPosition <- false
        request.Metadata["sha256"] <- digest

        try
            let reply = result (client.PutObjectAsync(request, cancellationToken))
            require (reply.HttpStatusCode = HttpStatusCode.OK) "cache payload PUT failed"
            printfn "cache=uploaded key=%s bytes=%d" key input.Length
        with :? AmazonS3Exception as error when error.StatusCode = HttpStatusCode.PreconditionFailed ->
            require (objectDigest key = Some digest) "existing cache payload digest mismatch"
            printfn "cache=present key=%s" key

let putManifest (assignment: Assignment) (payloadKey: string) =
    let manifest =
        { SchemaVersion = 1
          CacheVersion = assignment.CacheVersion
          InputHashes = assignment.InputHashes
          PayloadKey = payloadKey }

    use input = new MemoryStream(JsonSerializer.SerializeToUtf8Bytes(manifest, json))

    let request =
        PutObjectRequest(BucketName = bucket, Key = assignment.ExportKey, InputStream = input)

    request.ContentType <- "application/json"
    let reply = result (client.PutObjectAsync(request, cancellationToken))
    require (reply.HttpStatusCode = HttpStatusCode.OK) "cache manifest PUT failed"
    CacheGc.prune client bucket assignment.ExportKey payloadKey cancellationToken

let export (name: string) (assignment: Assignment) (root: string) (paths: string seq) =
    let archive = Path.Combine("/tmp", Path.GetRandomFileName() + ".tar.zst")

    try
        let digest = createZstd root paths archive
        let payloadKey = $"objects/{digest}.tar.zst"

        let prefix =
            assignment.ExportKey[.. assignment.ExportKey.Length - "_latest.json".Length - 1]

        upload (prefix + payloadKey) archive digest
        putManifest assignment payloadKey
        printfn "cache=exported component=%s" name
    finally
        if File.Exists archive then
            File.Delete archive

match Environment.GetEnvironmentVariable "CACHE_ACTION" with
| "pull" ->
    let revision = Environment.GetEnvironmentVariable "SOURCE_REVISION"
    let plan = create revision select
    write plan

    [ "toolchain", plan.Toolchain, "/toolchain"
      "go-build", plan.Build, "/build-cache" ]
    |> List.choose (fun (name, assignment, destination) ->
        if assignment.Origin = "cache" then
            Some(async { restore name assignment destination })
        else
            None)
    |> parallel
| "export" ->
    let plan = read ()

    [ async { export "toolchain" plan.Toolchain "/toolchain" [ "usr/lib/go"; "image-base"; "modules"; ".ready" ] }
      async { export "go-build" plan.Build "/build-cache" [ "." ] } ]
    |> parallel
| _ -> fail "invalid cache action"