// DNS response cache with TTL expiration and LRU eviction.
//
// Cache time is deliberately supplied by the caller.  This keeps the cache
// target-independent, makes expiry deterministic in tests, and lets the
// resolver use a monotonic clock supplied by its transport/runtime.

///|
pub enum NegativeCacheKind {
  NXDomain
  NoData
}

///|
pub enum CacheLookup {
  Miss
  Positive(DnsResult)
  Negative(NegativeCacheKind, DnsResult)
}

///|
pub struct DnsCache {
  hash : @hashmap.HashMap[String, DnsResult]
  expires : @hashmap.HashMap[String, Int64]
  negative : @hashmap.HashMap[String, NegativeCacheKind]
  access_order : @hashmap.HashMap[String, Int64]
  next_counter : Ref[Int64]
  max_size : Int
  hit_count : Ref[Int64]
  miss_count : Ref[Int64]
  expired_count : Ref[Int64]
}

///|
pub struct CacheStats {
  hits : Int
  misses : Int
  entries : Int
  expired : Int
}

// Default negative TTL in seconds when an authority SOA is absent.

///|
pub let negative_ttl : Int = 60

///|
fn clone_opt_options(options : Array[OptOption]) -> Array[OptOption] {
  let cloned : Array[OptOption] = Array::new(capacity=options.length())
  for option in options {
    cloned.push({ ..option, opt_data: option.opt_data.copy() })
  }
  cloned
}

///|
fn clone_rdata(data : RData) -> RData {
  match data {
    A(value) => A(value)
    AAAA(w1, w2, w3, w4) => AAAA(w1, w2, w3, w4)
    CNAME(name) => CNAME(name)
    NS(name) => NS(name)
    PTR(name) => PTR(name)
    MX(preference, exchange) => MX(preference, exchange)
    TXT(strings) => TXT(strings.copy())
    SOA(mname, rname, serial, refresh, retry, expire, minimum) =>
      SOA(mname, rname, serial, refresh, retry, expire, minimum)
    SRV(priority, weight, port, target) => SRV(priority, weight, port, target)
    OPT(options) => OPT(clone_opt_options(options))
    Unknown(raw) => Unknown(raw.copy())
  }
}

///|
fn clone_rr(record : RR) -> RR {
  { ..record, rdata: clone_rdata(record.rdata) }
}

///|
fn clone_dns_result(result : DnsResult) -> DnsResult {
  let answers : Array[RR] = Array::new(capacity=result.answers.length())
  for record in result.answers {
    answers.push(clone_rr(record))
  }
  { answers, cname_chain: result.cname_chain.copy() }
}

// DNS names are case-insensitive. The shared parser also gives escaped label
// octets a unique identity, while an unescaped terminal dot remains equivalent
// to the root-label spelling.

///|
fn canonical_dns_name(name : String) -> String {
  match split_labels_checked(name) {
    Ok(labels) => canonical_label_sequence(labels, 0)
    // Invalid inputs are rejected before they reach a wire encoder. Keep a
    // deterministic fallback here so cache lookup itself cannot abort first.
    Err(_) => "!invalid:" + name_to_lowercase(name)
  }
}

// Build a cache key from canonical DNS name and qtype.

///|
fn cache_key(name : String, qtype : UInt16) -> String {
  let canonical = canonical_dns_name(name)
  canonical.length().to_string() +
  ":" +
  canonical +
  ":TYPE:" +
  qtype.to_string()
}

///|
fn nxdomain_cache_key(name : String) -> String {
  let canonical = canonical_dns_name(name)
  canonical.length().to_string() + ":" + canonical + ":NXDOMAIN"
}

///|
pub fn DnsCache::new(max_size? : Int = 1024) -> DnsCache {
  {
    hash: HashMap([], capacity=0),
    expires: HashMap([], capacity=0),
    negative: HashMap([], capacity=0),
    access_order: HashMap([], capacity=0),
    next_counter: Ref(0),
    max_size: if max_size < 1 {
      1
    } else {
      max_size
    },
    hit_count: Ref(0),
    miss_count: Ref(0),
    expired_count: Ref(0),
  }
}

