# Nim-LibP2P
# Copyright (c) 2024 Status Research & Development GmbH
# Licensed under either of
#  * Apache License, version 2.0, ([LICENSE-APACHE](LICENSE-APACHE))
#  * MIT license ([LICENSE-MIT](LICENSE-MIT))
# at your option.
# This file may not be copied, modified, or distributed except according to
# those terms.

{.push raises: [].}

import std/[deques, sequtils]
import stew/[byteutils, results, endians2]
import chronos, chronos/transports/[osnet, ipnet], metrics
import ../[multiaddress, multicodec]
import ../switch
import ../wire
import ../utils/heartbeat
import ../crypto/crypto

logScope:
  topics = "libp2p wildcardresolverservice"

type
  WildcardAddressResolverService* = ref object of Service
    ## Service used to resolve wildcard addresses of the type "0.0.0.0" for IPv4 or "::" for IPv6.
    ## When used with a `Switch`, this service will be automatically started and stopped
    ## when the switch starts and stops. This is facilitated by adding the service to the switch's
    ## list of services using the `.withServices(@[svc])` method in the `SwitchBuilder`.
    networkInterfaceProvider: NetworkInterfaceProvider
    ## Provides a list of network addresses.
    addressMapper: AddressMapper
    ## An implementation of an address mapper that takes a list of listen addresses and expands each wildcard address
    ## to the respective list of interface addresses. As an example, if the listen address is 0.0.0.0:4001
    ## and the machine has 2 interfaces with IPs 172.217.11.174 and 64.233.177.113, the address mapper will
    ## expand the wildcard address to 172.217.11.174:4001 and 64.233.177.113:4001.
    scheduleHandle: Future[void]
    ## Represents the task that is scheduled to run the service.
    scheduleInterval: Opt[Duration]
    ## The interval at which the service should run.

  NetworkInterfaceProvider* = ref object of RootObj

proc isLoopbackOrUp(networkInterface: NetworkInterface): bool =
  if (networkInterface.ifType == IfSoftwareLoopback) or
      (networkInterface.state == StatusUp): true else: false

method getAddresses*(
    networkInterfaceProvider: NetworkInterfaceProvider, addrFamily: AddressFamily
): seq[InterfaceAddress] {.base.} =
  ## This method retrieves the addresses of network interfaces based on the specified address family.
  ##
  ## The `getAddresses` method filters the available network interfaces to include only
  ## those that are either loopback or up. It then collects all the addresses from these
  ## interfaces and filters them to match the provided address family.
  ##
  ## Parameters:
  ## - `networkInterfaceProvider`: A provider that offers access to network interfaces.
  ## - `addrFamily`: The address family to filter the network addresses (e.g., `AddressFamily.IPv4` or `AddressFamily.IPv6`).
  ##
  ## Returns:
  ## - A sequence of `InterfaceAddress` objects that match the specified address family.
  let
    interfaces = getInterfaces().filterIt(it.isLoopbackOrUp())
    flatInterfaceAddresses = concat(interfaces.mapIt(it.addresses))
    filteredInterfaceAddresses =
      flatInterfaceAddresses.filterIt(it.host.family == addrFamily)
  return filteredInterfaceAddresses

proc new*(
    T: typedesc[WildcardAddressResolverService],
    scheduleInterval: Opt[Duration] = Opt.none(Duration),
    networkInterfaceProvider: NetworkInterfaceProvider = new(NetworkInterfaceProvider),
): T =
  ## This procedure initializes a new `WildcardAddressResolverService` with the provided
  ## schedule interval and network interface provider.
  ##
  ## Parameters:
  ## - `T`: The type descriptor for `WildcardAddressResolverService`.
  ## - `scheduleInterval`: An optional interval at which the service should run. Defaults to none.
  ## - `networkInterfaceProvider`: A provider that offers access to network interfaces. Defaults to a new instance of `NetworkInterfaceProvider`.
  ##
  ## Returns:
  ## - A new instance of `WildcardAddressResolverService`.
  return T(
    scheduleInterval: scheduleInterval,
    networkInterfaceProvider: networkInterfaceProvider,
  )

proc schedule(
    service: WildcardAddressResolverService, switch: Switch, interval: Duration
) {.async.} =
  ## Schedules the WildcardAddressResolverService to run at regular intervals.
  ##
  ## Sets up a schedule for the WildcardAddressResolverService to execute periodically based
  ## on the specified interval. It continuously runs the service on the given switch at the defined time intervals.
  ##
  ## Parameters:
  ## - `service`: The instance of the WildcardAddressResolverService that will be scheduled.
  ## - `switch`: The Switch object that the service will operate on.
  ## - `interval`: The Duration specifying how often the service should run.
  heartbeat "Scheduling WildcardAddressResolverService run", interval:
    await service.run(switch)

proc getProtocolArgument*(ma: MultiAddress, codec: MultiCodec): MaResult[seq[byte]] =
  var buffer: seq[byte]
  for item in ma:
    let
      ritem = ?item
      code = ?ritem.protoCode()
    if code == codec:
      let arg = ?ritem.protoAddress()
      return ok(arg)

  err("Multiaddress codec has not been found")

proc getWildcardMultiAddresses(
    interfaceAddresses: seq[InterfaceAddress],
    protocol: Protocol,
    port: Port,
    suffix: MultiAddress,
): seq[MultiAddress] =
  var addresses: seq[MultiAddress]
  for ifaddr in interfaceAddresses:
    var address = ifaddr.host
    address.port = port
    MultiAddress.init(address, protocol).toOpt.withValue(maddress):
      maddress.concat(suffix).toOpt.withValue(a):
        addresses.add(a)
  addresses

