open System
open System.Collections.Concurrent
open System.Collections.Generic
open System.Timers

type CacheExpirationPolicy =
| NoExpiration
| AbsoluteExpiration of TimeSpan
| SlidingExpiration of TimeSpan

type CacheEntryExpiration =
| NeverExpires
| ExpiresAt of DateTime
| ExpiresAfter of TimeSpan

type CacheEntry<'key, 'value> =
    {
        Key: 'key
        Value: 'value
        Expiration: CacheEntryExpiration
        LastUsage: DateTime
    }

module CacheExpiration =
    let isExpired (entry: CacheEntry<_,_>) =
        match entry.Expiration with
        | NeverExpires -> false
        | ExpiresAt date -> DateTime.UtcNow > date
        | ExpiresAfter window -> (DateTime.UtcNow - entry.LastUsage) > window

type private IMemoryCacheStore<'key, 'value> =
    inherit IEnumerable<CacheEntry<'key, 'value>>
    abstract member Add: CacheEntry<'key, 'value> -> unit
    abstract member GetOrAdd: 'key -> ('key -> CacheEntry<'key, 'value>) -> CacheEntry<'key, 'value>
    abstract member Remove: 'key -> unit
    abstract member Contains: 'key -> bool
    abstract member Update: 'key -> (CacheEntry<'key, 'value> -> CacheEntry<'key, 'value>) -> unit
    abstract member TryFind: 'key -> CacheEntry<'key, 'value> option

type MemoryCache<'key, 'value> (?cacheExpirationPolicy) =
    let policy = defaultArg cacheExpirationPolicy NoExpiration
    let store = 
        let entries = ConcurrentDictionary<'key, CacheEntry<'key, 'value>>()
        let get, getEnumerator =
            let values = entries |> Seq.map (fun kvp -> kvp.Value)
            (fun () -> values), (fun () -> values.GetEnumerator())
        {new IMemoryCacheStore<'key, 'value> with
            member __.Add entry = entries.AddOrUpdate(entry.Key, entry, fun _ _ -> entry) |> ignore
            member __.GetOrAdd key getValue = entries.GetOrAdd(key, getValue)
            member __.Remove key = entries.TryRemove key |> ignore
            member __.Contains key = entries.ContainsKey key
            member __.Update key update = 
                match entries.TryGetValue(key) with
                | (true, entry) -> entries.AddOrUpdate(key, entry, fun _ entry -> update entry) |> ignore
                | _ -> ()
            member __.TryFind key =
                match entries.TryGetValue(key) with
                | (true, entry) -> Some entry
                | _ -> None
            member __.GetEnumerator () = getEnumerator ()
            member __.GetEnumerator () = getEnumerator () :> Collections.IEnumerator
        }
    
    let checkExpiration () =
        store 
        |> Seq.filter CacheExpiration.isExpired
        |> Seq.map (fun entry -> entry.Key)
        |> Seq.iter store.Remove

    let newCacheEntry key value =
        { Key = key
          Value = value
          Expiration = match policy with
                       | NoExpiration -> NeverExpires
                       | AbsoluteExpiration time -> ExpiresAt (DateTime.UtcNow + time)
                       | SlidingExpiration window -> ExpiresAfter window
          LastUsage = DateTime.UtcNow
        }

    let add key value =
        if key |> store.Contains
        then store.Update key (fun entry -> {entry with Value = value; LastUsage = DateTime.UtcNow})
        else store.Add <| newCacheEntry key value

    let remove key = 
        store.Remove key

    let get key =
        store.TryFind key |> Option.bind (fun entry -> Some entry.Value)

    let getOrAdd key value = 
        store.GetOrAdd key (fun _ -> newCacheEntry key value) 
        |> fun entry -> entry.Value

    let getOrAddResult key f =
        store.GetOrAdd key (fun _ -> newCacheEntry key <| f())
        |> fun entry -> entry.Value

    let getTimer (expiration: TimeSpan) =
        if expiration.TotalSeconds < 1.0
        then TimeSpan.FromMilliseconds 100.0
        elif expiration.TotalMinutes < 1.0 
        then TimeSpan.FromSeconds 1.0 
        else TimeSpan.FromMinutes 1.0
        |> fun interval -> new Timer(interval.TotalMilliseconds)

    let timer = 
       match policy with
       | NoExpiration -> None
       | AbsoluteExpiration time -> time |> getTimer |> Some
       | SlidingExpiration time -> time |> getTimer |> Some    

    let observer =
        match timer with
        | Some t -> 
            let disposable = t.Elapsed |> Observable.subscribe (fun _ -> checkExpiration())
            t.Start()
            Some disposable
        | None -> None

    member __.Add key value = add key value
    member __.Remove key = remove key
    member __.Get key = get key
    member __.GetOrAdd key value = getOrAdd key value    
    member __.GetOrAddResult key f = getOrAddResult key f

[<AutoOpen>]
module Memoization =
    type private MemoizationCache<'a,'b when 'a: equality> (?cacheExpirationPolicy) =
        let cache = 
            match cacheExpirationPolicy with
            | Some policy -> new MemoryCache<string,'b>(policy)
            | None -> new MemoryCache<string,'b>()
        let getKey (key: 'a) = 
            if key |> box |> isNull
            then typeof<'a>.FullName
            else sprintf "%s_%d" typeof<'a>.FullName <| key.GetHashCode()
        member __.TryGetValue key =
            let keyString = getKey key
            cache.Get keyString
        member __.GetOrAdd key getValue =
            let keyString = getKey key
            cache.GetOrAddResult keyString getValue

    let private memoizeWithCache f (cache: MemoizationCache<_,_>) x =
        (fun () -> f x) |> cache.GetOrAdd x            

    /// Create a new function that remembers the results of the given function
    /// for each unique input parameter, returning the cached result instead
    /// of recomputing it each time the function is called, optimizing the execution speed
    /// at the cost of increased memory usage.
    ///
    /// Note:  This optimization should only be used with referentially transparent functions
    let memoize f = 
        new MemoizationCache<_,_>() |> memoizeWithCache f

    /// Create a new function that remembers the results of the given function
    /// for each unique input parameter, returning the cached result instead
    /// of recomputing it each time the function is called, optimizing the execution speed
    /// at the cost of increased memory usage.  Uses a specified Cache Expiration Policy
    /// to limit the amount of time the results are stored, to allow memory to be free'd
    /// after a certain amount of time.
    ///
    /// Note:  This optimization should only be used with referentially transparent functions
    let memoizeWithExpiration policy f = 
        new MemoizationCache<_,_>(policy) |> memoizeWithCache f