///|
fn DnsCache::remove_key(self : DnsCache, key : String) -> Unit {
  self.hash.remove(key)
  self.expires.remove(key)
  self.negative.remove(key)
  self.access_order.remove(key)
}

///|
fn DnsCache::remove_name_entries(self : DnsCache, name : String) -> Unit {
  let canonical = canonical_dns_name(name)
  // Length-prefixing keeps names such as "a" and "a:child" disjoint even
  // though ':' is a legal DNS label octet.
  let prefix = canonical.length().to_string() + ":" + canonical + ":"
  let keys : Array[String] = []
  self.hash.each(fn(key, _) { if key.has_prefix(prefix) { keys.push(key) } })
  for key in keys {
    self.remove_key(key)
  }
}

///|
fn DnsCache::touch(self : DnsCache, key : String) -> Unit {
  if self.next_counter.val == 9_223_372_036_854_775_807L {
    self.renormalize_access_order()
  }
  self.access_order[key] = self.next_counter.val
  self.next_counter.val += 1L
}

///|
fn DnsCache::renormalize_access_order(self : DnsCache) -> Unit {
  let entries : Array[(String, Int64)] = []
  self.access_order.each(fn(key, order) { entries.push((key, order)) })
  entries.sort_by(fn(left, right) {
    let by_order = left.1.compare(right.1)
    if by_order != 0 {
      by_order
    } else {
      left.0.compare(right.0)
    }
  })
  for index, entry in entries {
    self.access_order[entry.0] = index.to_int64()
  }
  self.next_counter.val = entries.length().to_int64()
}

///|
fn increment_counter(counter : Ref[Int64]) -> Unit {
  if counter.val < 9_223_372_036_854_775_807L {
    counter.val += 1L
  }
}

///|
fn public_counter(value : Int64) -> Int {
  if value > 2_147_483_647L {
    2_147_483_647
  } else {
    value.to_int()
  }
}

///|
fn DnsCache::expiry_from_ttl(now_ms : Int64, ttl_seconds : Int64) -> Int64 {
  let ttl = if ttl_seconds < 0L { 0L } else { ttl_seconds }
  let delta = ttl * 1000L
  // Saturation is only relevant for caller-injected clocks near Int64::max.
  // Normal process-monotonic timestamps retain their full multi-year range.
  if now_ms > 9_223_372_036_854_775_807L - delta {
    9_223_372_036_854_775_807L
  } else {
    now_ms + delta
  }
}

///|
fn DnsCache::evict_one_if_full(self : DnsCache) -> Unit {
  if self.hash.length() < self.max_size {
    return
  }
  let oldest_key = Ref("")
  let oldest_order = Ref(9_223_372_036_854_775_807L)
  self.access_order.each(fn(key, order) {
    if order < oldest_order.val {
      oldest_key.val = key
      oldest_order.val = order
    }
  })
  if oldest_key.val != "" {
    self.remove_key(oldest_key.val)
  }
}

///|
fn DnsCache::remaining_ttl_seconds_ms64(
  self : DnsCache,
  name : String,
  qtype : UInt16,
  now_ms : Int64,
) -> Int64? {
  let nx_key = nxdomain_cache_key(name)
  let key = if self.expires.contains(nx_key) && now_ms < self.expires[nx_key] {
    nx_key
  } else {
    cache_key(name, qtype)
  }
  match self.expires.get(key) {
    Some(expiry) if now_ms < expiry => Some((expiry - now_ms) / 1000L)
    _ => None
  }
}

///|
fn result_with_remaining_ttl(
  result : DnsResult,
  remaining_seconds : Int64,
) -> DnsResult {
  let remaining = remaining_seconds.reinterpret_as_uint64().to_uint()
  let answers : Array[RR] = Array::new(capacity=result.answers.length())
  for record in result.answers {
    answers.push({
      ..record,
      ttl: if record.ttl < remaining {
        record.ttl
      } else {
        remaining
      },
      rdata: clone_rdata(record.rdata),
    })
  }
  { answers, cname_chain: result.cname_chain.copy() }
}