proc getWildcardAddress(
    maddress: MultiAddress,
    multiCodec: MultiCodec,
    anyAddr: openArray[uint8],
    addrFamily: AddressFamily,
    port: Port,
    networkInterfaceProvider: NetworkInterfaceProvider,
    peerId: MultiAddress,
): seq[MultiAddress] =
  var addresses: seq[MultiAddress]
  maddress.getProtocolArgument(multiCodec).toOpt.withValue(address):
    if address == anyAddr:
      let filteredInterfaceAddresses = networkInterfaceProvider.getAddresses(addrFamily)
      addresses.add(
        getWildcardMultiAddresses(filteredInterfaceAddresses, IPPROTO_TCP, port, peerId)
      )
    else:
      maddress.concat(peerId).toOpt.withValue(a):
        addresses.add(a)
  return addresses

proc expandWildcardAddresses(
    networkInterfaceProvider: NetworkInterfaceProvider,
    peerId: PeerId,
    listenAddrs: seq[MultiAddress],
): seq[MultiAddress] =
  let peerIdMa = MultiAddress.init(multiCodec("p2p"), peerId).valueOr:
    return default(seq[MultiAddress])

  var addresses: seq[MultiAddress]
  # In this loop we expand bounded addresses like `0.0.0.0` and `::` to list of interface addresses.
  for listenAddr in listenAddrs:
    if TCP_IP.matchPartial(listenAddr):
      listenAddr.getProtocolArgument(multiCodec("tcp")).toOpt.withValue(portArg):
        if len(portArg) == sizeof(uint16):
          let port = Port(uint16.fromBytesBE(portArg))
          if IP4.matchPartial(listenAddr):
            let wildcardAddresses = getWildcardAddress(
              listenAddr,
              multiCodec("ip4"),
              AnyAddress.address_v4,
              AddressFamily.IPv4,
              port,
              networkInterfaceProvider,
              peerIdMa,
            )
            addresses.add(wildcardAddresses)
          elif IP6.matchPartial(listenAddr):
            let wildcardAddresses = getWildcardAddress(
              listenAddr,
              multiCodec("ip6"),
              AnyAddress6.address_v6,
              AddressFamily.IPv6,
              port,
              networkInterfaceProvider,
              peerIdMa,
            )
            addresses.add(wildcardAddresses)
          else:
            listenAddr.concat(peerIdMa).withValue(ma):
              addresses.add(ma)
    else:
      let suffixed = listenAddr.concat(peerIdMa).valueOr:
        continue
      addresses.add(suffixed)
  addresses

method setup*(
    self: WildcardAddressResolverService, switch: Switch
): Future[bool] {.async.} =
  ## Sets up the `WildcardAddressResolverService`.
  ##
  ## This method adds the address mapper to the peer's list of address mappers and schedules the seervice to run
  ## at the specified interval.
  ##
  ## Parameters:
  ## - `self`: The instance of `WildcardAddressResolverService` being set up.
  ## - `switch`: The switch context in which the service operates.
  ##
  ## Returns:
  ## - A `Future[bool]` that resolves to `true` if the setup was successful, otherwise `false`.
  self.addressMapper = proc(
      listenAddrs: seq[MultiAddress]
  ): Future[seq[MultiAddress]] {.async.} =
    echo listenAddrs
    return expandWildcardAddresses(
      self.networkInterfaceProvider, switch.peerInfo.peerId, listenAddrs
    )

  debug "Setting up WildcardAddressResolverService"
  let hasBeenSetup = await procCall Service(self).setup(switch)
  if hasBeenSetup:
    self.scheduleInterval.withValue(interval):
      self.scheduleHandle = schedule(self, switch, interval)
    switch.peerInfo.addressMappers.add(self.addressMapper)
  return hasBeenSetup

method run*(self: WildcardAddressResolverService, switch: Switch) {.async, public.} =
  ## Runs the WildcardAddressResolverService for a given switch.
  ##
  ## It updates the peer information for the provided switch by running the registered address mapper. Any other
  ## address mappers that are registered with the switch will also be run.
  ##
  trace "Running WildcardAddressResolverService"
  await switch.peerInfo.update()

method stop*(
    self: WildcardAddressResolverService, switch: Switch
): Future[bool] {.async, public.} =
  ## Stops the WildcardAddressResolverService.
  ##
  ## Handles the shutdown process of the WildcardAddressResolverService for a given switch.
  ## It cancels the scheduled task, if any, and removes the address mapper from the switch's list of address mappers.
  ## It then updates the peer information for the provided switch. Any wildcard address wont be resolved anymore.
  ##
  ## Parameters:
  ## - `self`: The instance of the WildcardAddressResolverService.
  ## - `switch`: The Switch object associated with the service.
  ##
  ## Returns:
  ## - A future that resolves to `true` if the service was successfully stopped, otherwise `false`.
  debug "Stopping WildcardAddressResolverService"
  let hasBeenStopped = await procCall Service(self).stop(switch)
  if hasBeenStopped:
    if not isNil(self.scheduleHandle):
      self.scheduleHandle.cancel()
      self.scheduleHandle = nil
    switch.peerInfo.addressMappers.keepItIf(it != self.addressMapper)
    await switch.peerInfo.update()
  return hasBeenStopped