namespace PureGym

open System
open System.Threading
open System.Threading.Tasks

type private CacheMessage<'a> =
    | TriggerUpdate
    | UpdateStored of 'a Task
    | Get of AsyncReplyChannel<'a>
    | Quit of AsyncReplyChannel<unit>

type internal Cache<'a> (obtainNew : CancellationToken -> 'a Task, expiry : 'a -> DateTime option) =
    let cts = new CancellationTokenSource ()

    let initialValue = obtainNew cts.Token

    let rec handle (value : 'a Task) (mailbox : MailboxProcessor<CacheMessage<'a>>) : Async<unit> =
        async {
            let! message = mailbox.Receive ()

            match message with
            | Quit channel ->
                channel.Reply ()
                return ()
            | CacheMessage.UpdateStored newValue -> return! handle newValue mailbox
            | CacheMessage.TriggerUpdate ->
                async {
                    let! a = Async.AwaitTask (obtainNew cts.Token)
                    let expiry = expiry a

                    match expiry with
                    | None -> return ()
                    | Some expiry ->

                    // a bit sloppy but :shrug:
                    do! Async.Sleep ((expiry - DateTime.Now) - TimeSpan.FromMinutes 1.0)

                    try
                        mailbox.Post CacheMessage.TriggerUpdate
                    with _ ->
                        // Post during shutdown sequence: drop it on the floor
                        ()

                    return ()
                }
                |> fun a -> Async.Start (a, cancellationToken = cts.Token)

                return! handle value mailbox
            | CacheMessage.Get reply ->
                let! valueAwaited = Async.AwaitTask value
                reply.Reply valueAwaited
                return! handle value mailbox
        }

    let mailbox = new MailboxProcessor<_> (handle initialValue)

    do
        mailbox.Start ()
        mailbox.Post CacheMessage.TriggerUpdate

    let isDisposing = ref 0
    let hasDisposed = TaskCompletionSource<unit> ()

    member this.GetCurrentValue () =
        try
            mailbox.PostAndAsyncReply CacheMessage.Get
        with
        // TODO I think this is the right exception...
        | :? InvalidOperationException ->
            raise (ObjectDisposedException (nameof (Cache)))

    interface IDisposable with
        member _.Dispose () =
            if Interlocked.Increment isDisposing = 1 then
                mailbox.PostAndReply CacheMessage.Quit
                (mailbox :> IDisposable).Dispose ()
                // We can't terminate the CTS until the mailbox has processed all client requests.
                // Otherwise we terminate the mailbox's state Task before it has finished querying that
                // task on behalf of clients.
                cts.Cancel ()
                cts.Dispose ()
                hasDisposed.SetResult ()
            else
                hasDisposed.Task.Result