// Look up an entry and retain its negative semantic, which is required for
// RFC 2308 cache hits to remain errors rather than become an empty success.

///|
pub fn DnsCache::lookup(
  self : DnsCache,
  name : String,
  qtype : UInt16,
  now_ms : Int,
) -> CacheLookup {
  self.lookup_ms64(name, qtype, now_ms.to_int64())
}

///|
fn DnsCache::lookup_ms64(
  self : DnsCache,
  name : String,
  qtype : UInt16,
  now_ms : Int64,
) -> CacheLookup {
  // RFC 2308 caches NXDOMAIN by QNAME/QCLASS, whereas NODATA and positive
  // answers are keyed by QNAME/QTYPE/QCLASS. Resolver cache entries are IN
  // class, so an active name-wide NXDOMAIN entry takes precedence.
  let nx_key = nxdomain_cache_key(name)
  if self.hash.contains(nx_key) && self.expires.contains(nx_key) {
    if now_ms < self.expires[nx_key] {
      self.touch(nx_key)
      increment_counter(self.hit_count)
      return Negative(NXDomain, clone_dns_result(self.hash[nx_key]))
    }
    self.remove_key(nx_key)
    increment_counter(self.expired_count)
  }
  let key = cache_key(name, qtype)
  if !self.hash.contains(key) || !self.expires.contains(key) {
    increment_counter(self.miss_count)
    return Miss
  }
  if now_ms >= self.expires[key] {
    self.remove_key(key)
    increment_counter(self.expired_count)
    increment_counter(self.miss_count)
    return Miss
  }
  self.touch(key)
  increment_counter(self.hit_count)
  let remaining = (self.expires[key] - now_ms) / 1000L
  let result = result_with_remaining_ttl(self.hash[key], remaining)
  match self.negative.get(key) {
    Some(kind) => Negative(kind, result)
    None => Positive(result)
  }
}

// Backwards-compatible positive-only cache read.  Resolver code should use
// `lookup` so that a cached NXDOMAIN/NODATA is never flattened into success.

///|
pub fn DnsCache::get(
  self : DnsCache,
  name : String,
  qtype : UInt16,
  now_ms : Int,
) -> DnsResult? {
  match self.lookup(name, qtype, now_ms) {
    Positive(result) => Some(result)
    _ => None
  }
}

///|
pub fn DnsCache::put_positive(
  self : DnsCache,
  name : String,
  qtype : UInt16,
  result : DnsResult,
  ttl_seconds : Int,
  now_ms : Int,
) -> Unit {
  self.put_positive_ms64(
    name,
    qtype,
    result,
    ttl_seconds.to_int64(),
    now_ms.to_int64(),
  )
}

///|
fn DnsCache::put_positive_ms64(
  self : DnsCache,
  name : String,
  qtype : UInt16,
  result : DnsResult,
  ttl_seconds : Int64,
  now_ms : Int64,
) -> Unit {
  self.remove_key(nxdomain_cache_key(name))
  let key = cache_key(name, qtype)
  // RFC 1035 TTL=0 means do not cache. Remove an existing entry for the same
  // key, but never evict a different live entry merely to insert something
  // that is already expired.
  if ttl_seconds <= 0L {
    self.remove_key(key)
    return
  }
  if !self.hash.contains(key) {
    self.evict_one_if_full()
  }
  // Cache entries are immutable snapshots. DnsResult contains mutable arrays,
  // so retaining a caller-owned value would let later caller mutations poison
  // future resolver results.
  self.hash[key] = clone_dns_result(result)
  self.expires[key] = DnsCache::expiry_from_ttl(now_ms, ttl_seconds)
  self.negative.remove(key)
  self.touch(key)
}

