package com.kaspersky.his.common.ip

import java.net.{Inet4Address, Inet6Address, InetAddress}
import java.nio.ByteBuffer

sealed trait CIDR {
  def address: InetAddress
  def fixedBits: Int
  override def toString: String = s"${address.getHostAddress}/$fixedBits"
  def isInRange(check: InetAddress): Boolean

  def minIp: InetAddress
  def maxIp: InetAddress
}

object CIDR {

  private[ip] val cidr4Size = 32
  private[ip] val cidr6Size = 128

  def parse(str: String): CIDR = {
    val Array(addr, b) = str.split('/')
    val bits = b.toInt
    InetAddress.getByName(addr) match {
      case ia4: Inet4Address =>
        val bits4 =
          if (bits > cidr4Size)
            bits - cidr6Size + cidr4Size
          else
            bits
        CIDR4(ia4, bits4)
      case ia6: Inet6Address => CIDR6(ia6, b.toInt)
    }
  }
}

case class CIDR6(address: Inet6Address, fixedBits: Int) extends CIDR {

  private val size = CIDR.cidr6Size
  private val bytes = size / java.lang.Byte.SIZE

  private val hmask = byteMask(fixedBits.min(size / 2))
  private val lmask = byteMask((fixedBits - size / 2).max(0))
  private val (hcidr, lcidr) = getBytes(address)

  def isInRange(check: InetAddress): Boolean = {
    check match {
      case ia: Inet6Address =>
        val (h, l) = getBytes(ia)
        (h & hmask) == (hcidr & hmask) && (l & lmask) == (lcidr & lmask)
      case _ =>
        false
    }
  }

  private def getBytes(address: Inet6Address): (Long, Long) = {
    val bb = ByteBuffer.allocate(bytes).put(address.getAddress)
    (bb.getLong(0), bb.getLong(java.lang.Long.BYTES))
  }

  private def byteMask(length: Int) = {
    (1 to length).map(x => 1l << (size / 2 - x)).sum
  }

  def minIp: InetAddress = getBytes(address) match {
    case (h, l) =>
      val bb = ByteBuffer.allocate(bytes)
      bb.putLong(h & hmask)
      bb.putLong(l & lmask)
      InetAddress.getByAddress(bb.array())
  }

  def maxIp: InetAddress = getBytes(address) match {
    case (h, l) =>
      val bb = ByteBuffer.allocate(bytes)
      bb.putLong(h | ~hmask)
      bb.putLong(l | ~lmask)
      InetAddress.getByAddress(bb.array())
  }

}

case class CIDR4(address: Inet4Address, fixedBits: Int) extends CIDR {

  private val size = CIDR.cidr4Size
  private val bytes = size / java.lang.Byte.SIZE

  private val mask = byteMask(fixedBits)
  private val cidr = getBytes(address)

  def isInRange(check: InetAddress): Boolean = {
    check match {
      case ia: Inet4Address =>
        val byteAddr = getBytes(ia)
        (byteAddr & mask) == (cidr & mask)
      case _ =>
        false
    }
  }

  private def getBytes(address: Inet4Address): Int = {
    val bb = ByteBuffer.allocate(bytes)
    bb.put(address.getAddress)
    bb.getInt(0)
  }

  private def byteMask(length: Int) = {
    (1 to length).map(x => 1 << (size - x)).sum
  }

  def minIp: InetAddress = getBytes(address) match {
    case byteAddr =>
      InetAddress.getByAddress(ByteBuffer.allocate(bytes).putInt(byteAddr).array())
  }

  def maxIp: InetAddress = getBytes(address) match {
    case byteAddr =>
      InetAddress.getByAddress(ByteBuffer.allocate(bytes).putInt(byteAddr | ~mask).array())
  }
}