(*---------------------------------------------------------------------------
   Copyright (c) 2025 Thomas Gazagnaire. All rights reserved.
   SPDX-License-Identifier: MIT
  ---------------------------------------------------------------------------*)

let log_src = Logs.Src.create "cache"

module Log = (val Logs.src_log log_src : Logs.LOG)

module type KEY = sig
  type t

  val equal : t -> t -> bool
  val hash : t -> int
end

module type S = sig
  type key
  type 'a t

  val create :
    clock:_ Eio.Time.clock -> ?base_ttl:float -> ?jitter:float -> unit -> 'a t

  val get : 'a t -> key -> 'a option
  val set : 'a t -> key -> 'a -> unit
  val remove : 'a t -> key -> unit
  val get_or_compute : 'a t -> key -> (unit -> 'a) -> 'a
  val gc : 'a t -> unit
  val clear : 'a t -> unit
  val stats : 'a t -> int * int
end

module Make (K : KEY) : S with type key = K.t = struct
  type key = K.t

  module H = Hashtbl.Make (K)

  type 'a entry = { value : 'a; expires_at : float }

  type 'a t = {
    entries : 'a entry H.t;
    mutex : Eio.Mutex.t;
    base_ttl : float;
    jitter : float;
    now : unit -> float;
  }

  let create ~clock ?(base_ttl = 60.0) ?(jitter = 0.2) () =
    {
      entries = H.create 256;
      mutex = Eio.Mutex.create ();
      base_ttl;
      jitter;
      now = (fun () -> Eio.Time.now clock);
    }

  let now t = t.now ()

  let ttl_with_jitter t =
    let jitter_range = t.base_ttl *. t.jitter in
    let random_jitter = Random.float jitter_range -. (jitter_range /. 2.0) in
    t.base_ttl +. random_jitter

  let get t key =
    Eio.Mutex.use_ro t.mutex @@ fun () ->
    match H.find_opt t.entries key with
    | None -> None
    | Some entry ->
        if now t < entry.expires_at then Some entry.value
        else begin
          H.remove t.entries key;
          None
        end

  let set t key value =
    Eio.Mutex.use_rw ~protect:true t.mutex @@ fun () ->
    let entry = { value; expires_at = now t +. ttl_with_jitter t } in
    H.replace t.entries key entry

  let remove t key =
    Eio.Mutex.use_rw ~protect:true t.mutex @@ fun () -> H.remove t.entries key

  let get_or_compute t key compute =
    match get t key with
    | Some v -> v
    | None ->
        let v = compute () in
        set t key v;
        v

  let gc t =
    Eio.Mutex.use_rw ~protect:true t.mutex @@ fun () ->
    let current = now t in
    let to_remove =
      H.fold
        (fun k entry acc ->
          if current >= entry.expires_at then k :: acc else acc)
        t.entries []
    in
    List.iter (H.remove t.entries) to_remove;
    Log.debug (fun m ->
        m "Cache GC: removed %d expired entries" (List.length to_remove))

  let clear t =
    Eio.Mutex.use_rw ~protect:true t.mutex @@ fun () -> H.clear t.entries

  let stats t =
    Eio.Mutex.use_ro t.mutex @@ fun () ->
    let total = H.length t.entries in
    let current = now t in
    let valid =
      H.fold
        (fun _ entry acc -> if current < entry.expires_at then acc + 1 else acc)
        t.entries 0
    in
    (total, valid)
end

module String = Make (struct
  type t = string

  let equal = String.equal
  let hash = Hashtbl.hash
end)