///|
pub fn DnsCache::put_negative(
  self : DnsCache,
  name : String,
  qtype : UInt16,
  kind : NegativeCacheKind,
  ttl_seconds : Int,
  now_ms : Int,
) -> Unit {
  self.put_negative_ms64(
    name,
    qtype,
    kind,
    ttl_seconds.to_int64(),
    now_ms.to_int64(),
  )
}

///|
fn DnsCache::put_negative_ms64(
  self : DnsCache,
  name : String,
  qtype : UInt16,
  kind : NegativeCacheKind,
  ttl_seconds : Int64,
  now_ms : Int64,
) -> Unit {
  let key = if kind is NXDomain {
    // A fresh NXDOMAIN supersedes every type-specific entry for this name.
    self.remove_name_entries(name)
    nxdomain_cache_key(name)
  } else {
    self.remove_key(nxdomain_cache_key(name))
    cache_key(name, qtype)
  }
  if ttl_seconds <= 0L {
    if kind is NXDomain {
      self.remove_name_entries(name)
    } else {
      self.remove_key(key)
    }
    return
  }
  if !self.hash.contains(key) {
    self.evict_one_if_full()
  }
  self.hash[key] = { answers: [], cname_chain: [] }
  self.expires[key] = DnsCache::expiry_from_ttl(now_ms, ttl_seconds)
  self.negative[key] = kind
  self.touch(key)
}

// Legacy writer retained for users of the initial API.  New resolver code
// uses `put_positive` or `put_negative` explicitly.

///|
pub fn DnsCache::put(
  self : DnsCache,
  name : String,
  qtype : UInt16,
  result : DnsResult,
  ttl : Int,
  is_negative : Bool,
  now_ms : Int,
) -> Unit {
  if is_negative {
    self.put_negative(name, qtype, NXDomain, ttl, now_ms)
  } else {
    self.put_positive(name, qtype, result, ttl, now_ms)
  }
}

///|
pub fn DnsCache::prune_expired(self : DnsCache, now_ms : Int) -> Unit {
  self.prune_expired_ms64(now_ms.to_int64())
}

///|
fn DnsCache::prune_expired_ms64(self : DnsCache, now_ms : Int64) -> Unit {
  let expired_keys : Array[String] = Array::new(capacity=16)
  self.expires.each(fn(key, expiry) {
    if now_ms >= expiry {
      expired_keys.push(key)
    }
  })
  for key in expired_keys {
    self.remove_key(key)
    increment_counter(self.expired_count)
  }
}

///|
pub fn DnsCache::stats(self : DnsCache) -> CacheStats {
  {
    hits: public_counter(self.hit_count.val),
    misses: public_counter(self.miss_count.val),
    entries: self.hash.length(),
    expired: public_counter(self.expired_count.val),
  }
}

///|
pub fn DnsCache::len(self : DnsCache) -> Int {
  self.hash.length()
}

///|
pub fn DnsCache::is_empty(self : DnsCache) -> Bool {
  self.hash.length() == 0
}

///|
pub fn DnsCache::contains(
  self : DnsCache,
  name : String,
  qtype : UInt16,
  now_ms : Int,
) -> Bool {
  self.contains_ms64(name, qtype, now_ms.to_int64())
}

///|
fn DnsCache::contains_ms64(
  self : DnsCache,
  name : String,
  qtype : UInt16,
  now_ms : Int64,
) -> Bool {
  let nx_key = nxdomain_cache_key(name)
  if self.hash.contains(nx_key) &&
    self.expires.contains(nx_key) &&
    now_ms < self.expires[nx_key] {
    true
  } else {
    let key = cache_key(name, qtype)
    self.hash.contains(key) &&
    self.expires.contains(key) &&
    now_ms < self.expires[key]
  }
}

///|
pub fn DnsCache::remove_entry(
  self : DnsCache,
  name : String,
  qtype : UInt16,
) -> Unit {
  self.remove_key(cache_key(name, qtype))
  self.remove_key(nxdomain_cache_key(name))
}